diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c1ce52f..0dfb196 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,3 +21,4 @@ jobs: run: | python tests/test_audit.py python tests/test_obfuscator.py + python tests/test_stress_matrix.py diff --git a/README.md b/README.md index 0cdcd88..b8d4e49 100644 --- a/README.md +++ b/README.md @@ -222,14 +222,32 @@ python examples/hello_obf.py `examples/stress.py` is the torture test - decorators, generators, classes with `super()`, `match` statements, walrus, exceptions, `*args`/`**kwargs`, globals. If something doesn't work, it'll usually break here first. +`examples/modern.py` covers newer Python constructs: `dataclass`, structural +pattern matching with capture binders, mapping/list/star patterns, async +functions, context managers, properties, f-strings with format specs, bytes, +closures, and `nonlocal`. + +`examples/future_annotations.py` covers `from __future__ import annotations` +and verifies that observable annotation strings are preserved. + ## Tests ```sh python tests/test_obfuscator.py python tests/test_audit.py +python tests/test_stress_matrix.py ``` -Runs every example through the obfuscator at multiple seeds, executes original and obfuscated versions, and checks stdout matches byte-for-byte. Also confirms different seeds produce different output and the same seed produces stable output. +Runs every example through the obfuscator at multiple seeds, executes original +and obfuscated versions, and checks stdout matches byte-for-byte. Also confirms +different seeds produce different output and the same seed produces stable +output. + +The stress matrix generates deterministic Python programs with arithmetic, +bitwise operations, pattern matching, closures, comprehensions, strings, bytes, +and future annotations. It then runs them through multiple obfuscation seeds and +option combinations, including disabled lazy loading, disabled renaming, +disabled string encryption, and disabled opaque/junk transforms. ## Layout diff --git a/examples/future_annotations.py b/examples/future_annotations.py new file mode 100644 index 0000000..c1e82f4 --- /dev/null +++ b/examples/future_annotations.py @@ -0,0 +1,25 @@ +from __future__ import annotations + + +class Node: + def __init__(self, value: int, child: Node | None = None): + self.value = value + self.child = child + + +def flatten(node: Node | None) -> list[int]: + values: list[int] = [] + while node is not None: + values.append(node.value) + node = node.child + return values + + +def run() -> None: + chain = Node(1, Node(2, Node(3))) + print('annotations:', flatten(chain)) + print('future:', flatten.__annotations__) + + +if __name__ == '__main__': + run() diff --git a/examples/modern.py b/examples/modern.py new file mode 100644 index 0000000..a5028d6 --- /dev/null +++ b/examples/modern.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import asyncio +from contextlib import contextmanager +from dataclasses import dataclass + + +@dataclass +class Point: + x: int + y: int + + +@contextmanager +def tagged(name: str): + yield f'<{name}>' + + +def describe(value): + match value: + case Point(0, y): + return f'y-axis:{y}' + case Point(x=x, y=y) if x == y: + return f'diagonal:{x}' + case {'kind': 'pair', 'items': [first, *rest], **extra}: + return f'pair:{first}:{len(rest)}:{sorted(extra)}' + case ('wrap', inner as captured): + return f'wrapped:{inner}:{captured}' + case [head, *tail]: + return f'list:{head}:{sum(tail)}' + case _: + return 'unknown' + + +async def async_total(values): + total = 0 + for value in values: + await asyncio.sleep(0) + total += value + return total + + +def make_counter(start=0): + count = start + + def next_value(step=1): + nonlocal count + count += step + return count + + return next_value + + +class Box: + def __init__(self, value): + self.value = value + + @property + def doubled(self): + return self.value * 2 + + +def run() -> None: + items = [ + Point(0, 7), + Point(4, 4), + {'kind': 'pair', 'items': [3, 5, 8], 'tag': 'fib'}, + ('wrap', 'token'), + [1, 2, 3, 4], + object(), + ] + print('descriptions:', [describe(item) for item in items]) + print('async:', asyncio.run(async_total([1, 2, 3, 4]))) + counter = make_counter(10) + print('counter:', [counter(), counter(5), counter()]) + with tagged('ctx') as label: + print('context:', label) + print('box:', Box(9).doubled) + print('bytes:', bytes([65, 66, 67]).decode()) + print('format:', f'{255:#06x}:{3.14159:.2f}') + + +if __name__ == '__main__': + run() diff --git a/patchwork/core.py b/patchwork/core.py index 075f6b6..6daaa8c 100644 --- a/patchwork/core.py +++ b/patchwork/core.py @@ -1,79 +1,98 @@ -from __future__ import annotations -import ast -from pathlib import Path -from .lazy import encrypt_user_functions -from .loader import build_loader -from .packer import pack -from .transforms import IdentifierRenamer, NumberObfuscator, OpaquePredicateInjector, StringEncryptor, build_decrypt_helper, collect_skip_names -from .transforms.junk import JunkBranchInjector -from .transforms.mba import MBATransformer -from .transforms.opaque import build_seed_stmt -from .util import make_rng -_SEED_NAME = '_pw_seed_value' -_DEC_S_NAME = '_pw_dec_str' -_DEC_B_NAME = '_pw_dec_bytes' -_RESOLVE_NAME = '__pw_resolve_lazy__' - -class Obfuscator: - - def __init__(self, *, seed: int | None=None, rename: bool=True, encrypt_strings: bool=True, obfuscate_numbers: bool=True, opaque_predicates: bool=True, mba: bool=True, junk_branches: bool=True, lazy_funcs: bool=True, anti_debug: bool=True, layers: int=3, stage2_layers: int=3, keep: set[str] | None=None): - self.rng, self.seed = make_rng(seed) - self.rename = rename - self.encrypt_strings = encrypt_strings - self.obfuscate_numbers = obfuscate_numbers - self.opaque_predicates = opaque_predicates - self.mba = mba - self.junk_branches = junk_branches - self.lazy_funcs = lazy_funcs - self.anti_debug = anti_debug - self.layers = max(1, int(layers)) - self.stage2_layers = max(1, int(stage2_layers)) - self.keep_extra = set(keep or ()) - - def obfuscate(self, source: str) -> str: - tree = ast.parse(source) - skip = collect_skip_names(tree) - seed_used = False - if self.opaque_predicates: - opi = OpaquePredicateInjector(self.rng, _SEED_NAME) - tree = opi.visit(tree) - seed_used = seed_used or opi.injected - if self.junk_branches: - jbi = JunkBranchInjector(self.rng, _SEED_NAME) - tree = jbi.visit(tree) - seed_used = seed_used or jbi.injected - if seed_used: - tree.body.insert(0, build_seed_stmt(_SEED_NAME)) - if self.mba: - tree = MBATransformer(self.rng).visit(tree) - if self.obfuscate_numbers: - tree = NumberObfuscator(self.rng).visit(tree) - if self.encrypt_strings: - enc = StringEncryptor(self.rng, _DEC_S_NAME, _DEC_B_NAME) - tree = enc.visit(tree) - tree.body[:0] = build_decrypt_helper(_DEC_S_NAME, _DEC_B_NAME) - if self.rename: - renamer = IdentifierRenamer(self.rng, keep=skip | self.keep_extra | {_RESOLVE_NAME}) - tree = renamer.visit(tree) - ast.fix_missing_locations(tree) - user_code = compile(tree, '', 'exec') - if self.lazy_funcs: - user_code, lazy_blobs = encrypt_user_functions(user_code, self.rng, _RESOLVE_NAME) - else: - lazy_blobs = [] - user_payload, user_keys = pack(user_code, self.rng, layers=self.layers) - return build_loader(user_payload, user_keys, lazy_blobs, resolver_name=_RESOLVE_NAME, rng=self.rng, anti_debug=self.anti_debug, stage2_layers=self.stage2_layers) - -def obfuscate(source: str, **kwargs) -> str: - return Obfuscator(**kwargs).obfuscate(source) - -def obfuscate_file(input_path: str | Path, output_path: str | Path | None=None, **kwargs) -> Path: - inp = Path(input_path) - src = inp.read_text(encoding='utf-8') - out = obfuscate(src, **kwargs) - if output_path is None: - outp = inp.with_name(inp.stem + '_obf.py') - else: - outp = Path(output_path) - outp.write_text(out, encoding='utf-8') - return outp +from __future__ import annotations +import ast +from pathlib import Path +from .lazy import encrypt_user_functions +from .loader import build_loader +from .packer import pack +from .transforms import IdentifierRenamer, NumberObfuscator, OpaquePredicateInjector, StringEncryptor, build_decrypt_helper, collect_skip_names +from .transforms.junk import JunkBranchInjector +from .transforms.mba import MBATransformer +from .transforms.opaque import build_seed_stmt +from .util import make_rng +_SEED_NAME = '_pw_seed_value' +_DEC_S_NAME = '_pw_dec_str' +_DEC_B_NAME = '_pw_dec_bytes' +_RESOLVE_NAME = '__pw_resolve_lazy__' + +def _insert_runtime_stmts(tree: ast.Module, stmts: list[ast.stmt]) -> None: + if not stmts: + return + insert_at = 0 + if ( + tree.body + and isinstance(tree.body[0], ast.Expr) + and isinstance(tree.body[0].value, ast.Constant) + and isinstance(tree.body[0].value.value, str) + ): + insert_at = 1 + while ( + insert_at < len(tree.body) + and isinstance(tree.body[insert_at], ast.ImportFrom) + and tree.body[insert_at].module == '__future__' + ): + insert_at += 1 + tree.body[insert_at:insert_at] = stmts + +class Obfuscator: + + def __init__(self, *, seed: int | None=None, rename: bool=True, encrypt_strings: bool=True, obfuscate_numbers: bool=True, opaque_predicates: bool=True, mba: bool=True, junk_branches: bool=True, lazy_funcs: bool=True, anti_debug: bool=True, layers: int=3, stage2_layers: int=3, keep: set[str] | None=None): + self.rng, self.seed = make_rng(seed) + self.rename = rename + self.encrypt_strings = encrypt_strings + self.obfuscate_numbers = obfuscate_numbers + self.opaque_predicates = opaque_predicates + self.mba = mba + self.junk_branches = junk_branches + self.lazy_funcs = lazy_funcs + self.anti_debug = anti_debug + self.layers = max(1, int(layers)) + self.stage2_layers = max(1, int(stage2_layers)) + self.keep_extra = set(keep or ()) + + def obfuscate(self, source: str) -> str: + tree = ast.parse(source) + skip = collect_skip_names(tree) + seed_used = False + if self.opaque_predicates: + opi = OpaquePredicateInjector(self.rng, _SEED_NAME) + tree = opi.visit(tree) + seed_used = seed_used or opi.injected + if self.junk_branches: + jbi = JunkBranchInjector(self.rng, _SEED_NAME) + tree = jbi.visit(tree) + seed_used = seed_used or jbi.injected + if seed_used: + _insert_runtime_stmts(tree, [build_seed_stmt(_SEED_NAME)]) + if self.mba: + tree = MBATransformer(self.rng).visit(tree) + if self.obfuscate_numbers: + tree = NumberObfuscator(self.rng).visit(tree) + if self.encrypt_strings: + enc = StringEncryptor(self.rng, _DEC_S_NAME, _DEC_B_NAME) + tree = enc.visit(tree) + _insert_runtime_stmts(tree, build_decrypt_helper(_DEC_S_NAME, _DEC_B_NAME)) + if self.rename: + renamer = IdentifierRenamer(self.rng, keep=skip | self.keep_extra | {_RESOLVE_NAME}) + tree = renamer.visit(tree) + ast.fix_missing_locations(tree) + user_code = compile(tree, '', 'exec') + if self.lazy_funcs: + user_code, lazy_blobs = encrypt_user_functions(user_code, self.rng, _RESOLVE_NAME) + else: + lazy_blobs = [] + user_payload, user_keys = pack(user_code, self.rng, layers=self.layers) + return build_loader(user_payload, user_keys, lazy_blobs, resolver_name=_RESOLVE_NAME, rng=self.rng, anti_debug=self.anti_debug, stage2_layers=self.stage2_layers) + +def obfuscate(source: str, **kwargs) -> str: + return Obfuscator(**kwargs).obfuscate(source) + +def obfuscate_file(input_path: str | Path, output_path: str | Path | None=None, **kwargs) -> Path: + inp = Path(input_path) + src = inp.read_text(encoding='utf-8') + out = obfuscate(src, **kwargs) + if output_path is None: + outp = inp.with_name(inp.stem + '_obf.py') + else: + outp = Path(output_path) + outp.write_text(out, encoding='utf-8') + return outp diff --git a/patchwork/transforms/identifiers.py b/patchwork/transforms/identifiers.py index f0b932e..1e6b5d6 100644 --- a/patchwork/transforms/identifiers.py +++ b/patchwork/transforms/identifiers.py @@ -1,167 +1,218 @@ -from __future__ import annotations -import ast -import builtins -import keyword -import random -from ..util import gen_name -_DUNDERS = {'__name__', '__main__', '__file__', '__doc__', '__builtins__', '__package__', '__path__', '__loader__', '__spec__', '__cached__', '__all__', '__init__', '__new__', '__del__', '__init_subclass__', '__subclasshook__', '__set_name__', '__class_getitem__', '__enter__', '__exit__', '__aenter__', '__aexit__', '__await__', '__aiter__', '__anext__', '__call__', '__repr__', '__str__', '__bytes__', '__format__', '__hash__', '__bool__', '__len__', '__length_hint__', '__sizeof__', '__dir__', '__eq__', '__ne__', '__lt__', '__le__', '__gt__', '__ge__', '__add__', '__sub__', '__mul__', '__truediv__', '__floordiv__', '__mod__', '__divmod__', '__pow__', '__lshift__', '__rshift__', '__and__', '__or__', '__xor__', '__radd__', '__rsub__', '__rmul__', '__rtruediv__', '__rfloordiv__', '__rmod__', '__rdivmod__', '__rpow__', '__rlshift__', '__rrshift__', '__rand__', '__ror__', '__rxor__', '__iadd__', '__isub__', '__imul__', '__itruediv__', '__ifloordiv__', '__imod__', '__ipow__', '__ilshift__', '__irshift__', '__iand__', '__ior__', '__ixor__', '__matmul__', '__rmatmul__', '__imatmul__', '__neg__', '__pos__', '__abs__', '__invert__', '__int__', '__float__', '__complex__', '__round__', '__index__', '__trunc__', '__floor__', '__ceil__', '__getattr__', '__setattr__', '__delattr__', '__getattribute__', '__getitem__', '__setitem__', '__delitem__', '__missing__', '__iter__', '__next__', '__reversed__', '__contains__', '__copy__', '__deepcopy__', '__reduce__', '__reduce_ex__', '__getstate__', '__setstate__', '__getnewargs__', '__getnewargs_ex__', '__class__', '__instancecheck__', '__subclasscheck__', '__slots__', '__dict__', '__weakref__', '__module__', '__qualname__', '__defaults__', '__kwdefaults__', '__annotations__', '__wrapped__', '__signature__', '__match_args__', '__post_init__', '__fspath__', '__index__', '__buffer__', '__release_buffer__'} -RESERVED: set[str] = set(dir(builtins)) | set(keyword.kwlist) | _DUNDERS | {'self', 'cls', 'metaclass'} - -def collect_skip_names(tree: ast.AST) -> set[str]: - skip: set[str] = set() - for node in ast.walk(tree): - if isinstance(node, ast.Assign): - for target in node.targets: - if isinstance(target, ast.Name) and target.id == '__all__': - if isinstance(node.value, (ast.List, ast.Tuple, ast.Set)): - for elt in node.value.elts: - if isinstance(elt, ast.Constant) and isinstance(elt.value, str): - skip.add(elt.value) - elif isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and (node.func.id in ('getattr', 'setattr', 'hasattr', 'delattr')): - if len(node.args) >= 2: - arg = node.args[1] - if isinstance(arg, ast.Constant) and isinstance(arg.value, str): - skip.add(arg.value) - elif isinstance(node, ast.ClassDef): - for stmt in node.body: - if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)): - skip.add(stmt.name) - elif isinstance(stmt, ast.Assign): - for tgt in stmt.targets: - if isinstance(tgt, ast.Name): - skip.add(tgt.id) - elif isinstance(tgt, (ast.Tuple, ast.List)): - for elt in tgt.elts: - if isinstance(elt, ast.Name): - skip.add(elt.id) - elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name): - skip.add(stmt.target.id) - elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): - args = node.args - for a in args.args + args.posonlyargs + args.kwonlyargs: - skip.add(a.arg) - if args.vararg: - skip.add(args.vararg.arg) - if args.kwarg: - skip.add(args.kwarg.arg) - return skip - -class IdentifierRenamer(ast.NodeTransformer): - - def __init__(self, rng: random.Random, keep: set[str] | None=None): - self.rng = rng - self.mapping: dict[str, str] = {} - self.reserved = RESERVED | (keep or set()) - self._used_outputs: set[str] = set() - - def _new(self) -> str: - n = gen_name(self.rng, used=self._used_outputs) - self._used_outputs.add(n) - return n - - def _rename(self, name: str) -> str: - if not name or name in self.reserved: - return name - if name.startswith('__') and name.endswith('__') and (len(name) >= 4): - return name - if name not in self.mapping: - self.mapping[name] = self._new() - return self.mapping[name] - - def visit_Name(self, node: ast.Name) -> ast.AST: - node.id = self._rename(node.id) - return node - - def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST: - node.name = self._rename(node.name) - node.decorator_list = [self.visit(d) for d in node.decorator_list] - if node.returns is not None: - node.returns = self.visit(node.returns) - node.args = self.visit(node.args) - node.body = [self.visit(stmt) for stmt in node.body] - return node - visit_AsyncFunctionDef = visit_FunctionDef - - def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST: - node.name = self._rename(node.name) - node.bases = [self.visit(b) for b in node.bases] - for kw in node.keywords: - kw.value = self.visit(kw.value) - node.decorator_list = [self.visit(d) for d in node.decorator_list] - node.body = [self.visit(stmt) for stmt in node.body] - return node - - def visit_Lambda(self, node: ast.Lambda) -> ast.AST: - node.args = self.visit(node.args) - node.body = self.visit(node.body) - return node - - def visit_arg(self, node: ast.arg) -> ast.AST: - node.arg = self._rename(node.arg) - if node.annotation is not None: - node.annotation = self.visit(node.annotation) - return node - - def visit_arguments(self, node: ast.arguments) -> ast.AST: - node.args = [self.visit(a) for a in node.args] - node.posonlyargs = [self.visit(a) for a in node.posonlyargs] - node.kwonlyargs = [self.visit(a) for a in node.kwonlyargs] - if node.vararg is not None: - node.vararg = self.visit(node.vararg) - if node.kwarg is not None: - node.kwarg = self.visit(node.kwarg) - node.defaults = [self.visit(d) for d in node.defaults] - node.kw_defaults = [self.visit(d) if d is not None else None for d in node.kw_defaults] - return node - - def visit_Attribute(self, node: ast.Attribute) -> ast.AST: - node.value = self.visit(node.value) - return node - - def visit_keyword(self, node: ast.keyword) -> ast.AST: - if node.value is not None: - node.value = self.visit(node.value) - return node - - def visit_Global(self, node: ast.Global) -> ast.AST: - node.names = [self._rename(n) for n in node.names] - return node - - def visit_Nonlocal(self, node: ast.Nonlocal) -> ast.AST: - node.names = [self._rename(n) for n in node.names] - return node - - def visit_ExceptHandler(self, node: ast.ExceptHandler) -> ast.AST: - if node.type is not None: - node.type = self.visit(node.type) - if node.name: - node.name = self._rename(node.name) - node.body = [self.visit(s) for s in node.body] - return node - - def visit_Import(self, node: ast.Import) -> ast.AST: - for alias in node.names: - if '.' in alias.name: - if alias.asname: - new = self._rename(alias.asname) - alias.asname = new - continue - if alias.asname: - alias.asname = self._rename(alias.asname) - else: - if alias.name not in self.mapping: - self.mapping[alias.name] = self._new() - alias.asname = self.mapping[alias.name] - return node - - def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST: - for alias in node.names: - if alias.name == '*': - continue - if alias.asname: - alias.asname = self._rename(alias.asname) - else: - if alias.name not in self.mapping: - self.mapping[alias.name] = self._new() - alias.asname = self.mapping[alias.name] - return node +from __future__ import annotations +import ast +import builtins +import keyword +import random +from ..util import gen_name +_DUNDERS = {'__name__', '__main__', '__file__', '__doc__', '__builtins__', '__package__', '__path__', '__loader__', '__spec__', '__cached__', '__all__', '__init__', '__new__', '__del__', '__init_subclass__', '__subclasshook__', '__set_name__', '__class_getitem__', '__enter__', '__exit__', '__aenter__', '__aexit__', '__await__', '__aiter__', '__anext__', '__call__', '__repr__', '__str__', '__bytes__', '__format__', '__hash__', '__bool__', '__len__', '__length_hint__', '__sizeof__', '__dir__', '__eq__', '__ne__', '__lt__', '__le__', '__gt__', '__ge__', '__add__', '__sub__', '__mul__', '__truediv__', '__floordiv__', '__mod__', '__divmod__', '__pow__', '__lshift__', '__rshift__', '__and__', '__or__', '__xor__', '__radd__', '__rsub__', '__rmul__', '__rtruediv__', '__rfloordiv__', '__rmod__', '__rdivmod__', '__rpow__', '__rlshift__', '__rrshift__', '__rand__', '__ror__', '__rxor__', '__iadd__', '__isub__', '__imul__', '__itruediv__', '__ifloordiv__', '__imod__', '__ipow__', '__ilshift__', '__irshift__', '__iand__', '__ior__', '__ixor__', '__matmul__', '__rmatmul__', '__imatmul__', '__neg__', '__pos__', '__abs__', '__invert__', '__int__', '__float__', '__complex__', '__round__', '__index__', '__trunc__', '__floor__', '__ceil__', '__getattr__', '__setattr__', '__delattr__', '__getattribute__', '__getitem__', '__setitem__', '__delitem__', '__missing__', '__iter__', '__next__', '__reversed__', '__contains__', '__copy__', '__deepcopy__', '__reduce__', '__reduce_ex__', '__getstate__', '__setstate__', '__getnewargs__', '__getnewargs_ex__', '__class__', '__instancecheck__', '__subclasscheck__', '__slots__', '__dict__', '__weakref__', '__module__', '__qualname__', '__defaults__', '__kwdefaults__', '__annotations__', '__wrapped__', '__signature__', '__match_args__', '__post_init__', '__fspath__', '__index__', '__buffer__', '__release_buffer__'} +RESERVED: set[str] = set(dir(builtins)) | set(keyword.kwlist) | _DUNDERS | {'self', 'cls', 'metaclass'} + +def _has_future_annotations(tree: ast.AST) -> bool: + if not isinstance(tree, ast.Module): + return False + for stmt in tree.body: + if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant) and isinstance(stmt.value.value, str): + continue + if isinstance(stmt, ast.ImportFrom) and stmt.module == '__future__': + if any(alias.name == 'annotations' for alias in stmt.names): + return True + continue + break + return False + +def _names_in_annotation(node: ast.AST | None) -> set[str]: + if node is None: + return set() + return {child.id for child in ast.walk(node) if isinstance(child, ast.Name)} + +def collect_skip_names(tree: ast.AST) -> set[str]: + skip: set[str] = set() + preserve_annotation_names = _has_future_annotations(tree) + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == '__all__': + if isinstance(node.value, (ast.List, ast.Tuple, ast.Set)): + for elt in node.value.elts: + if isinstance(elt, ast.Constant) and isinstance(elt.value, str): + skip.add(elt.value) + elif isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and (node.func.id in ('getattr', 'setattr', 'hasattr', 'delattr')): + if len(node.args) >= 2: + arg = node.args[1] + if isinstance(arg, ast.Constant) and isinstance(arg.value, str): + skip.add(arg.value) + elif isinstance(node, ast.ClassDef): + for stmt in node.body: + if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)): + skip.add(stmt.name) + elif isinstance(stmt, ast.Assign): + for tgt in stmt.targets: + if isinstance(tgt, ast.Name): + skip.add(tgt.id) + elif isinstance(tgt, (ast.Tuple, ast.List)): + for elt in tgt.elts: + if isinstance(elt, ast.Name): + skip.add(elt.id) + elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name): + skip.add(stmt.target.id) + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + args = node.args + if preserve_annotation_names: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + skip.update(_names_in_annotation(node.returns)) + for a in args.args + args.posonlyargs + args.kwonlyargs: + skip.update(_names_in_annotation(a.annotation)) + if args.vararg: + skip.update(_names_in_annotation(args.vararg.annotation)) + if args.kwarg: + skip.update(_names_in_annotation(args.kwarg.annotation)) + for a in args.args + args.posonlyargs + args.kwonlyargs: + skip.add(a.arg) + if args.vararg: + skip.add(args.vararg.arg) + if args.kwarg: + skip.add(args.kwarg.arg) + elif preserve_annotation_names and isinstance(node, ast.AnnAssign): + skip.update(_names_in_annotation(node.annotation)) + return skip + +class IdentifierRenamer(ast.NodeTransformer): + + def __init__(self, rng: random.Random, keep: set[str] | None=None): + self.rng = rng + self.mapping: dict[str, str] = {} + self.reserved = RESERVED | (keep or set()) + self._used_outputs: set[str] = set() + + def _new(self) -> str: + n = gen_name(self.rng, used=self._used_outputs) + self._used_outputs.add(n) + return n + + def _rename(self, name: str) -> str: + if not name or name in self.reserved: + return name + if name.startswith('__') and name.endswith('__') and (len(name) >= 4): + return name + if name not in self.mapping: + self.mapping[name] = self._new() + return self.mapping[name] + + def visit_Name(self, node: ast.Name) -> ast.AST: + node.id = self._rename(node.id) + return node + + def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST: + node.name = self._rename(node.name) + node.decorator_list = [self.visit(d) for d in node.decorator_list] + if node.returns is not None: + node.returns = self.visit(node.returns) + node.args = self.visit(node.args) + node.body = [self.visit(stmt) for stmt in node.body] + return node + visit_AsyncFunctionDef = visit_FunctionDef + + def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST: + node.name = self._rename(node.name) + node.bases = [self.visit(b) for b in node.bases] + for kw in node.keywords: + kw.value = self.visit(kw.value) + node.decorator_list = [self.visit(d) for d in node.decorator_list] + node.body = [self.visit(stmt) for stmt in node.body] + return node + + def visit_Lambda(self, node: ast.Lambda) -> ast.AST: + node.args = self.visit(node.args) + node.body = self.visit(node.body) + return node + + def visit_arg(self, node: ast.arg) -> ast.AST: + node.arg = self._rename(node.arg) + if node.annotation is not None: + node.annotation = self.visit(node.annotation) + return node + + def visit_arguments(self, node: ast.arguments) -> ast.AST: + node.args = [self.visit(a) for a in node.args] + node.posonlyargs = [self.visit(a) for a in node.posonlyargs] + node.kwonlyargs = [self.visit(a) for a in node.kwonlyargs] + if node.vararg is not None: + node.vararg = self.visit(node.vararg) + if node.kwarg is not None: + node.kwarg = self.visit(node.kwarg) + node.defaults = [self.visit(d) for d in node.defaults] + node.kw_defaults = [self.visit(d) if d is not None else None for d in node.kw_defaults] + return node + + def visit_Attribute(self, node: ast.Attribute) -> ast.AST: + node.value = self.visit(node.value) + return node + + def visit_keyword(self, node: ast.keyword) -> ast.AST: + if node.value is not None: + node.value = self.visit(node.value) + return node + + def visit_Global(self, node: ast.Global) -> ast.AST: + node.names = [self._rename(n) for n in node.names] + return node + + def visit_Nonlocal(self, node: ast.Nonlocal) -> ast.AST: + node.names = [self._rename(n) for n in node.names] + return node + + def visit_ExceptHandler(self, node: ast.ExceptHandler) -> ast.AST: + if node.type is not None: + node.type = self.visit(node.type) + if node.name: + node.name = self._rename(node.name) + node.body = [self.visit(s) for s in node.body] + return node + + def visit_Import(self, node: ast.Import) -> ast.AST: + for alias in node.names: + if '.' in alias.name: + if alias.asname: + new = self._rename(alias.asname) + alias.asname = new + continue + if alias.asname: + alias.asname = self._rename(alias.asname) + else: + if alias.name not in self.mapping: + self.mapping[alias.name] = self._new() + alias.asname = self.mapping[alias.name] + return node + + def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST: + if node.module == '__future__': + return node + for alias in node.names: + if alias.name == '*': + continue + if alias.asname: + alias.asname = self._rename(alias.asname) + else: + if alias.name not in self.mapping: + self.mapping[alias.name] = self._new() + alias.asname = self.mapping[alias.name] + return node + + def visit_MatchAs(self, node: ast.MatchAs) -> ast.AST: + if node.pattern is not None: + node.pattern = self.visit(node.pattern) + if node.name is not None: + node.name = self._rename(node.name) + return node + + def visit_MatchStar(self, node: ast.MatchStar) -> ast.AST: + if node.name is not None: + node.name = self._rename(node.name) + return node + + def visit_MatchMapping(self, node: ast.MatchMapping) -> ast.AST: + node.keys = [self.visit(key) for key in node.keys] + node.patterns = [self.visit(pattern) for pattern in node.patterns] + if node.rest is not None: + node.rest = self._rename(node.rest) + return node diff --git a/tests/test_obfuscator.py b/tests/test_obfuscator.py index d098f53..76b5759 100644 --- a/tests/test_obfuscator.py +++ b/tests/test_obfuscator.py @@ -1,58 +1,73 @@ -from __future__ import annotations -import subprocess -import sys -from pathlib import Path -ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(ROOT)) -from patchwork import Obfuscator -EXAMPLES = ROOT / 'examples' -EXAMPLE_NAMES = ['hello', 'fizzbuzz', 'classes', 'stress'] -SEEDS = [1, 7, 42, 1234, 999999] - -def run(path: Path) -> str: - res = subprocess.run([sys.executable, str(path)], capture_output=True, text=True) - if res.returncode != 0: - raise RuntimeError(f'{path} exited {res.returncode}\nstderr:\n{res.stderr}') - return res.stdout - -def check_one(name: str, seed: int) -> None: - src_path = EXAMPLES / f'{name}.py' - expected = run(src_path) - obf_src = Obfuscator(seed=seed).obfuscate(src_path.read_text(encoding='utf-8')) - obf_path = EXAMPLES / f'{name}_test_{seed}.py' - obf_path.write_text(obf_src, encoding='utf-8') - try: - actual = run(obf_path) - finally: - obf_path.unlink(missing_ok=True) - if actual != expected: - raise AssertionError(f'{name} seed={seed}: output mismatch\n--- expected ---\n{expected}--- actual ---\n{actual}') - -def main() -> int: - fails = 0 - for name in EXAMPLE_NAMES: - for seed in SEEDS: - try: - check_one(name, seed) - print(f'ok: {name:10s} seed={seed}') - except Exception as e: - fails += 1 - print(f'FAIL: {name:10s} seed={seed}: {e}') - src = (EXAMPLES / 'hello.py').read_text(encoding='utf-8') - a = Obfuscator(seed=1).obfuscate(src) - b = Obfuscator(seed=2).obfuscate(src) - c = Obfuscator(seed=1).obfuscate(src) - if a == b: - print('FAIL: polymorphism (seed=1 vs seed=2 produced identical output)') - fails += 1 - else: - print('ok: polymorphism (seeds 1 vs 2 differ)') - if a != c: - print('FAIL: reproducibility (same seed gave different output)') - fails += 1 - else: - print('ok: reproducibility (seed=1 stable)') - print(f'\n{fails} failures' if fails else '\nall good') - return 1 if fails else 0 -if __name__ == '__main__': - raise SystemExit(main()) +from __future__ import annotations +import subprocess +import sys +from pathlib import Path +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) +from patchwork import Obfuscator +EXAMPLES = ROOT / 'examples' +EXAMPLE_NAMES = ['hello', 'fizzbuzz', 'classes', 'stress', 'modern', 'future_annotations'] +SEEDS = [1, 7, 42, 1234, 999999] + +def run(path: Path) -> str: + res = subprocess.run([sys.executable, str(path)], capture_output=True, text=True, timeout=20) + if res.returncode != 0: + raise RuntimeError(f'{path} exited {res.returncode}\nstderr:\n{res.stderr}') + return res.stdout + +def check_one(name: str, seed: int) -> None: + src_path = EXAMPLES / f'{name}.py' + expected = run(src_path) + obf_src = Obfuscator(seed=seed).obfuscate(src_path.read_text(encoding='utf-8')) + obf_path = EXAMPLES / f'{name}_test_{seed}.py' + obf_path.write_text(obf_src, encoding='utf-8') + try: + actual = run(obf_path) + finally: + obf_path.unlink(missing_ok=True) + if actual != expected: + raise AssertionError(f'{name} seed={seed}: output mismatch\n--- expected ---\n{expected}--- actual ---\n{actual}') + +def main() -> int: + fails = 0 + for name in EXAMPLE_NAMES: + for seed in SEEDS: + try: + check_one(name, seed) + print(f'ok: {name:10s} seed={seed}') + except Exception as e: + fails += 1 + print(f'FAIL: {name:10s} seed={seed}: {e}') + src = (EXAMPLES / 'hello.py').read_text(encoding='utf-8') + a = Obfuscator(seed=1).obfuscate(src) + b = Obfuscator(seed=2).obfuscate(src) + c = Obfuscator(seed=1).obfuscate(src) + if a == b: + print('FAIL: polymorphism (seed=1 vs seed=2 produced identical output)') + fails += 1 + else: + print('ok: polymorphism (seeds 1 vs 2 differ)') + if a != c: + print('FAIL: reproducibility (same seed gave different output)') + fails += 1 + else: + print('ok: reproducibility (seed=1 stable)') + future_src = (EXAMPLES / 'future_annotations.py').read_text(encoding='utf-8') + for kwargs in ({'encrypt_strings': False}, {'opaque_predicates': False, 'junk_branches': False}): + try: + obf = Obfuscator(seed=77, **kwargs).obfuscate(future_src) + obf_path = EXAMPLES / f'future_annotations_options_{len(kwargs)}.py' + obf_path.write_text(obf, encoding='utf-8') + try: + if run(obf_path) != run(EXAMPLES / 'future_annotations.py'): + raise AssertionError('output mismatch') + finally: + obf_path.unlink(missing_ok=True) + print(f'ok: future annotations options {kwargs}') + except Exception as e: + fails += 1 + print(f'FAIL: future annotations options {kwargs}: {e}') + print(f'\n{fails} failures' if fails else '\nall good') + return 1 if fails else 0 +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/test_stress_matrix.py b/tests/test_stress_matrix.py new file mode 100644 index 0000000..209aab0 --- /dev/null +++ b/tests/test_stress_matrix.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +import random +import subprocess +import sys +from pathlib import Path +import tempfile +import textwrap + +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +from patchwork import Obfuscator + + +GENERATED_SEEDS = [3, 11, 29, 101, 1009] +OBFUSCATION_SEEDS = [5, 17, 257] + + +def run(path: Path) -> str: + result = subprocess.run([sys.executable, str(path)], capture_output=True, text=True, timeout=20) + if result.returncode != 0: + raise RuntimeError(f"{path} exited {result.returncode}\nstderr:\n{result.stderr}") + return result.stdout + + +def generated_program(seed: int) -> str: + rng = random.Random(seed) + numbers = [rng.randint(-5000, 5000) for _ in range(16)] + strings = [f"s{idx}_{rng.randrange(1 << 20):x}" for idx in range(6)] + pairs = [(rng.randint(-20, 20), rng.randint(-20, 20)) for _ in range(8)] + return textwrap.dedent( + f""" + from __future__ import annotations + + NUMBERS = {numbers!r} + STRINGS = {strings!r} + PAIRS = {pairs!r} + + def fold(values): + acc = 0 + for idx, value in enumerate(values): + if idx % 3 == 0: + acc ^= value + elif idx % 3 == 1: + acc += value << (idx % 5) + else: + acc -= value >> (idx % 4) + return acc + + def classify(item): + match item: + case (0, y): + return f"zero:{{y}}" + case (x, y) if x == y: + return f"same:{{x}}" + case (x, y): + return f"pair:{{x + y}}" + case _: + return "other" + + def closure(base): + state = base + def inner(delta): + nonlocal state + state = (state + delta) ^ (delta << 2) + return state + return inner + + def run(): + fn = closure(7) + print("fold", fold(NUMBERS)) + print("strings", "|".join(s[::-1] for s in STRINGS)) + print("pairs", [classify(pair) for pair in PAIRS]) + print("closure", [fn(x) for x in range(6)]) + print("bytes", bytes(range(8)).hex()) + + if __name__ == "__main__": + run() + """ + ).lstrip() + + +def check_source(source: str, *, label: str, seed: int, **kwargs) -> None: + obf_source = Obfuscator(seed=seed, **kwargs).obfuscate(source) + with tempfile.TemporaryDirectory() as raw: + directory = Path(raw) + src_path = directory / f"{label}.py" + obf_path = directory / f"{label}_obf.py" + src_path.write_text(source, encoding="utf-8") + obf_path.write_text(obf_source, encoding="utf-8") + expected = run(src_path) + actual = run(obf_path) + if actual != expected: + raise AssertionError(f"{label} seed={seed} mismatch\n--- expected ---\n{expected}--- actual ---\n{actual}") + + +def main() -> int: + failures = 0 + matrix = [ + {}, + {"lazy_funcs": False}, + {"anti_debug": False}, + {"rename": False}, + {"encrypt_strings": False}, + {"opaque_predicates": False, "junk_branches": False}, + ] + for generated_seed in GENERATED_SEEDS: + source = generated_program(generated_seed) + for obfuscation_seed in OBFUSCATION_SEEDS: + for options in matrix: + label = f"generated_{generated_seed}_{obfuscation_seed}_{len(options)}" + try: + check_source(source, label=label, seed=obfuscation_seed, **options) + print(f"ok: {label} {options}") + except Exception as exc: + failures += 1 + print(f"FAIL: {label} {options}: {exc}") + print(f"\n{failures} failures" if failures else "\nstress matrix all good") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main())