import base64
import fnmatch
import importlib.util
import json
import logging
import shutil
import typing
import warnings
from pathlib import Path

import yaml

try:
    from yaml import CSafeDumper as Dumper
    from yaml import CSafeLoader as Loader
except ImportError:
    # Fall back to the pure-Python safe loader/dumper when libyaml is
    # unavailable. These must stay Safe* variants: module YAML is loaded with
    # them and the full Loader allows arbitrary object instantiation (S506).
    from yaml import SafeDumper as Dumper
    from yaml import SafeLoader as Loader

from packaging.version import parse
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.orm import Session

from empire.server.api.v2.module.module_dto import (
    ModuleBulkUpdateRequest,
    ModuleUpdateRequest,
)
from empire.server.common import helpers
from empire.server.core.config.config_manager import DATA_DIR
from empire.server.core.db import models
from empire.server.core.db.base import SessionLocal
from empire.server.core.exceptions import (
    ModuleExecutionException,
    ModuleValidationException,
)
from empire.server.core.module_models import (
    EmpireModule,
    EmpireModuleOption,
    LanguageEnum,
)
from empire.server.utils import data_util
from empire.server.utils.bof_packer import process_arguments
from empire.server.utils.option_util import convert_module_options, validate_options
from empire.server.utils.string_util import slugify

if typing.TYPE_CHECKING:
    from empire.server.common.empire import MainMenu
    from empire.server.core.download_service import DownloadService
    from empire.server.core.obfuscation_service import ObfuscationService

log = logging.getLogger(__name__)


class ModuleExecutionRequest(BaseModel):
    command: str
    data: str
    files: list[Path] = []


