bikini-patchwork/patchwork/loader.py

from __future__ import annotations
import base64
import marshal
import random
import zlib
from . import crypto
from .packer import pack
from .util import gen_bytes, gen_name
_MOD_KEYS = ('sys', 'os', 'marshal', 'zlib', 'base64', 'hashlib', 'rand_mod')
_FUNC_KEYS = ('kstream', 'xor_dec', 'perm_inv', 'perm_dec', 'rot_dec')
_VAR_KEYS = ('data', 'code_obj', 'ns', 'blobs', 'tmp', 'i', 'k', 'p', 'g')

def _bytes_repr(b: bytes) -> str:
    return f'bytes.fromhex({b.hex()!r})'

def _chunked_bytes(b: bytes, rng: random.Random) -> str:
    if len(b) < 64:
        return _bytes_repr(b)
    nchunks = rng.randint(4, 9)
    chunk_size = (len(b) + nchunks - 1) // nchunks
    chunks: list[bytes] = []
    pos = 0
    while pos < len(b):
        chunks.append(b[pos:pos + chunk_size])
        pos += chunk_size
    order = list(range(len(chunks)))
    rng.shuffle(order)
    shuffled = [chunks[i] for i in order]
    inv = [0] * len(order)
    for i, o in enumerate(order):
        inv[o] = i
    parts_src = ', '.join((_bytes_repr(c) for c in shuffled))
    perm_src = ', '.join((str(x) for x in inv))
    return f"b''.join((({parts_src},)[__pwi] for __pwi in ({perm_src},)))"

def _junk(rng: random.Random) -> str:
    target = gen_name(rng)
    n = rng.randint(0, 1 << 16)
    forms = [f'{target} = (lambda _x: _x ^ _x)({n})', f'{target} = sum(({n},))', f'{target} = ({n} & {n}) | 0', f'{target} = pow({n} & 0xFF, 1, 1) + 0', f'{target} = len(()) + {n} - {n}', f'{target} = ({n} * 0) ^ 0', f'{target} = (~{n} ^ ~{n}) & 0xFF']
    return rng.choice(forms)

def _name_pool(rng: random.Random, keys: tuple[str, ...]) -> dict[str, str]:
    used: set[str] = set()
    out: dict[str, str] = {}
    for k in keys:
        n = gen_name(rng, used=used)
        used.add(n)
        out[k] = n
    return out

def _decrypt_call(N: dict[str, str], op: str, var: str, key: bytes) -> str:
    kbytes = _bytes_repr(key)
    if op == 'xor':
        return f"{var} = {N['xor_dec']}({var}, {kbytes})"
    if op == 'perm':
        return f"{var} = {N['perm_dec']}({var}, {kbytes})"
    if op == 'rot':
        return f"{var} = {N['rot_dec']}({var}, {kbytes})"
    raise ValueError(op)

def _emit_helpers(N: dict[str, str]) -> list[str]:
    return [f"def {N['kstream']}(k, n):\n    r = bytearray()\n    c = 0\n    while len(r) < n:\n        r += {N['hashlib']}.sha256(k + c.to_bytes(8, 'big')).digest()\n        c += 1\n    return bytes(r[:n])", f"def {N['xor_dec']}(d, k):\n    s = {N['kstream']}(k, len(d))\n    return bytes(a ^ b for a, b in zip(d, s))", f"def {N['perm_inv']}(k):\n    g = {N['rand_mod']}.Random(int.from_bytes({N['hashlib']}.sha256(k).digest(), 'big'))\n    p = list(range(256))\n    g.shuffle(p)\n    inv = bytearray(256)\n    for i, v in enumerate(p):\n        inv[v] = i\n    return bytes(inv)", f"def {N['perm_dec']}(d, k):\n    inv = {N['perm_inv']}(k)\n    return bytes(inv[b] for b in d)", f"def {N['rot_dec']}(d, k):\n    out = bytearray()\n    L = len(k)\n    for i, b in enumerate(d):\n        n = (-k[i % L]) & 7\n        if n:\n            out.append(((b << n) | (b >> (8 - n))) & 0xFF)\n        else:\n            out.append(b)\n    return bytes(out)"]

def _emit_imports(N: dict[str, str], rng: random.Random) -> list[str]:
    lines = [f"{N['sys']} = __import__('sys')", f"{N['os']} = __import__('os')", f"{N['marshal']} = __import__('marshal')", f"{N['zlib']} = __import__('zlib')", f"{N['base64']} = __import__('base64')", f"{N['hashlib']} = __import__('hashlib')", f"{N['rand_mod']} = __import__('random')"]
    rng.shuffle(lines)
    return lines

def _emit_anti_debug(N: dict[str, str], rng: random.Random) -> list[str]:
    checks = [f"({N['sys']}.gettrace() is None) or {N['os']}._exit(1)", f"({N['sys']}.getprofile() is None) or {N['os']}._exit(1)", f"all(_m not in {N['sys']}.modules for _m in ('pdb','bdb','pydevd','debugpy','pydevd_bundle','pydevd_bundle.pydevd_resolver','ptvsd','rpdb','wdb')) or {N['os']}._exit(1)", f"({N['os']}.environ.get('PYTHONBREAKPOINT', '') in ('', '0', 'pdb.set_trace')) or {N['os']}._exit(1)"]
    rng.shuffle(checks)
    return checks

def _emit_audit_hook(N: dict[str, str]) -> list[str]:
    return [f"def __pw_audit(event, args):\n    if event == 'import':\n        _n = args[0] if args else ''\n        if _n in ('pdb','bdb','pydevd','debugpy','pydevd_bundle','ptvsd','rpdb','wdb'):\n            {N['os']}._exit(1)\n    elif event in ('sys.settrace','sys.setprofile'):\n        {N['os']}._exit(1)", f"{N['sys']}.addaudithook(__pw_audit)"]

def _emit_frame_check(N: dict[str, str]) -> list[str]:
    return [f"_f = {N['sys']}._getframe(0)\nwhile _f is not None:\n    _gn = _f.f_globals.get('__name__', '')\n    if _gn.startswith(('pdb', 'bdb', 'pydevd', 'debugpy', '_pydev')):\n        {N['os']}._exit(1)\n    _f = _f.f_back"]

def _build_stage2_source(user_payload: bytes, user_layer_keys: list[tuple[str, bytes]], lazy_blobs: list[tuple[bytes, bytes]], resolver_name: str, mod_names: dict[str, str], rng: random.Random, *, anti_debug: bool=True) -> str:
    N = mod_names
    F = _name_pool(rng, _FUNC_KEYS)
    V = _name_pool(rng, _VAR_KEYS)
    parts: list[str] = []
    parts.append("_pw_types = __import__('types')")
    parts.extend(_emit_helpers({**N, **F}))
    if anti_debug:
        parts.extend(_emit_anti_debug({**N, **F}, rng))
    parts.append(_junk(rng))
    if lazy_blobs:
        blob_lines = []
        for k, e in lazy_blobs:
            blob_lines.append(f'({_chunked_bytes(k, rng)}, {_chunked_bytes(e, rng)})')
        parts.append(f"__pw_blobs = [{', '.join(blob_lines)}]")
    else:
        parts.append('__pw_blobs = []')
    parts.append(f"_pw_code_cache = {{}}\n_pw_done = set()\ndef _pw_find_func(g, target_code):\n    seen = set()\n    stack = list(g.values())\n    while stack:\n        v = stack.pop()\n        oid = id(v)\n        if oid in seen:\n            continue\n        seen.add(oid)\n        if isinstance(v, _pw_types.FunctionType) and v.__code__ is target_code:\n            return v\n        if isinstance(v, (classmethod, staticmethod)):\n            inner = v.__func__\n            if isinstance(inner, _pw_types.FunctionType) and inner.__code__ is target_code:\n                return inner\n            stack.append(inner)\n            continue\n        if isinstance(v, property):\n            for fn in (v.fget, v.fset, v.fdel):\n                if fn is not None:\n                    stack.append(fn)\n            continue\n        if isinstance(v, type):\n            for vv in v.__dict__.values():\n                stack.append(vv)\n    return None\ndef {resolver_name}(idx, args, kw):\n    fr = {N['sys']}._getframe(1)\n    g = fr.f_globals\n    stub_code = fr.f_code\n    if idx not in _pw_code_cache:\n        key, enc = __pw_blobs[idx]\n        _pw_code_cache[idx] = {N['marshal']}.loads({F['xor_dec']}(enc, key))\n    real = _pw_code_cache[idx]\n    target = _pw_find_func(g, stub_code)\n    if target is not None and id(target) not in _pw_done:\n        target.__code__ = real\n        _pw_done.add(id(target))\n        return target(*args, **kw)\n    fn = _pw_types.FunctionType(real, g)\n    return fn(*args, **kw)")
    user_encoded = base64.b85encode(user_payload)
    parts.append(f"{V['data']} = {N['base64']}.b85decode({_chunked_bytes(user_encoded, rng)})")
    for op, key in reversed(user_layer_keys):
        parts.append(_decrypt_call({**N, **F}, op, V['data'], key))
    parts.append(f"{V['data']} = {N['zlib']}.decompress({V['data']})")
    parts.append(f"{V['code_obj']} = {N['marshal']}.loads({V['data']})")
    parts.append(_junk(rng))
    if anti_debug:
        parts.append(f"({N['sys']}.gettrace() is None) or {N['os']}._exit(1)")
    parts.append(f"{V['ns']} = {{'__name__': '__main__', '__builtins__': __builtins__, '__file__': __file__, '__package__': None, '__doc__': None, '{resolver_name}': {resolver_name}}}")
    parts.append(f"exec({V['code_obj']}, {V['ns']})")
    return '\n'.join(parts) + '\n'

