#!/usr/bin/env python3
"""ApiOid PE static-triage helper for reverse-skill.

Pure-Python Windows PE triage using pefile (+ optional capstone). No external
RE tools required. Emits a compact JSON + human summary: identity, hashes,
sections (entropy), imports (categorized), notable strings, TLS/exports,
signature presence, packer hints.

Usage:  python omp-pe-triage.py <path-to.exe> [--strings-min 6] [--max-strings 60]
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import re
import sys

try:
    import pefile
except ImportError:
    print("ERR: pefile not installed (python -m pip install pefile)", file=sys.stderr)
    sys.exit(2)

# Import-name buckets that hint at behavior.
API_BUCKETS = {
    "process/inject": [
        "VirtualAlloc", "VirtualAllocEx", "WriteProcessMemory", "CreateRemoteThread",
        "NtCreateThreadEx", "QueueUserAPC", "SetThreadContext", "OpenProcess",
        "CreateProcess", "ShellExecute", "WinExec", "NtUnmapViewOfSection",
    ],
    "network": [
        "WSAStartup", "socket", "connect", "send", "recv", "InternetOpen",
        "InternetConnect", "HttpOpenRequest", "HttpSendRequest", "URLDownloadToFile",
        "WinHttpOpen", "WinHttpConnect", "getaddrinfo", "gethostbyname",
    ],
    "crypto": [
        "CryptAcquireContext", "CryptEncrypt", "CryptDecrypt", "CryptGenKey",
        "CryptHashData", "BCryptEncrypt", "BCryptDecrypt", "CryptStringToBinary",
    ],
    "filesystem": [
        "CreateFile", "WriteFile", "ReadFile", "DeleteFile", "MoveFile",
        "CopyFile", "FindFirstFile", "SHGetFolderPath", "GetTempPath",
    ],
    "registry": [
        "RegOpenKey", "RegSetValue", "RegCreateKey", "RegQueryValue", "RegDeleteKey",
    ],
    "persistence": [
        "CreateService", "OpenSCManager", "SetWindowsHookEx", "RegisterServiceCtrlHandler",
    ],
    "anti-analysis": [
        "IsDebuggerPresent", "CheckRemoteDebuggerPresent", "NtQueryInformationProcess",
        "GetTickCount", "QueryPerformanceCounter", "OutputDebugString", "Sleep",
    ],
    "privilege": [
        "AdjustTokenPrivileges", "OpenProcessToken", "LookupPrivilegeValue",
        "LogonUser", "ImpersonateLoggedOnUser",
    ],
}


def entropy(data: bytes) -> float:
    if not data:
        return 0.0
    counts = [0] * 256
    for b in data:
        counts[b] += 1
    n = len(data)
    ent = 0.0
    for c in counts:
        if c:
            p = c / n
            ent -= p * math.log2(p)
    return round(ent, 3)


def hashes(path: str) -> dict:
    h_md5, h_sha1, h_sha256 = hashlib.md5(), hashlib.sha1(), hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h_md5.update(chunk)
            h_sha1.update(chunk)
            h_sha256.update(chunk)
    return {"md5": h_md5.hexdigest(), "sha1": h_sha1.hexdigest(), "sha256": h_sha256.hexdigest()}


def ascii_strings(data: bytes, minlen: int):
    return re.findall(rb"[\x20-\x7e]{%d,}" % minlen, data)


def unicode_strings(data: bytes, minlen: int):
    return re.findall(rb"(?:[\x20-\x7e]\x00){%d,}" % minlen, data)


INTERESTING = re.compile(
    rb"(https?://|www\.|\.onion|[A-Za-z]:\\|/etc/|HKEY_|SOFTWARE\\|"
    rb"\.(exe|dll|dat|bin|sys|ps1|bat|vbs)\b|token|password|api[_-]?key|"
    rb"bearer|-----BEGIN|[0-9]{1,3}(?:\.[0-9]{1,3}){3})",
    re.IGNORECASE,
)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("path")
    ap.add_argument("--strings-min", type=int, default=6)
    ap.add_argument("--max-strings", type=int, default=60)
    args = ap.parse_args()

    out: dict = {"path": args.path}
    out["hashes"] = hashes(args.path)

    try:
        pe = pefile.PE(args.path, fast_load=False)
    except pefile.PEFormatError as e:
        print(f"ERR: not a valid PE: {e}", file=sys.stderr)
        return 2

    fh = pe.FILE_HEADER
    oh = pe.OPTIONAL_HEADER
    is64 = hex(oh.Magic) == "0x20b"
    out["identity"] = {
        "machine": hex(fh.Machine),
        "bits": 64 if is64 else 32,
        "subsystem": pefile.SUBSYSTEM_TYPE.get(oh.Subsystem, str(oh.Subsystem)),
        "dll": bool(fh.Characteristics & 0x2000),
        "entry_point": hex(oh.AddressOfEntryPoint),
        "image_base": hex(oh.ImageBase),
        "timestamp": fh.TimeDateStamp,
        "num_sections": fh.NumberOfSections,
    }

    out["sections"] = []
    for s in pe.sections:
        raw = s.get_data()
        name = s.Name.rstrip(b"\x00").decode("latin1", "replace")
        out["sections"].append({
            "name": name,
            "vsize": s.Misc_VirtualSize,
            "rawsize": s.SizeOfRawData,
            "entropy": entropy(raw),
            "high_entropy": entropy(raw) > 7.0,
        })

    # Imports (categorized)
    buckets = {k: [] for k in API_BUCKETS}
    all_imports = []
    dlls = []
    if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            dll = entry.dll.decode("latin1", "replace") if entry.dll else "?"
            dlls.append(dll)
            for imp in entry.imports:
                nm = imp.name.decode("latin1", "replace") if imp.name else f"ord_{imp.ordinal}"
                all_imports.append(nm)
                for bucket, needles in API_BUCKETS.items():
                    if any(nm.startswith(n) for n in needles):
                        buckets[bucket].append(nm)
    out["import_dlls"] = sorted(set(dlls))
    out["import_count"] = len(all_imports)
    out["import_categories"] = {k: sorted(set(v)) for k, v in buckets.items() if v}

    # Exports
    exports = []
    if hasattr(pe, "DIRECTORY_ENTRY_EXPORT"):
        for exp in pe.DIRECTORY_ENTRY_EXPORT.symbols:
            if exp.name:
                exports.append(exp.name.decode("latin1", "replace"))
    out["exports"] = exports[:50]
    out["has_tls"] = hasattr(pe, "DIRECTORY_ENTRY_TLS")

    # Authenticode signature presence (directory entry 4)
    try:
        sec = oh.DATA_DIRECTORY[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_SECURITY"]]
        out["signed"] = sec.VirtualAddress != 0 and sec.Size != 0
    except Exception:
        out["signed"] = None

    # Packer hint: few imports + high-entropy exec section
    highent = [s for s in out["sections"] if s["high_entropy"]]
    out["packer_hint"] = bool(highent) and out["import_count"] < 20

    # Interesting strings (sampled from whole file)
    with open(args.path, "rb") as f:
        blob = f.read()
    found = []
    seen = set()
    for s in ascii_strings(blob, args.strings_min) + [
        u.replace(b"\x00", b"") for u in unicode_strings(blob, args.strings_min)
    ]:
        if INTERESTING.search(s):
            txt = s.decode("latin1", "replace")
            if txt not in seen:
                seen.add(txt)
                found.append(txt)
        if len(found) >= args.max_strings:
            break
    out["interesting_strings"] = found

    print(json.dumps(out, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    sys.exit(main())