Pennyw0rth-NetExec/cme/loaders/moduleloader.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import importlib
import os
import cme
from cme.context import Context
from cme.logger import CMEAdapter
from cme.paths import CME_PATH


class ModuleLoader:
    def __init__(self, args, db, logger):
        self.args = args
        self.db = db
        self.logger = logger

    def module_is_sane(self, module, module_path):
        """
        Check if a module has the proper attributes
        """
        module_error = False
        if not hasattr(module, 'name'):
            self.logger.error('{} missing the name variable'.format(module_path))
            module_error = True
        elif not hasattr(module, 'description'):
            self.logger.error('{} missing the description variable'.format(module_path))
            module_error = True
        elif not hasattr(module, 'supported_protocols'):
            self.logger.error('{} missing the supported_protocols variable'.format(module_path))
            module_error = True
        elif not hasattr(module, 'opsec_safe'):
            self.logger.error('{} missing the opsec_safe variable'.format(module_path))
            module_error = True
        elif not hasattr(module, 'multiple_hosts'):
            self.logger.error('{} missing the multiple_hosts variable'.format(module_path))
            module_error = True
        elif not hasattr(module, 'options'):
            self.logger.error('{} missing the options function'.format(module_path))
            module_error = True
        elif not hasattr(module, 'on_login') and not (module, 'on_admin_login'):
            self.logger.error('{} missing the on_login/on_admin_login function(s)'.format(module_path))
            module_error = True
        # elif not hasattr(module, 'chain_support'):
        #    self.logger.error('{} missing the chain_support variable'.format(module_path))
        #    module_error = True

        if module_error:
            return False
        return True

    def load_module(self, module_path):
        """
        Load a module, initializing it and checking that it has the proper attributes
        """
        try:
            spec = importlib.util.spec_from_file_location("CMEModule", module_path)
            module = spec.loader.load_module().CMEModule()

            if self.module_is_sane(module, module_path):
                return module
        except Exception as e:
            self.logger.error('Failed loading module at {}: {}'.format(module_path, e))
        return None

    def init_module(self, module_path):
        """
        Initialize a module for execution
        """
        module = None
        module = self.load_module(module_path)

        if module and (self.args.protocol in module.supported_protocols):
            try:
                module_logger = CMEAdapter(extra={'module_name': module.name.upper()})
            except Exception as e:
                self.logger.error(f"Error loading CMEAdaptor for module {module.name.upper()}: {e}")
            context = Context(self.db, module_logger, self.args)
            module_options = {}

            for option in self.args.module_options:
                key, value = option.split('=', 1)
                module_options[str(key).upper()] = value

            module.options(context, module_options)
        return module

    def get_module_info(self, module_path):
        """
        Get the path, description, and options from a module
        """
        try:
            spec = importlib.util.spec_from_file_location("CMEModule", module_path)
            module_spec = spec.loader.load_module().CMEModule

            module = {
                f"{module_spec.name.lower()}":
                    {
                        "path": module_path,
                        "description": module_spec.description,
                        "options": module_spec.options.__doc__
                    }
            }
            if self.module_is_sane(module_spec, module_path):
                return module
        except Exception as e:
            self.logger.error('Failed loading module at {}: {}'.format(module_path, e))
        return None

    def list_modules(self):
        """
        List modules without initializing them
        """
        modules = {}
        modules_paths = [os.path.join(os.path.dirname(cme.__file__), 'modules'), os.path.join(CME_PATH, 'modules')]

        for path in modules_paths:
            for module in os.listdir(path):
                if module[-3:] == '.py' and module != 'example_module.py':
                    module_path = os.path.join(path, module)
                    module_data = self.get_module_info(module_path)
                    modules.update(module_data)
        return modules