Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| import re | |
| import random | |
| import time | |
| import ollama | |
| from langfuse.openai import openai | |
| import xml.etree.ElementTree as ET | |
| client = ollama.Client( | |
| host="https://owl-light-racer.ngrok-free.app" | |
| ) | |
| 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( | |
| model="bilel_cherif/Dante-mlwr", | |
| messages=[ | |
| {"role": "user", "content": prompt} | |
| ], | |
| options={ | |
| "temperature": 1.0, | |
| "top_p": 0.95, | |
| "num_predict": 30 * 1024, | |
| } | |
| ) | |
| end_time = time.time() | |
| content = response["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 = """Pick an example and run to infer the model""" | |
| 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="GRPO Shell generation Demo", | |
| description=description, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |