#!/usr/bin/env python3
"""
patch_anti_debug.py - Automatically patch common anti-debugging checks in PE binaries.

Detects and patches:
  - IsDebuggerPresent calls
  - NtQueryInformationProcess (ProcessDebugPort, ProcessDebugObjectHandle)
  - PEB.BeingDebugged direct access
  - Timing checks (GetTickCount, QueryPerformanceCounter, rdtsc)
  - NtQuerySystemInformation (SystemKernelDebuggerInformation)
  - CheckRemoteDebuggerPresent

Always preserves the original binary (--no-backup to disable).

Usage:
    python3 patch_anti_debug.py --input sample.exe --output sample_patched.exe
    python3 patch_anti_debug.py --input sample.exe --dry-run --verbose
    python3 patch_anti_debug.py --input sample.exe --techniques isdebuggerpresent,peb
"""

from __future__ import annotations

import argparse
import logging
import os
import shutil
import struct
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Patch definitions
# ---------------------------------------------------------------------------

NOP = b"\x90"
RET_0_32 = b"\x31\xc0\xc3"         # xor eax, eax; ret
RET_0_PAD4 = b"\x31\xc0\xc3\x90"   # xor eax, eax; ret; nop


@dataclass
class PatchResult:
    """Result of a single patch operation."""
    technique: str
    offset: int
    original: bytes
    patched: bytes
    description: str


@dataclass
class PatchReport:
    """Aggregated report of all patches applied."""
    input_file: str
    output_file: str
    patches: list = field(default_factory=list)
    errors: list = field(default_factory=list)

    @property
    def total_patches(self) -> int:
        return len(self.patches)

    def summary(self) -> str:
        lines = [
            f"Patch Report for: {self.input_file}",
            f"Output: {self.output_file}",
            f"Total patches applied: {self.total_patches}",
            "",
        ]
        for p in self.patches:
            lines.append(
                f"  [{p.technique}] offset 0x{p.offset:08x}: {p.description}"
            )
        if self.errors:
            lines.append("")
            lines.append("Errors:")
            for e in self.errors:
                lines.append(f"  - {e}")
        return "\n".join(lines)


# ---------------------------------------------------------------------------
# PE helper (minimal, no pefile dependency required)
# ---------------------------------------------------------------------------

class MinimalPE:
    """Minimal PE parser sufficient for locating imports and code patterns."""

    def __init__(self, data: bytes):
        self.data = bytearray(data)
        self._parse_headers()

    def _parse_headers(self):
        if self.data[:2] != b"MZ":
            raise ValueError("Not a valid PE file (missing MZ header)")

        pe_offset = struct.unpack_from("<I", self.data, 0x3C)[0]
        if self.data[pe_offset:pe_offset + 4] != b"PE\x00\x00":
            raise ValueError("Not a valid PE file (missing PE signature)")

        self.pe_offset = pe_offset
        magic_offset = pe_offset + 0x18
        self.magic = struct.unpack_from("<H", self.data, magic_offset)[0]
        self.is_64bit = self.magic == 0x20B

        # Number of sections
        self.num_sections = struct.unpack_from("<H", self.data, pe_offset + 6)[0]

        # Optional header size
        self.opt_header_size = struct.unpack_from("<H", self.data, pe_offset + 0x14)[0]

        # Section table offset
        self.section_table_offset = pe_offset + 0x18 + self.opt_header_size

        # Image base
        if self.is_64bit:
            self.image_base = struct.unpack_from("<Q", self.data, magic_offset + 0x18)[0]
        else:
            self.image_base = struct.unpack_from("<I", self.data, magic_offset + 0x1C)[0]

        self._parse_sections()
        self._parse_imports()

    def _parse_sections(self):
        self.sections = []
        offset = self.section_table_offset
        for _ in range(self.num_sections):
            name = self.data[offset:offset + 8].rstrip(b"\x00").decode("ascii", errors="replace")
            virtual_size = struct.unpack_from("<I", self.data, offset + 8)[0]
            virtual_addr = struct.unpack_from("<I", self.data, offset + 12)[0]
            raw_size = struct.unpack_from("<I", self.data, offset + 16)[0]
            raw_offset = struct.unpack_from("<I", self.data, offset + 20)[0]
            self.sections.append({
                "name": name,
                "virtual_size": virtual_size,
                "virtual_address": virtual_addr,
                "raw_size": raw_size,
                "raw_offset": raw_offset,
            })
            offset += 40

    def rva_to_offset(self, rva: int) -> Optional[int]:
        for s in self.sections:
            if s["virtual_address"] <= rva < s["virtual_address"] + s["raw_size"]:
                return rva - s["virtual_address"] + s["raw_offset"]
        return None

    def _parse_imports(self):
        """Extract imported function names for anti-debug detection."""
        self.imports = {}
        try:
            import pefile
            pe = pefile.PE(data=bytes(self.data))
            if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
                for entry in pe.DIRECTORY_ENTRY_IMPORT:
                    dll = entry.dll.decode("ascii", errors="replace").lower()
                    for imp in entry.imports:
                        if imp.name:
                            name = imp.name.decode("ascii", errors="replace")
                            self.imports[name] = {
                                "dll": dll,
                                "thunk_rva": imp.address - self.image_base if imp.address else 0,
                            }
        except ImportError:
            logger.debug("pefile not available; falling back to string-based detection")
            self._find_import_strings()
        except Exception as e:
            logger.debug(f"Import parsing failed: {e}; falling back to string-based detection")
            self._find_import_strings()

    def _find_import_strings(self):
        """Fallback: search for API name strings in the binary."""
        api_names = [
            b"IsDebuggerPresent",
            b"CheckRemoteDebuggerPresent",
            b"NtQueryInformationProcess",
            b"NtQuerySystemInformation",
            b"GetTickCount",
            b"QueryPerformanceCounter",
            b"OutputDebugString",
            b"NtSetInformationThread",
        ]
        for api in api_names:
            idx = self.data.find(api)
            if idx != -1:
                self.imports[api.decode()] = {"dll": "unknown", "offset": idx}

    def find_all(self, pattern: bytes) -> list:
        """Find all occurrences of a byte pattern."""
        results = []
        start = 0
        while True:
            idx = self.data.find(pattern, start)
            if idx == -1:
                break
            results.append(idx)
            start = idx + 1
        return results

    def patch_bytes(self, offset: int, new_bytes: bytes) -> None:
        """Patch bytes at the given offset."""
        self.data[offset:offset + len(new_bytes)] = new_bytes

    def get_bytes(self) -> bytes:
        return bytes(self.data)


# ---------------------------------------------------------------------------
# Patching techniques
# ---------------------------------------------------------------------------

def patch_isdebuggerpresent(pe: MinimalPE, report: PatchReport, dry_run: bool = False) -> None:
    """Patch calls to IsDebuggerPresent to always return 0."""
    # Pattern: FF 15 xx xx xx xx (call [IsDebuggerPresent])
    # We NOP the call and set eax = 0 instead

    if "IsDebuggerPresent" in pe.imports:
        logger.info("IsDebuggerPresent import detected")

    # Search for common call patterns followed by test eax, eax
    # call dword ptr [addr]; test eax, eax  =>  xor eax, eax; nop...; test eax, eax
    patterns_found = 0

    # Pattern 1: FF 15 (indirect call via IAT)
    for offset in pe.find_all(b"\xFF\x15"):
        # Verify this is plausibly an IsDebuggerPresent call
        # Check if the bytes after the 6-byte call instruction include test eax,eax (85 C0)
        after_call = pe.data[offset + 6:offset + 12]
        if b"\x85\xc0" in after_call:
            desc = "IsDebuggerPresent indirect call -> xor eax,eax + NOPs"
            original = bytes(pe.data[offset:offset + 6])
            patch = b"\x31\xc0" + NOP * 4  # xor eax, eax; nop; nop; nop; nop
            if not dry_run:
                pe.patch_bytes(offset, patch)
            report.patches.append(PatchResult(
                technique="isdebuggerpresent",
                offset=offset,
                original=original,
                patched=patch,
                description=desc,
            ))
            patterns_found += 1
            if patterns_found >= 10:  # Safety limit
                break

    # Pattern 2: Direct PEB access (fs:[30h] on x86, gs:[60h] on x64)
    # 64 A1 30 00 00 00 (mov eax, fs:[0x30]) - x86 PEB access
    for offset in pe.find_all(b"\x64\xA1\x30\x00\x00\x00"):
        # Check if followed by movzx/mov of BeingDebugged (offset 2)
        after = pe.data[offset + 6:offset + 10]
        if b"\x02" in after:
            desc = "PEB.BeingDebugged via fs:[30h] -> xor eax,eax + NOPs"
            region = bytes(pe.data[offset:offset + 10])
            patch = b"\x31\xc0" + NOP * 8
            if not dry_run:
                pe.patch_bytes(offset, patch)
            report.patches.append(PatchResult(
                technique="isdebuggerpresent",
                offset=offset,
                original=region,
                patched=patch,
                description=desc,
            ))

    if patterns_found == 0:
        logger.info("No IsDebuggerPresent call patterns found")


