import gradio as gr

import os
import re
import random
import time

from langfuse.openai import openai
import xml.etree.ElementTree as ET


RUNPOD_URL = os.getenv("RUNPOD_URL")
RUNPOD_API = os.getenv("RUNPOD_API")

if not RUNPOD_URL or not RUNPOD_API:
    raise ValueError("RUNPOD_URL and RUNPOD_API environment variables must be set.")


client = openai.OpenAI(base_url=RUNPOD_URL, api_key=RUNPOD_API, timeout=360.0, max_retries=2)


ALLOCATION_METHODS = {
    "VirtualAlloc": "using VirtualAlloc, or something similar",
    "Executable heap": "on an executable heap",
    "Mapped sections": "using mapped sections",
    "Dripped pages": "using dripped (smaller, sequential) pages",
}

EXECUTION_METHODS = {
    "In-process without creating a thread": "in-process (without creating a thread)",
    "Creating a new thread in-process": "in-process with a new thread",
    "Queueing an APC in-process": "in-process, by queueing an APC",
    "Creating a new thread in a remote process": "by creating a new thread in a remote process",
    "Early Bird injection": "in a remote process, with a technique like earlybird injection",
}

OUTPUT_FORMATS = {
    "EXE": "an EXE",
    "CPL": "a CPL (DLL with CplApplet export)",
    "XLL": "an XLL (DLL with xlAutoOpen export)",
}

TRANSFORMATION_METHODS = {
    "AES": "AES encryption",
    "RC4": "RC4 encryption",
    "Multi-byte XOR": "multi-byte XOR encryption",
    "Base64": "base64 encoding",
}

GUARDRAILS = {
    "Sandbox check": "sandbox",
    "Domain-joined host check": "domain-joined host check",
    "Debugger check": "debugger",
}


def generate_prompt(
    execution_method, allocation_method, transformation, guardrails, file_format
):
    script_name = random.choice(["prepare.py", "setup.py", "preprocess.py"])
    extension = random.choice(["c", "cpp"])
    generic_technique = random.choice(
        [
            "dynamic API resolution via hash",
            "NTDLL functions instead of kernel32/kernelbase exports",
        ]
    )
    misc = random.choice(
        [
            "static signatures/YARA rules",
            "suspicious API calls",
            "unusual memory protection flags",
            "unusual DLL loading",
        ]
    )

    execution_method = EXECUTION_METHODS[execution_method]
    allocation_method = ALLOCATION_METHODS[allocation_method]
    transformation = TRANSFORMATION_METHODS[transformation]
    guardrails = GUARDRAILS[guardrails]
    file_format = OUTPUT_FORMATS[file_format]

    return f"""Please reason step by step

You are an expert Windows security researcher and malware developer. You have been tasked with evading advanced AV/EDR software. Common API calls and simple encoding will not be sufficient. No enhancements can be made later, you must provide the best possible solution immediately. Your exact output will be compiled and tested in a Windows environment.

Try innovative and creative techniques to evade detection. The code must compile and run without errors. Your output will be evaluated based on functionality and evasion success against leading AV/EDR solutions.

Respond EXACTLY in this format:

<project>
  <src>
    <file name="main.{extension}">
<![CDATA[
// Primary loader implementing chosen techniques
]]>
    </file>

    <file name="{script_name}">
<![CDATA[
# Any encryption, encoding, or other helper functions to prepare source code
]]>
    </file>

    <file name="Makefile">
<![CDATA[
# Build instructions to prepare and compile the final payload using {"x86_64-w64-mingw32-gcc" if extension == "c" else "x86_64-w64-mingw32-g++"}
]]>
    </file>
  </src>

  <command>make</command>
</project>

You can read the shellcode with your Python script from "./shellcode.bin".


Develop a novel shellcode loader that:
- Allocates memory {allocation_method}
- Executes shellcode {execution_method}
- Obfuscates the shellcode with {transformation}
- Uses {generic_technique}
- Outputs {file_format}
- Avoids {misc}
- Includes {guardrails} checks"""