class ModuleService:
    def __init__(self, main_menu: "MainMenu"):
        self.main_menu = main_menu
        self.obfuscation_service: ObfuscationService = main_menu.obfuscationv2
        self.download_service: DownloadService = main_menu.downloadsv2
        self.dotnet_compiler = main_menu.dotnet_compiler

        self.module_source_path = main_menu.install_path / "data/module_source"
        self._obfuscated_module_source_path = DATA_DIR / "obfuscated_module_source"

        self.modules = {}

        with SessionLocal.begin() as db:
            self.load_modules(db)

    def get_all(self, hide_disabled: bool = False):
        if hide_disabled:
            return {k: v for k, v in self.modules.items() if v.enabled}
        return self.modules

    def get_by_id(self, uid: str):
        return self.modules.get(uid)

    def _load_custom_generate_class(self, module: EmpireModule):
        """
        Lazily import the .py file referenced by `module.advanced.custom_generate_path`
        and instantiate its `Module` class. Subsequent calls reuse the cached instance.
        Import / instantiation errors are re-raised as `ModuleValidationException`
        so the caller surfaces a specific user-facing message instead of the
        generic "Error generating script." fallback.

        Intentionally not thread-safe: two concurrent first-execute calls for
        the same module can both run `exec_module`. Last write wins, both
        threads get a valid instance. The in-tree custom_generate modules use
        a `@staticmethod` `generate`, so the duplicate instantiation is benign.
        """
        if module.advanced.generate_class is not None:
            return module.advanced.generate_class
        path = module.advanced.custom_generate_path
        spec = importlib.util.spec_from_file_location(f"{module.id}.py", path)
        if spec is None or spec.loader is None:
            raise ModuleValidationException(
                f"custom_generate module {module.id!r}: cannot build import spec for {path}"
            )
        try:
            imp_mod = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(imp_mod)
            module.advanced.generate_class = imp_mod.Module()
        except Exception as e:
            raise ModuleValidationException(
                f"custom_generate module {module.id!r}: failed to load {path}: {e}"
            ) from e
        return module.advanced.generate_class

    def update_module(
        self, db: Session, module: EmpireModule, module_req: ModuleUpdateRequest
    ):
        db_module: models.Module = db.scalars(
            select(models.Module).where(models.Module.id == module.id)
        ).first()
        db_module.enabled = module_req.enabled

        self.modules.get(module.id).enabled = module_req.enabled

    def update_modules(self, db: Session, module_req: ModuleBulkUpdateRequest):
        db_modules: list[models.Module] = db.scalars(
            select(models.Module).where(models.Module.id.in_(module_req.modules))
        ).all()

        for db_module in db_modules:
            db_module.enabled = module_req.enabled

        for db_module in db_modules:
            self.modules.get(db_module.id).enabled = module_req.enabled

    def execute_module(  # noqa: PLR0913 PLR0912 PLR0915
        self,
        db: Session,
        agent: models.Agent,
        module_id: str,
        params: dict,
        ignore_language_version_check: bool = False,
        ignore_admin_check: bool = False,
        modified_input: str | None = None,
        background_override: bool | None = None,
    ) -> tuple[ModuleExecutionRequest | None, str | None]:
        """
        Execute the module. Note this doesn't actually add the task to the queue,
        it only generates the module data needed for a task to be created.
        :param module_id: str
        :param params: the execution parameters
        :param user_id: the user executing the module
        :return: tuple with the response and an error message (if applicable)
        """
        module = self.get_by_id(module_id)

        if not module:
            return None, f"Module not found for id {module_id}"
        if not module.enabled:
            return None, "Cannot execute disabled module"

        if modified_input:
            module = self._create_modified_module(module, modified_input)

        cleaned_options, err = self._validate_module_params(
            db, module, agent, params, ignore_language_version_check, ignore_admin_check
        )

        if err:
            raise ModuleValidationException(err)

        module_data = self._generate_script(
            db,
            module,
            cleaned_options,
            agent.language,
        )
        if isinstance(module_data, tuple):
            warnings.warn(
                "Returning a tuple on errors from module generation is deprecated. Raise exceptions instead."
                "https://bc-security.gitbook.io/empire-wiki/module-development/powershell-modules#custom-generate",
                DeprecationWarning,
                stacklevel=5,
            )
            (module_data, err) = module_data
        else:
            # Not all modules return a tuple. If they just return a single value,
            # we don't want to throw an unpacking error.
            err = None

        # Should standardize on the return type.
        if not module_data:
            # This should probably be a ModuleExecutionException, but
            # for backwards compatability with 5.x, it needs to raise a 400
            raise ModuleValidationException(err or "module produced an empty script")

        if type(module_data) is not ModuleExecutionRequest:
            module_data = ModuleExecutionRequest(command="", data=module_data)

        if not module_data.data.isascii():
            # This previously returned 'None, 'module source contains non-ascii characters'
            # Was changed in 4.3 to print a warning.
            log.warning(f"Module source for {module_id} contains non-ascii characters")

        if module.language == LanguageEnum.powershell:
            module_data.data = helpers.strip_powershell_comments(module_data.data)

        if module.language == LanguageEnum.python:
            module_data.data = helpers.strip_python_comments(module_data.data)

        extension = module.output_extension.rjust(5) if module.output_extension else ""

        effective_background = (
            background_override
            if background_override is not None
            else module.background
        )

        if agent.language in ("ironpython", "python"):
            if module.language == "python":
                if effective_background:
                    module_data.command = "TASK_PYTHON_CMD_JOB"
                else:
                    command, data = self._handle_save_file_command(
                        "TASK_PYTHON", module.name, extension, module_data.data
                    )
                    module_data.command = command
                    module_data.data = data
            elif module.language == "powershell":
                if effective_background:
                    module_data.command = "TASK_POWERSHELL_CMD_JOB"
                else:
                    command, data = self._handle_save_file_command(
                        "TASK_POWERSHELL", module.name, extension, module_data.data
                    )
                    module_data.command = command
                    module_data.data = data
            elif module.language in ("csharp", "bof"):
                if effective_background:
                    module_data.command = "TASK_CSHARP_CMD_JOB"
                else:
                    module_data.command = "TASK_CSHARP_CMD_WAIT"
            else:
                log.error(
                    f"Unsupported module language {module.language} for agent {agent.language}"
                )

        elif agent.language == "csharp":
            if module.language in ("csharp", "bof"):
                if effective_background:
                    module_data.command = "TASK_CSHARP_CMD_JOB"
                else:
                    module_data.command = "TASK_CSHARP_CMD_WAIT"
            elif module.language == "powershell":
                module_data.command = "TASK_POWERSHELL_CMD_JOB"
            else:
                log.error(
                    f"Unsupported module language {module.language} for agent {agent.language}"
                )

        elif agent.language == "powershell":
            if module.language == "powershell":
                if effective_background:
                    module_data.command = "TASK_POWERSHELL_CMD_JOB"
                else:
                    command, data = self._handle_save_file_command(
                        "TASK_POWERSHELL", module.name, extension, module_data.data
                    )
                    module_data.command = command
                    module_data.data = data
            elif module.language in ("csharp", "bof"):
                if effective_background:
                    module_data.command = "TASK_CSHARP_CMD_JOB"
                else:
                    module_data.command = "TASK_CSHARP_CMD_WAIT"
            else:
                log.error(
                    f"Unsupported module language {module.language} for agent {agent.language}"
                )
        elif agent.language == "go":
            if module.language == "powershell":
                if effective_background:
                    module_data.command = "TASK_POWERSHELL_CMD_JOB"
                else:
                    command, data = self._handle_save_file_command(
                        "TASK_POWERSHELL", module.name, extension, module_data.data
                    )
                    module_data.command = command
                    module_data.data = data
            elif module.language == "csharp":
                if effective_background:
                    module_data.command = "TASK_CSHARP_CMD_JOB"
                else:
                    module_data.command = "TASK_CSHARP_CMD_WAIT"
            elif module.language == "bof":
                module_data.command = "TASK_BOF_CMD_WAIT"
            elif module.language == "pe":
                module_data.command = "TASK_PE_CMD_WAIT"
            else:
                log.error(
                    f"Unsupported module language {module.language} for agent {agent.language}"
                )
        else:
            log.error(f"Unsupported agent language {agent.language}")
            return None, f"Unsupported agent language: {agent.language}"

        return module_data, None

    @staticmethod
    def _handle_save_file_command(cmd_type, module_name, extension, module_data):
        if extension:
            save_file_prefix = Path(module_name).name[:15]
            module_data = save_file_prefix.rjust(15) + extension + module_data
            return f"{cmd_type}_CMD_WAIT_SAVE", module_data
        return f"{cmd_type}_CMD_WAIT", module_data

    def _validate_module_params(  # noqa: PLR0913
        self,
        db: Session,
        module: EmpireModule,
        agent: models.Agent,
        params: dict[str, str],
        ignore_language_version_check: bool = False,
        ignore_admin_check: bool = False,
    ) -> tuple[dict[str, str] | None, str | None]:
        """
        Given a module and execution params, validate the input and return back a clean Dict for execution.
        :param module: EmpireModule
        :param params: the execution parameters set by the user
        :return: tuple with options and the error message (if applicable)
        """

        # Define valid agent/module language compatibility
        valid_language_mapping = {
            "go": {"bof", "powershell", "csharp"},
            "ironpython": {"bof", "powershell", "csharp", "python"},
            "powershell": {"bof", "powershell", "csharp"},
            "csharp": {"bof", "powershell", "csharp"},
            "python": {"python"},
        }

        # Ensure the agent's language is supported
        if agent.language not in valid_language_mapping:
            raise ModuleValidationException(
                f"Unsupported agent language '{agent.language}'."
            )

        # Ensure the module language is compatible with the agent's language
        if module.language not in valid_language_mapping.get(agent.language, {}):
            raise ModuleValidationException(
                f"agent language '{agent.language}' cannot run module language '{module.language}'."
            )

        converted_options = convert_module_options(module.options)

        options, err = validate_options(
            converted_options, params, db, self.download_service
        )

        if err:
            return None, err

        if not ignore_language_version_check and module.language == agent.language:
            module_version = parse(module.min_language_version or "0")
            agent_version = parse(agent.language_version or "0")

            # check if the agent/module PowerShell versions are compatible
            if module_version > agent_version:
                raise ModuleValidationException(
                    f"module requires language version {module.min_language_version} but agent running language version {agent.language_version}",
                )

        if module.needs_admin and not ignore_admin_check and not agent.high_integrity:
            raise ModuleValidationException(
                "module needs to run in an elevated context"
            )

        return options, None

    def _generate_script(  # noqa: PLR0911, PLR0912
        self,
        db: Session,
        module: EmpireModule,
        params: dict,
        agent_language: str,
        obfuscation_config: models.ObfuscationConfig = None,
    ) -> tuple[ModuleExecutionRequest | None, str | None]:
        """
        Generate the script to execute
        :param module: the execution parameters (already validated)
        :param params: the execution parameters
        :param obfuscation_config: the obfuscation config. If not provided, will look up from the db.
        :return: tuple containing the generated script and an error if it exists
        """
        if not obfuscation_config:
            obfuscation_config = self.obfuscation_service.get_obfuscation_config(
                db, module.language
            )
        if not obfuscation_config:
            obfuscation_enabled = False
            obfuscation_command = None
        else:
            obfuscation_enabled = obfuscation_config.enabled
            obfuscation_command = obfuscation_config.command

        if module.advanced.custom_generate:
            # In a future release we could refactor the modules to accept an obuscation_config,
            #  but there's little benefit to doing so at this point. So I'm saving myself the pain.
            try:
                kwargs = {}
                if module.language == LanguageEnum.bof:
                    kwargs["agent_language"] = agent_language
                generate_class = self._load_custom_generate_class(module)
                return generate_class.generate(
                    self.main_menu,
                    module,
                    params,
                    obfuscation_enabled,
                    obfuscation_command,
                    **kwargs,
                )
            except (ModuleValidationException, ModuleExecutionException):
                raise
            except Exception as e:
                log.error(f"Error generating script: {e}", exc_info=True)
                return None, "Error generating script."
        elif module.language == LanguageEnum.powershell:
            resp = self._generate_script_powershell(module, params, obfuscation_config)
            return ModuleExecutionRequest(command="", data=resp), None
        # We don't have obfuscation for other languages yet, but when we do,
        # we can pass it in here.
        elif module.language == LanguageEnum.python:
            resp = self._generate_script_python(module, params, obfuscation_config)
            return ModuleExecutionRequest(command="", data=resp), None
        elif module.language == LanguageEnum.csharp:
            return self.generate_script_csharp(module, params, obfuscation_config), None
        elif module.language == LanguageEnum.bof:
            if agent_language == "go":
                resp = self.generate_go_bof(module, params)
                return ModuleExecutionRequest(command="", data=resp), None
            if not obfuscation_config:
                obfuscation_config = self.obfuscation_service.get_obfuscation_config(
                    db, LanguageEnum.csharp
                )
            return self.generate_script_bof(module, params, obfuscation_enabled), None

        return None, "Unsupported language"

    def generate_script_bof(
        self,
        module: EmpireModule,
        params: dict,
        obfuscate: bool = False,
    ) -> ModuleExecutionRequest:
        bof_module = self.modules["csharp_code_execution_runcoff"]

        if params["Architecture"] == "x86":
            script_path = self.module_source_path / module.bof.x86
        else:
            script_path = self.module_source_path / module.bof.x64

        bof_data = script_path.read_bytes()
        b64_bof_data = base64.b64encode(bof_data).decode("utf-8")

        script_file = self.dotnet_compiler.compile_task(
            bof_module.compiler_yaml,
            bof_module.name,
            dot_net_version="net40",
            confuse=obfuscate,
        )

        filtered_params = {
            key: (
                value if value != "" else " "
            )  # Replace empty values with a blank space
            for key, value in params.items()
            if key.lower()
            not in [
                "agent",
                "dotnetversion",
                "architecture",
                "entrypoint",
            ]
        }

        formatted_args = " ".join(
            f'"{value}"' if " " in str(value) else str(value)
            for value in filtered_params.values()
        )

        params_dict = {}
        params_dict["Entrypoint"] = module.bof.entry_point or "go"
        params_dict["File"] = b64_bof_data
        params_dict["HexData"] = process_arguments(
            module.bof.format_string, formatted_args
        )

        final_base64_json = base64.b64encode(
            json.dumps(params_dict).encode("utf-8")
        ).decode("utf-8")

        return ModuleExecutionRequest(
            command="",
            data=f"{script_file}|,{final_base64_json}",
            files=[script_file],
        )

    def format_bof_output(
        self,
        bof_data_b64: str,
        hex_data: str,
        agent_language: str,
        obfuscate: bool = False,
        entry_point: str = "go",
    ) -> str | ModuleExecutionRequest:
        """
        Build the final output string for a BOF module.

        For Go agents, returns base64-encoded JSON with File and HexData.
        For .NET agents, compiles the RunCOFF wrapper and returns
        the compiled file path with base64-encoded JSON.
        """
        if agent_language == "go":
            params_dict = {
                "File": bof_data_b64,
                "HexData": hex_data,
            }
            return base64.b64encode(json.dumps(params_dict).encode("utf-8")).decode(
                "utf-8"
            )

        bof_module = self.modules["csharp_code_execution_runcoff"]
        script_file = self.dotnet_compiler.compile_task(
            bof_module.compiler_yaml,
            bof_module.name,
            dot_net_version="net40",
            confuse=obfuscate,
        )

        params_dict = {
            "Entrypoint": entry_point,
            "File": bof_data_b64,
            "HexData": hex_data,
        }

        final_base64_json = base64.b64encode(
            json.dumps(params_dict).encode("utf-8")
        ).decode("utf-8")

        return ModuleExecutionRequest(
            command="",
            data=f"{script_file}|,{final_base64_json}",
            files=[script_file],
        )

    def generate_go_bof(
        self,
        module: EmpireModule,
        params: dict,
        skip_params=False,
    ) -> str:
        if params["Architecture"] == "x86":
            script_path = self.module_source_path / module.bof.x86
        else:
            script_path = self.module_source_path / module.bof.x64

        bof_data = script_path.read_bytes()
        b64_bof_data = base64.b64encode(bof_data).decode("utf-8")

        filtered_params = {
            key: (
                value if value != "" else " "
            )  # Replace empty values with a blank space
            for key, value in params.items()
            if key.lower()
            not in [
                "agent",
                "dotnetversion",
                "architecture",
                "entrypoint",
            ]
        }

        formatted_args = " ".join(
            f'"{value}"' if " " in str(value) else str(value)
            for value in filtered_params.values()
        )

        if not skip_params:
            params_dict = {}
            params_dict["File"] = b64_bof_data
            params_dict["HexData"] = process_arguments(
                module.bof.format_string, formatted_args
            )
        else:
            params_dict = params

        final_base64_json = base64.b64encode(
            json.dumps(params_dict).encode("utf-8")
        ).decode("utf-8")

        return f"{final_base64_json}"

    def generate_go_pe(
        self,
        module: EmpireModule,
        params: dict[str, str],
        skip_params=False,
    ) -> str:
        """
        Generates a base64-encoded JSON structure for a PE file and its arguments.

        :param module: EmpireModule object containing PE file paths and format info.
        :param params: Dictionary of parameters, including architecture and arguments.
        :param skip_params: If True, bypasses argument formatting and uses params as is.
        :return: Base64-encoded JSON string.
        """
        # Determine the file path based on architecture
        if params["Architecture"] == "x86":
            script_path = self.module_source_path / module.pe.x86
        else:
            script_path = self.module_source_path / module.pe.x64

        # Read the PE file and encode it in base64
        pe_data = script_path.read_bytes()
        b64_pe_data = base64.b64encode(pe_data).decode("utf-8")

        # Filter and prepare arguments
        filtered_params = {
            key: (
                value if value != "" else " "
            )  # Replace empty values with a blank space
            for key, value in params.items()
            if key.lower()
            not in [
                "agent",
                "dotnetversion",
                "architecture",
                "entrypoint",
            ]
        }

        # Create a list of arguments
        formatted_args = [
            f'"{value}"' if " " in str(value) else str(value)
            for value in filtered_params.values()
        ]

        if not skip_params:
            params_dict = {"File": b64_pe_data, "Args": formatted_args}
        else:
            params_dict = params

        return base64.b64encode(json.dumps(params_dict).encode("utf-8")).decode("utf-8")

    def _generate_script_python(
        self,
        module: EmpireModule,
        params: dict,
        obfuscation_config: models.ObfuscationConfig,
    ) -> str:
        obfuscate = (
            obfuscation_config.enabled if obfuscation_config is not None else False
        )

        if module.script_path:
            script_path = self.module_source_path / module.script_path
            script = script_path.read_text()
        else:
            script = module.script

        for key, value in params.items():
            if key.lower() != "agent":
                script = script.replace("{{ " + key + " }}", value).replace(
                    "{{" + key + "}}", value
                )

        if obfuscate:
            script = self.obfuscation_service.python_obfuscate(script)

        return script

    def _generate_script_powershell(
        self,
        module: EmpireModule,
        params: dict,
        obfuscation_config: models.ObfuscationConfig,
    ) -> str:
        obfuscate = (
            obfuscation_config.enabled if obfuscation_config is not None else False
        )
        obfuscate_command = (
            obfuscation_config.command if obfuscation_config is not None else ""
        )

        if module.script_path:
            script, err = self.get_module_source(
                module_name=module.script_path,
                obfuscate=obfuscate,
                obfuscate_command=obfuscate_command,
            )

            if err:
                raise ModuleValidationException(err)
        elif obfuscate:
            script = self.obfuscation_service.obfuscate(
                module.script, obfuscate_command
            )
        else:
            script = module.script

        script_end = f" {module.script_end} "
        option_strings = []

        # This is where the code goes for all the modules that do not have a custom generate function.
        for key, value in params.items():
            if key.lower() not in ["agent", "outputfunction"] and value and value != "":
                if value.lower() == "true":
                    # if we're just adding a switch
                    # wannabe mustache templating.
                    # If we want to get more advanced, we can import a library for it.
                    this_option = module.advanced.option_format_string_boolean.replace(
                        "{{ KEY }}", str(key)
                    ).replace("{{KEY}}", str(key))
                    option_strings.append(f"{this_option}")
                elif value.lower() == "false":
                    # Have to add a continue for false statements, else it adds -option 'False'
                    continue
                else:
                    this_option = (
                        module.advanced.option_format_string.replace(
                            "{{ KEY }}", str(key)
                        )
                        .replace("{{KEY}}", str(key))
                        .replace("{{ VALUE }}", str(value))
                        .replace("{{VALUE}}", str(value))
                    )
                    option_strings.append(f"{this_option}")

        script_end = (
            script_end.replace("{{ PARAMS }}", " ".join(option_strings))
            .replace("{{PARAMS}}", " ".join(option_strings))
            .replace(
                "{{ OUTPUT_FUNCTION }}", params.get("OutputFunction", "Out-String")
            )
            .replace("{{OUTPUT_FUNCTION}}", params.get("OutputFunction", "Out-String"))
        )

        # When obfuscation is enabled, the script source is already
        # obfuscated (from get_module_source or inline obfuscation above),
        # so tell finalize_module to only obfuscate the invoke command.
        return self.finalize_module(
            script=script,
            script_end=script_end,
            obfuscate=obfuscate,
            obfuscation_command=obfuscate_command,
            script_already_obfuscated=obfuscate,
        )

    def generate_script_csharp(
        self,
        module: EmpireModule,
        params: dict,
        obfuscation_config: models.ObfuscationConfig,
    ) -> ModuleExecutionRequest:
        try:
            obfuscate = (
                obfuscation_config.enabled if obfuscation_config is not None else False
            )
            script_file = self.dotnet_compiler.compile_task(
                module.compiler_yaml,
                module.name,
                dot_net_version=params["DotNetVersion"].lower(),
                confuse=obfuscate,
            )
            filtered_params = {}
            for key, value in params.items():
                if (
                    key.lower() not in ["agent", "dotnetversion"]
                    and value
                    and value != ""
                ):
                    if key.lower() == "file":
                        base64_assembly = value.get_base64_file()
                        filtered_params[key] = base64_assembly
                    else:
                        filtered_params[key] = value

            param_json = json.dumps(filtered_params)
            base64_json = base64.b64encode(param_json.encode("utf-8")).decode("utf-8")
            return ModuleExecutionRequest(
                command="",
                data=f"{script_file}|,{base64_json}",
                files=[script_file],
            )
        except (ModuleValidationException, ModuleExecutionException):
            raise
        except Exception as e:
            log.exception("dotnet compile error")
            raise ModuleExecutionException("dotnet compile error") from e

    def _create_modified_module(self, module: EmpireModule, modified_input: str):
        """
        Return a copy of the original module with the input modified.
        """
        modified_module = module.model_copy(deep=True)
        modified_module.script = modified_input
        modified_module.script_path = None

        if modified_module.language == LanguageEnum.csharp:
            compiler_dict = yaml.load(modified_module.compiler_yaml, Loader=Loader)
            compiler_dict[0]["Code"] = modified_input
            modified_module.compiler_yaml = yaml.dump(compiler_dict, Dumper=Dumper)

        return modified_module

    def load_modules(self, db: Session):
        root_path = self.main_menu.install_path / "modules"
        log.info(f"v2: Loading modules from: {root_path}")

        # Pre-load all existing module records to avoid per-module DB queries
        existing_modules = {
            mod.id: mod for mod in db.scalars(select(models.Module)).all()
        }

        for file_path in root_path.rglob("*.y*ml"):
            filename = file_path.name
            if fnmatch.fnmatch(filename, "*template.yaml"):
                continue

            # instantiate the module and save it to the internal cache
            try:
                yaml2 = yaml.load(file_path.read_text(), Loader=Loader)
                yaml_module = {k: v for k, v in yaml2.items() if v is not None}
                self._load_module(
                    db, yaml_module, root_path, file_path, existing_modules
                )
            except Exception:
                log.exception(f"Error loading module {filename}")

    def _load_module(  # noqa: PLR0912
        self,
        db: Session,
        yaml_module,
        root_path: Path,
        file_path: Path,
        existing_modules: dict | None = None,
    ):
        module_name = file_path.relative_to(root_path).with_suffix("").as_posix()
        yaml_module["techniques"].extend(
            self._get_interpreter_technique(yaml_module["language"])
        )

        if yaml_module["language"] == "csharp":
            yaml_module["id"] = slugify(module_name)

            # TODO: Remove this from EmpireCompiler so we dont need to build all the extra unused fields
            dict_yaml = yaml_module.get("csharp", {}).copy()
            dict_yaml.update(
                {
                    "Name": yaml_module.get("name", ""),
                    "Language": yaml_module.get("language", ""),
                    "TokenTask": False,
                }
            )

            dict_yaml["ReferenceSourceLibraries"] = [
                {
                    "Name": ref_lib.get("Name", ""),
                    "Description": ref_lib.get("Description", ""),
                    "Location": ref_lib.get("Location", ""),
                    "Language": ref_lib.get("Language", "CSharp"),
                    "CompatibleDotNetVersions": ref_lib.get(
                        "CompatibleDotNetVersions", []
                    ),
                    "ReferenceAssemblies": ref_lib.get("ReferenceAssemblies", []),
                    "EmbeddedResources": ref_lib.get("EmbeddedResources", []),
                }
                for ref_lib in dict_yaml.get("ReferenceSourceLibraries", [])
            ]

            compiler_yaml = yaml.dump(
                [dict_yaml],
                sort_keys=False,
                default_flow_style=False,
                allow_unicode=True,
                Dumper=Dumper,
            )

            my_model = EmpireModule(**yaml_module)
            my_model.compiler_yaml = compiler_yaml

            my_model.options.append(
                EmpireModuleOption(
                    name="DotNetVersion",
                    value=my_model.csharp.CompatibleDotNetVersions[0],
                    description=".NET version to compile against",
                    required=True,
                    suggested_values=my_model.csharp.CompatibleDotNetVersions,
                    strict=True,
                )
            )
        else:
            yaml_module["id"] = slugify(module_name)
            my_model = EmpireModule(**yaml_module)

        if my_model.advanced.custom_generate:
            custom_py = file_path.with_suffix(".py")
            if not custom_py.exists():
                raise Exception("No File to use for custom generate.")
            # Defer the importlib + Module() construction to first execute.
            # See _load_custom_generate_class below.
            my_model.advanced.custom_generate_path = str(custom_py)
        elif my_model.script_path:
            script_path = self.module_source_path / my_model.script_path
            if not script_path.exists():
                raise Exception(
                    f"File provided in script_path does not exist: {module_name}"
                )
        elif my_model.script:
            pass
        elif my_model.language == LanguageEnum.bof:
            if not (self.module_source_path / my_model.bof.x86).exists():
                raise Exception(f"x86 bof file provided does not exist: {module_name}")
            if not (self.module_source_path / my_model.bof.x64).exists():
                raise Exception(f"x64 bof file provided does not exist: {module_name}")
        elif my_model.language == LanguageEnum.csharp:
            pass
        else:
            raise Exception(
                "Must provide a valid script, script_path, or custom generate function"
            )

        if existing_modules is not None:
            mod = existing_modules.get(my_model.id)
        else:
            mod = db.scalars(
                select(models.Module).where(models.Module.id == my_model.id)
            ).first()

        if not mod:
            mod = models.Module(
                id=my_model.id,
                name=module_name,
                enabled=True,
                tactic=my_model.tactics,
                technique=my_model.techniques,
                software=my_model.software,
            )
            db.add(mod)

        self.modules[slugify(module_name)] = my_model
        self.modules[slugify(module_name)].enabled = mod.enabled

    def _get_interpreter_technique(self, language):
        if language == LanguageEnum.powershell:
            return ["T1059.001"]
        if language == LanguageEnum.csharp:
            return ["T1620"]
        if language == LanguageEnum.python:
            return ["T1059.006"]
        if language == LanguageEnum.ironpython:
            return ["T1059.006", "T1620"]
        if language == LanguageEnum.bof:
            return ["T1620"]
        return []

    def get_module_script(self, module_id: str):
        mod: EmpireModule = self.modules.get(module_id)

        if not mod:
            return None

        if mod.script_path:
            script_path = self.module_source_path / mod.script_path
            script = script_path.read_text()
        else:
            script = mod.script

        return script

    def get_module_source(
        self, module_name: str, obfuscate: bool = False, obfuscate_command: str = ""
    ) -> tuple[str | None, str | None]:
        """
        Get the obfuscated/unobfuscated module source code.
        """
        try:
            if obfuscate:
                module_path = self._obfuscated_module_source_path / module_name
                # If pre-obfuscated module exists then return code
                if module_path.exists():
                    return module_path.read_text(), None

                # If pre-obfuscated module does not exist then generate obfuscated code and return it
                module_path = self.module_source_path / module_name
                module_code = module_path.read_text()
                obfuscated_module_code = self.obfuscation_service.obfuscate(
                    module_code, obfuscate_command
                )
                return obfuscated_module_code, None

            # Use regular/unobfuscated code
            module_path = self.module_source_path / module_name
            module_code = module_path.read_text()
        except Exception:
            return (
                None,
                f"[!] Could not read module source path at: {self.module_source_path}",
            )
        else:
            return module_code, None

    def preobfuscate_modules(self, language: str, reobfuscate=False):
        """
        Preobfuscate PowerShell module_source files
        """
        if not data_util.is_powershell_installed():
            err = "PowerShell is not installed and is required to use obfuscation, please install it first."
            log.error(err)
            return err

        with SessionLocal.begin() as db:
            db_obf_config = self.obfuscation_service.get_obfuscation_config(
                db, language
            )
            files = self._get_module_source_files()

            for file in files:
                if reobfuscate or not self.is_obfuscated(file):
                    message = f"Obfuscating {file.name}..."
                    log.info(message)
                else:
                    log.warning(
                        f"{file.name} was already obfuscated. Not reobfuscating."
                    )
                self.obfuscate_module(file, db_obf_config.command, reobfuscate)
            return None

    def preobfuscate_module_by_id(
        self, module_id: str, reobfuscate=False
    ) -> str | None:
        """Pre-obfuscate a single module's source file.

        Follows the same pattern as ``preobfuscate_modules`` but for a
        single module looked up by ID.  Returns None on success or an
        error message string.
        """
        module = self.get_by_id(module_id)
        if not module:
            log.error("Pre-obfuscation: module not found: %s", module_id)
            return f"Module not found: {module_id}"
        if not module.script_path:
            log.warning("Pre-obfuscation: %s has no script_path", module_id)
            return f"Module {module_id} has no script_path"

        source_path = self.module_source_path / module.script_path
        with SessionLocal.begin() as db:
            config = self.obfuscation_service.get_obfuscation_config(
                db, module.language
            )
            if not config:
                log.error(
                    "Pre-obfuscation: no obfuscation config for language %s",
                    module.language,
                )
                return f"No obfuscation config for language: {module.language}"
            obfuscation_command = config.command

        log.info("Pre-obfuscating module %s (%s)", module_id, module.script_path)
        self.obfuscate_module(source_path, obfuscation_command, reobfuscate)
        return None

    # this is still written in a way that its only used for PowerShell
    # to make it work for other languages, we probably want to just pass in the db_obf_config
    # and delegate to language specific functions
    def obfuscate_module(
        self, module_source: Path, obfuscation_command="", force_reobfuscation=False
    ):
        if self.is_obfuscated(module_source) and not force_reobfuscation:
            return None

        try:
            module_code = module_source.read_text()
        except Exception:
            log.exception(f"Could not read module source path at: {module_source}")
            return ""

        # Get the random function name generated at install and patch the stager with the proper function name
        module_code = self.obfuscation_service.obfuscate_keywords(module_code)

        # obfuscate and write to obfuscated source path
        obfuscated_code = self.obfuscation_service.obfuscate(
            module_code, obfuscation_command
        )

        relative_path = module_source.relative_to(self.module_source_path)
        obfuscated_source = self._obfuscated_module_source_path / relative_path

        try:
            obfuscated_source.parent.mkdir(parents=True, exist_ok=True)
            obfuscated_source.write_text(obfuscated_code)
        except Exception:
            log.exception(
                f"Could not write obfuscated module source path at: {obfuscated_source}"
            )
            return ""

    def is_obfuscated(self, module_source: Path):
        # Get the file path of the module_source, but only relative to the module_source directory
        # Then append to the obfuscated_module_source directory
        relative_path = module_source.relative_to(self.module_source_path)
        return (self._obfuscated_module_source_path / relative_path).exists()

    def _get_module_source_files(self) -> list[Path]:
        paths = []
        pattern = "*.ps1"
        for root, _dirs, files in self.module_source_path.walk():
            for filename in fnmatch.filter(files, pattern):
                paths.append(root / filename)

        return paths

    def remove_preobfuscated_modules(self, _language: str):
        shutil.rmtree(self._obfuscated_module_source_path, ignore_errors=True)

    def finalize_module(
        self,
        script: str,
        script_end: str,
        obfuscate: bool = False,
        obfuscation_command: str = "",
        script_already_obfuscated: bool = False,
    ) -> str:
        """Combine script and script_end with obfuscation if needed.

        When ``script_already_obfuscated`` is True, the script body has
        already been through Invoke-Obfuscation (via ``get_module_source``
        or ``auto_get_source``), so only ``script_end`` is obfuscated.
        When False (the default), the full combined script is obfuscated
        as a single unit — this is the safe default for any caller that
        constructs a script without pre-obfuscating it (e.g.,
        custom_generate modules that call ``finalize_module`` directly).

        Keyword obfuscation (lightweight string replacement) is always
        applied to the full combined result regardless.
        """
        if "PowerSploit File: PowerView.ps1" in script:
            module_name = script_end.lstrip().split(" ")[0]
            script = helpers.generate_dynamic_powershell_script(script, module_name)

        if obfuscate:
            if script_already_obfuscated:
                script_end = self.obfuscation_service.obfuscate(
                    script_end, obfuscation_command
                )
            else:
                script += script_end
                script = self.obfuscation_service.obfuscate(script, obfuscation_command)
                return self.obfuscation_service.obfuscate_keywords(script)

        script += script_end
        return self.obfuscation_service.obfuscate_keywords(script)

    def delete_all_modules(self, db: Session):
        for module in list(self.modules.values()):
            db_module: models.Module = db.scalars(
                select(models.Module).where(models.Module.id == module.id)
            ).first()
            if db_module:
                db.delete(db_module)
            del self.modules[module.id]
        db.flush()


def auto_get_source(func):
    def wrapper(*args, **kwargs):
        main_menu = args[0]
        module = args[1]
        obfuscate = args[3]
        obfuscation_command = args[4]

        script, err = main_menu.modulesv2.get_module_source(
            module_name=module.script_path,
            obfuscate=obfuscate,
            obfuscate_command=obfuscation_command,
        )

        if err:
            raise ModuleValidationException(err)

        return func(*args, script=script, **kwargs)

    return wrapper


def auto_finalize(func):
    def wrapper(*args, **kwargs):
        script, script_end = func(*args, **kwargs)

        main_menu = args[0]
        obfuscate = args[3]
        obfuscation_command = args[4]

        return main_menu.modulesv2.finalize_module(
            script=script,
            script_end=script_end,
            obfuscate=obfuscate,
            obfuscation_command=obfuscation_command,
            script_already_obfuscated=obfuscate,
        )

    return wrapper
