mirror of
https://github.com/bikini/patchwork
synced 2026-06-27 08:08:41 +00:00
bikini
89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
from __future__ import annotations
|
|
import marshal
|
|
import random
|
|
import types
|
|
from .crypto import stream_xor
|
|
from .util import gen_bytes
|
|
LAZY_PREFIX = '<pw'
|
|
LAZY_SUFFIX = '>'
|
|
_CO_GENERATOR = 0x20
|
|
_CO_COROUTINE = 0x100
|
|
_CO_ASYNC_GENERATOR = 0x200
|
|
|
|
|
|
def _classify_code(code: types.CodeType) -> str:
|
|
if not code.co_flags & 1:
|
|
return 'class'
|
|
if code.co_name.startswith('<') and code.co_name.endswith('>'):
|
|
return 'special'
|
|
return 'function'
|
|
|
|
|
|
def _stub_source(kind: str, idx: int, resolver_name: str) -> str:
|
|
if kind == 'async_gen':
|
|
return (
|
|
f'async def _stub(*_a, **_k):\n'
|
|
f' async for _v in {resolver_name}({idx}, _a, _k):\n'
|
|
f' yield _v\n'
|
|
)
|
|
if kind == 'coroutine':
|
|
return (
|
|
f'async def _stub(*_a, **_k):\n'
|
|
f' return await {resolver_name}({idx}, _a, _k)\n'
|
|
)
|
|
if kind == 'generator':
|
|
return (
|
|
f'def _stub(*_a, **_k):\n'
|
|
f' yield from {resolver_name}({idx}, _a, _k)\n'
|
|
)
|
|
return (
|
|
f'def _stub(*_a, **_k):\n'
|
|
f' return {resolver_name}({idx}, _a, _k)\n'
|
|
)
|
|
|
|
|
|
def _kind_for(code: types.CodeType) -> str:
|
|
f = code.co_flags
|
|
if f & _CO_ASYNC_GENERATOR:
|
|
return 'async_gen'
|
|
if f & _CO_COROUTINE:
|
|
return 'coroutine'
|
|
if f & _CO_GENERATOR:
|
|
return 'generator'
|
|
return 'function'
|
|
|
|
|
|
def _make_stub_code(idx: int, original: types.CodeType, resolver_name: str) -> types.CodeType:
|
|
kind = _kind_for(original)
|
|
src = _stub_source(kind, idx, resolver_name)
|
|
mod = compile(src, '<pwstub>', 'exec')
|
|
inner = next((c for c in mod.co_consts if isinstance(c, types.CodeType)))
|
|
return inner.replace(
|
|
co_name=f'{LAZY_PREFIX}{idx:x}{LAZY_SUFFIX}',
|
|
co_freevars=original.co_freevars,
|
|
co_cellvars=original.co_cellvars,
|
|
)
|
|
|
|
|
|
def encrypt_user_functions(top: types.CodeType, rng: random.Random, resolver_name: str, key_size: int=24) -> tuple[types.CodeType, list[tuple[bytes, bytes]]]:
|
|
blobs: list[tuple[bytes, bytes]] = []
|
|
|
|
def transform(code: types.CodeType, parent_kind: str) -> types.CodeType:
|
|
new_consts: list = []
|
|
for c in code.co_consts:
|
|
if isinstance(c, types.CodeType):
|
|
ck = _classify_code(c)
|
|
if parent_kind in ('module', 'class') and ck == 'function':
|
|
idx = len(blobs)
|
|
key = gen_bytes(rng, key_size)
|
|
enc = stream_xor(marshal.dumps(c), key)
|
|
blobs.append((key, enc))
|
|
new_consts.append(_make_stub_code(idx, c, resolver_name))
|
|
else:
|
|
new_consts.append(transform(c, ck))
|
|
else:
|
|
new_consts.append(c)
|
|
return code.replace(co_consts=tuple(new_consts))
|
|
new_top = transform(top, 'module')
|
|
return (new_top, blobs)
|