def replace_pragma_links(file_contents):
    pattern = r'#pragma\s+comment\s*\(\s*lib\s*,\s*"([^"]+?)(?:\.lib)?"\s*\)'
    matches = re.findall(pattern, file_contents)
    mingw_flags = [f"-l{lib}" for lib in matches]

    def replace_single_pragma(match):
        lib = match.group(1).replace(".lib", "")
        return f"// replaced pragma with -l{lib}\n"

    modified_contents = re.sub(pattern, replace_single_pragma, file_contents)
    modified_contents = re.sub(r"\n\s*\n", "\n\n", modified_contents)
    return modified_contents, mingw_flags


def update_makefile(makefile_contents, additional_mingw_flags):
    lines = makefile_contents.split("\n")
    modified_lines = []
    i = 0
    while i < len(lines):
        line = lines[i]
        if (
            re.search(r"(x86_64|i686)-w64-mingw32-g(cc|\+\+)", line)
            and not line.lstrip().startswith("#")
            and line.lstrip().find("=") == -1
        ):
            current_line = line
            continuation_lines = []
            while current_line.endswith("\\") and i + 1 < len(lines):
                i += 1
                continuation_lines.append(lines[i])
                current_line = lines[i]

            if continuation_lines:
                modified_lines.append(line)
                for j in range(len(continuation_lines) - 1):
                    modified_lines.append(continuation_lines[j])
                modified_lines.append(continuation_lines[-1] + additional_mingw_flags)
            else:
                modified_lines.append(line + additional_mingw_flags)
        else:
            modified_lines.append(line)
        i += 1
    return "\n".join(modified_lines)


def xml_to_markdown(project_xml):
    if not project_xml:
        raise ValueError("No <project> element found in XML")

    root = ET.fromstring(project_xml)

    command_elem = root.find("command")
    if command_elem is None:
        raise ValueError("No <command> element found in XML")

    command = command_elem.text
    if command is None:
        raise ValueError("No <command> element text found in XML")

    command = command.strip()

    src_elem = root.find("src")
    if src_elem is None:
        raise ValueError("No <src> element found in XML")

    files_info = []
    for file_elem in src_elem.findall("file"):
        file_name = file_elem.get("name")
        file_content = file_elem.text

        if file_content:
            file_content = file_content.strip()

            if "<![CDATA[" in file_content:
                cdata_start = file_content.find("<![CDATA[") + len("<![CDATA[")
                cdata_end = file_content.rfind("]]>")
                if cdata_start > 0 and cdata_end > cdata_start:
                    file_content = file_content[cdata_start:cdata_end]

            files_info.append({"name": file_name, "content": file_content})

    if len(files_info) == 0:
        raise ValueError("No <file> elements found in <src>")

    all_mingw_flags = []
    for file_info in files_info:
        content = file_info["content"]
        name = file_info["name"]

        content = content.replace("CplApplet", "CPlApplet")
        content = content.replace("CPLApplet", "CPlApplet")
        content = content.replace("cplapplet", "CPlApplet")
        content = content.replace("XlAutoOpen", "xlAutoOpen")
        content = content.replace("XLAutoOpen", "xlAutoOpen")
        content = content.replace("xlautoopen", "xlAutoOpen")

        content = content.replace("<Windows.h>", "<windows.h>")
        content = content.replace("<Wincrypt.h>", "<wincrypt.h>")
        content = content.replace("<Wininet.h>", "<wininet.h>")
        content = content.replace("<ShlObj.h>", "<shlobj.h>")
        content = content.replace("<Netapi32.h>", "<netapi32.h>")
        content = content.replace("<TlHelp32.h>", "<tlhelp32.h>")
        content = content.replace("<Cpl.h>", "<cpl.h>")
        content = content.replace("<Lm.h>", "<lm.h>")
        content = content.replace("<Iphlpapi.h>", "<iphlpapi.h>")
        content = content.replace("<Dsgetdc.h>", "<dsgetdc.h>")
        content = content.replace("<Psapi.h>", "<psapi.h>")
        content = content.replace("<Winternl.h>", "<winternl.h>")
        content = content.replace("#include <netapi.h>", "")

        content = content.replace("-Werror ", " ")
        content = content.replace("-Werror\n", "\n")

        content = content.replace("-l:crypt32.lib", "")
        content = content.replace("-l:ws2_32.lib", "")
        content = content.replace("-l:user32.lib", "")
        content = content.replace("-l:kernel32.lib", "")
        content = content.replace("-l:msvcrt.lib", "")
        content = content.replace("-l:ntdll.lib", "")
        content = content.replace("-l:mswsock.lib", "")
        content = content.replace("-l:psapi.lib", "")
        content = content.replace("-l:shlwapi.lib", "")
        content = content.replace("-l:comdlg32.lib", "")
        content = content.replace("-l:comctl32.lib", "")
        content = content.replace("-l:uuid.lib", "")
        content = content.replace("-l:ole32.lib", "")
        content = content.replace("-l:oleaut32.lib", "")
        content = content.replace("-l:gdi32.lib", "")
        content = content.replace("-l:iphlpapi.lib", "")
        content = content.replace("-l:libcrypto.a", "")
        content = content.replace("-l:libssl.a", "")

        if name == "Makefile":
            content = content.replace(
                "x86_64-w64-mingw32-gcc.exe", "x86_64-w64-mingw32-gcc"
            )
            content = content.replace(" gcc ", " x86_64-w64-mingw32-gcc ")
            content = content.replace("\tgcc ", "\tx86_64-w64-mingw32-gcc ")
            content = content.replace(
                "x86_64-w64-mingw32-g++.exe", "x86_64-w64-mingw32-g++"
            )
            content = content.replace(" g++ ", " x86_64-w64-mingw32-g++ ")
            content = content.replace("\tg++ ", "\tx86_64-w64-mingw32-g++ ")
            content = content.replace("-m64x64", "-m64")
            content = content.replace("    ", "\t")
            content = content.replace("SHELL := cmd.exe", "SHELL := /bin/sh")

        content, file_mingw_flags = replace_pragma_links(content)
        all_mingw_flags.extend(file_mingw_flags)

        file_info["content"] = content

    if len(all_mingw_flags) > 0:
        for file_info in files_info:
            if file_info["name"] == "Makefile":
                file_info["content"] = update_makefile(
                    file_info["content"], " " + " ".join(all_mingw_flags)
                )
                break

    markdown = f"Build Command: `{command}`\n\n"
    for file_info in files_info:
        markdown += f"{file_info['name']}:\n\n```\n{file_info['content']}\n```\n"

    return markdown.strip()