def patch_peb_checks(pe: MinimalPE, report: PatchReport, dry_run: bool = False) -> None:
    """Patch direct PEB flag checks (BeingDebugged, NtGlobalFlag, Heap Flags)."""
    # PEB access via fs:[30h] (32-bit) or gs:[60h] (64-bit)
    peb_patterns = [
        # 32-bit: mov eax, dword ptr fs:[30h]
        (b"\x64\xA1\x30\x00\x00\x00", 6, "PEB access via fs:[30h]"),
        # 32-bit: mov reg, dword ptr fs:[30h] (various registers)
        (b"\x64\x8B\x0D\x30\x00\x00\x00", 7, "PEB access via fs:[30h] to ecx"),
        (b"\x64\x8B\x15\x30\x00\x00\x00", 7, "PEB access via fs:[30h] to edx"),
        # 64-bit: mov rax, qword ptr gs:[60h]
        (b"\x65\x48\x8B\x04\x25\x60\x00\x00\x00", 9, "PEB access via gs:[60h]"),
    ]

    for pattern, length, desc in peb_patterns:
        for offset in pe.find_all(pattern):
            # Check next ~10 bytes for NtGlobalFlag offset (0x68 for 32-bit, 0xBC for 64-bit)
            after = pe.data[offset + length:offset + length + 15]

            # NtGlobalFlag check
            if b"\x68" in after[:8] or b"\xBC" in after[:8]:
                full_desc = f"{desc} -> NtGlobalFlag check neutralized"
                original = bytes(pe.data[offset:offset + length])
                patch = b"\x31\xc0" + NOP * (length - 2)
                if not dry_run:
                    pe.patch_bytes(offset, patch)
                report.patches.append(PatchResult(
                    technique="peb",
                    offset=offset,
                    original=original,
                    patched=patch,
                    description=full_desc,
                ))


def patch_timing_checks(pe: MinimalPE, report: PatchReport, dry_run: bool = False) -> None:
    """Patch timing-based anti-debug checks."""
    # rdtsc instruction: 0F 31
    for offset in pe.find_all(b"\x0F\x31"):
        # Check context: if there are two rdtsc within ~100 bytes, it's likely a timing check
        nearby = pe.data[offset + 2:offset + 150]
        if b"\x0F\x31" in nearby:
            desc = "rdtsc timing check -> NOP"
            original = bytes(pe.data[offset:offset + 2])
            patch = NOP * 2
            if not dry_run:
                pe.patch_bytes(offset, patch)
            report.patches.append(PatchResult(
                technique="timing",
                offset=offset,
                original=original,
                patched=patch,
                description=desc,
            ))

    # GetTickCount calls followed by comparison
    if "GetTickCount" in pe.imports:
        logger.info("GetTickCount import detected - look for call sites manually")
        for offset in pe.find_all(b"\xFF\x15"):
            # After call, check for sub/cmp pattern within 20 bytes
            after = pe.data[offset + 6:offset + 26]
            if b"\x2B" in after or b"\x3D" in after:
                # Could be GetTickCount delta check
                desc = "Potential GetTickCount timing call -> xor eax,eax + NOPs"
                original = bytes(pe.data[offset:offset + 6])
                # Only patch if we haven't already patched this offset
                already = any(p.offset == offset for p in report.patches)
                if not already:
                    patch = b"\x31\xc0" + NOP * 4
                    if not dry_run:
                        pe.patch_bytes(offset, patch)
                    report.patches.append(PatchResult(
                        technique="timing",
                        offset=offset,
                        original=original,
                        patched=patch,
                        description=desc,
                    ))


def patch_ntquery(pe: MinimalPE, report: PatchReport, dry_run: bool = False) -> None:
    """Patch NtQueryInformationProcess debug checks."""
    if "NtQueryInformationProcess" not in pe.imports:
        logger.info("NtQueryInformationProcess not imported, skipping")
        return

    # Look for push 7 (ProcessDebugPort) or push 0x1E (ProcessDebugObjectHandle)
    # followed by a call
    debug_port_patterns = [
        (b"\x6A\x07", "ProcessDebugPort (push 7)"),
        (b"\x6A\x1E", "ProcessDebugObjectHandle (push 0x1E)"),
        (b"\x6A\x1F", "ProcessDebugFlags (push 0x1F)"),
    ]

    for pattern, desc in debug_port_patterns:
        for offset in pe.find_all(pattern):
            # Check if followed within 20 bytes by a call (FF 15 or E8)
            after = pe.data[offset + 2:offset + 22]
            if b"\xFF\x15" in after or b"\xE8" in after:
                full_desc = f"NtQueryInformationProcess {desc} -> NOP push"
                original = bytes(pe.data[offset:offset + 2])
                patch = NOP * 2
                if not dry_run:
                    pe.patch_bytes(offset, patch)
                report.patches.append(PatchResult(
                    technique="ntquery",
                    offset=offset,
                    original=original,
                    patched=patch,
                    description=full_desc,
                ))