def build_loader(user_payload: bytes, user_layer_keys: list[tuple[str, bytes]], lazy_blobs: list[tuple[bytes, bytes]], resolver_name: str, rng: random.Random, *, anti_debug: bool=True, stage2_layers: int=3, junk_count: int=8) -> str:
    N = _name_pool(rng, _MOD_KEYS)
    F = _name_pool(rng, _FUNC_KEYS)
    V = _name_pool(rng, _VAR_KEYS)
    stage2_src = _build_stage2_source(user_payload, user_layer_keys, lazy_blobs, resolver_name=resolver_name, mod_names=N, rng=rng, anti_debug=anti_debug)
    stage2_code = compile(stage2_src, '<patchwork:stage2>', 'exec')
    stage2_payload, stage2_keys = pack(stage2_code, rng, layers=stage2_layers)
    parts: list[str] = []
    parts.extend(_emit_imports(N, rng))
    parts.append(_junk(rng))
    if anti_debug:
        parts.extend(_emit_anti_debug({**N, **F}, rng))
        parts.extend(_emit_audit_hook(N))
        parts.extend(_emit_frame_check(N))
    parts.append(_junk(rng))
    parts.extend(_emit_helpers({**N, **F}))
    parts.append(_junk(rng))
    stage2_encoded = base64.b85encode(stage2_payload)
    parts.append(f"{V['data']} = {N['base64']}.b85decode({_chunked_bytes(stage2_encoded, rng)})")
    for op, key in reversed(stage2_keys):
        parts.append(_decrypt_call({**N, **F}, op, V['data'], key))
    parts.append(f"{V['data']} = {N['zlib']}.decompress({V['data']})")
    if anti_debug:
        parts.append(f"({N['sys']}.gettrace() is None) or {N['os']}._exit(1)")
    parts.append(f"{V['code_obj']} = {N['marshal']}.loads({V['data']})")
    parts.append(_junk(rng))
    parts.append(f"{V['ns']} = {{'__builtins__': __builtins__, '__file__': __file__, '__name__': __name__, '{N['sys']}': {N['sys']}, '{N['os']}': {N['os']}, '{N['marshal']}': {N['marshal']}, '{N['zlib']}': {N['zlib']}, '{N['base64']}': {N['base64']}, '{N['hashlib']}': {N['hashlib']}, '{N['rand_mod']}': {N['rand_mod']}}}")
    parts.append(f"exec({V['code_obj']}, {V['ns']})")
    body = list(parts)
    for _ in range(junk_count):
        idx = rng.randint(0, len(body))
        body.insert(idx, _junk(rng))
    return '\n'.join(body) + '\n'