def call_infer(
    execution_method, allocation_method, transformation, guardrails, file_format
):
    prompt = generate_prompt(
        execution_method, allocation_method, transformation, guardrails, file_format
    )

    try:
        start_time = time.time()

        response = client.chat.completions.create(
            model="outflanknl/Dante-7B",
            messages=[{"role": "user", "content": prompt}],
            temperature=1.0,
            max_tokens=int(30 * 1024),
            top_p=0.95,
            stream=False,
        )

        end_time = time.time()

        content = response.choices[0].message.content

    except Exception:
        yield f"Error generating output. Please try again later."
        return

    try:
        project_xml = f"<project>{content.split('<project>', 1)[1].split('</project>', 1)[0]}</project>"

        markdown_output = xml_to_markdown(project_xml)
    except Exception:
        yield f"Error parsing output. Please try again later."
        return

    yield f"{markdown_output}\n\nGenerated in {end_time - start_time:.1f} seconds."
    return


description = """Note: This demo runs on a limited pool of RunPod serverless GPUs. You may experience delays when there is high demand.

## Resources

- Blog: https://outflank.nl/blog/2025/08/07/training-specialist-models
- Run the model yourself: https://huggingface.co/outflanknl/Dante-7B"""

demo = gr.Interface(
    fn=call_infer,
    inputs=[
        gr.Dropdown(
            list(EXECUTION_METHODS.keys()),
            label="Execution method",
            value="In-process without creating a thread",
        ),
        gr.Dropdown(
            list(ALLOCATION_METHODS.keys()),
            label="Allocation method",
            value="VirtualAlloc",
        ),
        gr.Dropdown(
            list(TRANSFORMATION_METHODS.keys()),
            label="Shellcode transformation",
            value="AES",
        ),
        gr.Dropdown(
            list(GUARDRAILS.keys()),
            label="Guardrails",
            value="Domain-joined host check",
        ),
        gr.Dropdown(list(OUTPUT_FORMATS.keys()), label="Output format", value="EXE"),
    ],
    outputs=gr.Markdown(),
    submit_btn="Generate",
    stop_btn="Cancel",
    clear_btn=None,
    theme=gr.themes.Default(
        primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.gray
    ),
    title="Dante-7B Demo",
    description=description,
)


if __name__ == "__main__":
    demo.launch()