def patch_checkremote(pe: MinimalPE, report: PatchReport, dry_run: bool = False) -> None:
    """Patch CheckRemoteDebuggerPresent calls."""
    if "CheckRemoteDebuggerPresent" not in pe.imports:
        return

    logger.info("CheckRemoteDebuggerPresent import detected")
    # Similar to IsDebuggerPresent - NOP the call and zero eax
    for offset in pe.find_all(b"\xFF\x15"):
        after = pe.data[offset + 6:offset + 16]
        # Check for test/cmp on the output parameter
        if b"\x85" in after or b"\x83" in after:
            desc = "CheckRemoteDebuggerPresent call -> xor eax,eax + NOPs"
            original = bytes(pe.data[offset:offset + 6])
            already = any(p.offset == offset for p in report.patches)
            if not already:
                patch = b"\x31\xc0" + NOP * 4
                if not dry_run:
                    pe.patch_bytes(offset, patch)
                report.patches.append(PatchResult(
                    technique="checkremote",
                    offset=offset,
                    original=original,
                    patched=patch,
                    description=desc,
                ))
                break  # Usually only one call


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

TECHNIQUES = {
    "isdebuggerpresent": patch_isdebuggerpresent,
    "peb": patch_peb_checks,
    "timing": patch_timing_checks,
    "ntquery": patch_ntquery,
    "checkremote": patch_checkremote,
}


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Automatically patch anti-debugging checks in PE binaries.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --input sample.exe --output sample_patched.exe
  %(prog)s --input sample.exe --dry-run --verbose
  %(prog)s --input sample.exe --techniques isdebuggerpresent,peb
        """,
    )
    parser.add_argument("--input", "-i", required=True, help="Input PE file")
    parser.add_argument("--output", "-o", help="Output patched PE file (default: <input>_patched.exe)")
    parser.add_argument(
        "--techniques", "-t",
        default="all",
        help=f"Comma-separated techniques to patch: {','.join(TECHNIQUES.keys())} or 'all' (default: all)",
    )
    parser.add_argument("--backup", action="store_true", default=True, help="Create backup of original (default: true)")
    parser.add_argument("--no-backup", action="store_true", help="Do not create a backup")
    parser.add_argument("--dry-run", action="store_true", help="Show patches without modifying file")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    parser.add_argument("--format", default="json", choices=["json", "text", "csv"],
                        help="Output format (default: json)")
    return parser.parse_args()


def main() -> None:
    args = parse_arguments()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s: %(message)s",
    )

    input_path = Path(args.input)
    if not input_path.is_file():
        logger.error(f"Input file not found: {input_path}")
        sys.exit(1)

    output_path = Path(args.output) if args.output else input_path.with_stem(input_path.stem + "_patched")

    # Read input
    data = input_path.read_bytes()
    logger.info(f"Loaded {input_path} ({len(data)} bytes)")

    # Parse PE
    try:
        pe = MinimalPE(data)
    except ValueError as e:
        logger.error(f"Failed to parse PE: {e}")
        sys.exit(1)

    logger.info(f"PE format: {'PE32+' if pe.is_64bit else 'PE32'}")
    logger.info(f"Sections: {', '.join(s['name'] for s in pe.sections)}")
    logger.info(f"Known imports: {len(pe.imports)}")

    # Select techniques
    if args.techniques == "all":
        selected = list(TECHNIQUES.keys())
    else:
        selected = [t.strip().lower() for t in args.techniques.split(",")]
        for t in selected:
            if t not in TECHNIQUES:
                logger.error(f"Unknown technique: {t}. Available: {', '.join(TECHNIQUES.keys())}")
                sys.exit(1)

    # Apply patches
    report = PatchReport(input_file=str(input_path), output_file=str(output_path))

    for tech_name in selected:
        logger.info(f"Scanning for {tech_name}...")
        try:
            TECHNIQUES[tech_name](pe, report, dry_run=args.dry_run)
        except Exception as e:
            msg = f"Error in {tech_name}: {e}"
            logger.warning(msg)
            report.errors.append(msg)

    # Report
    print()
    print(report.summary())
    print()

    if args.dry_run:
        logger.info("Dry run complete. No files modified.")
        return

    if report.total_patches == 0:
        logger.info("No patches applied. The binary may not contain recognized anti-debug patterns.")
        return

    # Backup
    if args.backup and not args.no_backup:
        backup_path = input_path.with_suffix(input_path.suffix + ".bak")
        shutil.copy2(input_path, backup_path)
        logger.info(f"Backup saved to {backup_path}")

    # Write output
    output_path.write_bytes(pe.get_bytes())
    logger.info(f"Patched binary written to {output_path}")


if __name__ == "__main__":
    main()