Add pooled literal obfuscation

patchwork/__init__.py

@@ -1,7 +1,7 @@
 from .core import Obfuscator, obfuscate, obfuscate_file
 from .audit import analyze_source, build_manifest, build_stats, verify_manifest
 
-__version__ = '0.4.0'
+__version__ = '0.5.0'
 
 __all__ = [
     'Obfuscator',

patchwork/audit.py

@@ -77,7 +77,7 @@ def _tool_version() -> str:
     try:
         return version("patchwork")
     except PackageNotFoundError:
-        return "0.4.0"
+        return "0.5.0"
 
 
 def _call_name(node: ast.AST) -> str | None:

patchwork/core.py

@@ -12,6 +12,7 @@ from .util import make_rng
 _SEED_NAME = '_pw_seed_value'
 _DEC_S_NAME = '_pw_dec_str'
 _DEC_B_NAME = '_pw_dec_bytes'
+_CONST_POOL_NAME = '_pw_literal_pool'
 _RESOLVE_NAME = '__pw_resolve_lazy__'
 
 def _insert_runtime_stmts(tree: ast.Module, stmts: list[ast.stmt]) -> None:
@@ -68,9 +69,14 @@ class Obfuscator:
         if self.obfuscate_numbers:
             tree = NumberObfuscator(self.rng).visit(tree)
         if self.encrypt_strings:
-            enc = StringEncryptor(self.rng, _DEC_S_NAME, _DEC_B_NAME)
+            enc = StringEncryptor(self.rng, _DEC_S_NAME, _DEC_B_NAME, _CONST_POOL_NAME)
             tree = enc.visit(tree)
-            _insert_runtime_stmts(tree, build_decrypt_helper(_DEC_S_NAME, _DEC_B_NAME))
+            runtime_stmts = []
+            pool_stmt = enc.build_pool_stmt()
+            if pool_stmt is not None:
+                runtime_stmts.append(pool_stmt)
+            runtime_stmts.extend(build_decrypt_helper(_DEC_S_NAME, _DEC_B_NAME))
+            _insert_runtime_stmts(tree, runtime_stmts)
         if self.rename:
             renamer = IdentifierRenamer(self.rng, keep=skip | self.keep_extra | {_RESOLVE_NAME})
             tree = renamer.visit(tree)

patchwork/transforms/identifiers.py

@@ -179,6 +179,8 @@ class IdentifierRenamer(ast.NodeTransformer):
             if alias.asname:
                 alias.asname = self._rename(alias.asname)
             else:
+                if alias.name in self.reserved:
+                    continue
                 if alias.name not in self.mapping:
                     self.mapping[alias.name] = self._new()
                 alias.asname = self.mapping[alias.name]
@@ -193,6 +195,8 @@ class IdentifierRenamer(ast.NodeTransformer):
             if alias.asname:
                 alias.asname = self._rename(alias.asname)
             else:
+                if alias.name in self.reserved:
+                    continue
                 if alias.name not in self.mapping:
                     self.mapping[alias.name] = self._new()
                 alias.asname = self.mapping[alias.name]

patchwork/transforms/numbers.py

@@ -1,70 +1,93 @@
-from __future__ import annotations
-import ast
-import random
-
-class NumberObfuscator(ast.NodeTransformer):
-
-    def __init__(self, rng: random.Random, prob: float=0.85):
-        self.rng = rng
-        self.prob = prob
-
-    def visit_Constant(self, node: ast.Constant) -> ast.AST:
-        v = node.value
-        if isinstance(v, bool):
-            return node
-        if isinstance(v, int):
-            if self.rng.random() <= self.prob:
-                return self._obfuscate(v)
-        return node
-
-    def visit_match_case(self, node: ast.match_case) -> ast.AST:
-        if node.guard is not None:
-            node.guard = self.visit(node.guard)
-        node.body = [self.visit(stmt) for stmt in node.body]
-        return node
-
-    def _obfuscate(self, n: int) -> ast.AST:
-        choices = ['xor', 'sum', 'shift', 'frombytes', 'affine', 'split']
-        kind = self.rng.choice(choices)
-        try:
-            return getattr(self, f'_t_{kind}')(n)
-        except Exception:
-            return ast.Constant(value=n)
-
-    def _t_xor(self, n: int) -> ast.AST:
-        k = self.rng.randrange(1, 1 << 32)
-        return self._parse_expr(f'(({n ^ k}) ^ {k})')
-
-    def _t_sum(self, n: int) -> ast.AST:
-        a = self.rng.randint(-(1 << 16), 1 << 16)
-        return self._parse_expr(f'(({n - a}) + ({a}))')
-
-    def _t_shift(self, n: int) -> ast.AST:
-        if abs(n) >= 1 << 60:
-            return self._t_xor(n)
-        s = self.rng.randint(1, 8)
-        return self._parse_expr(f'(({n} << {s}) >> {s})')
-
-    def _t_frombytes(self, n: int) -> ast.AST:
-        signed = n < 0
-        bit_len = max(1, n.bit_length() + (1 if signed else 0))
-        byte_len = (bit_len + 7) // 8
-        try:
-            data = n.to_bytes(byte_len, 'big', signed=signed)
-        except OverflowError:
-            return self._t_xor(n)
-        signed_arg = 'True' if signed else 'False'
-        return self._parse_expr(f"int.from_bytes({data!r}, 'big', signed={signed_arg})")
-
-    def _t_affine(self, n: int) -> ast.AST:
-        a = self.rng.randint(2, 64)
-        b = self.rng.randint(-(1 << 12), 1 << 12)
-        return self._parse_expr(f'((({n * a + b}) - ({b})) // {a})')
-
-    def _t_split(self, n: int) -> ast.AST:
-        p = self.rng.randint(-(1 << 16), 1 << 16)
-        q = n - p
-        return self._parse_expr(f'({p} + ({q}))')
-
-    def _parse_expr(self, src: str) -> ast.AST:
-        return ast.parse(src, mode='eval').body
+from __future__ import annotations
+import ast
+import random
+
+class NumberObfuscator(ast.NodeTransformer):
+
+    def __init__(self, rng: random.Random, prob: float=0.85):
+        self.rng = rng
+        self.prob = prob
+
+    def visit_Constant(self, node: ast.Constant) -> ast.AST:
+        v = node.value
+        if isinstance(v, bool):
+            return node
+        if isinstance(v, int):
+            if self.rng.random() <= self.prob:
+                return self._obfuscate(v)
+        return node
+
+    def visit_match_case(self, node: ast.match_case) -> ast.AST:
+        if node.guard is not None:
+            node.guard = self.visit(node.guard)
+        node.body = [self.visit(stmt) for stmt in node.body]
+        return node
+
+    def _obfuscate(self, n: int) -> ast.AST:
+        choices = ['xor', 'sum', 'shift', 'frombytes', 'affine', 'split', 'invert', 'bitmask', 'divmod', 'table']
+        kind = self.rng.choice(choices)
+        try:
+            return getattr(self, f'_t_{kind}')(n)
+        except Exception:
+            return ast.Constant(value=n)
+
+    def _t_xor(self, n: int) -> ast.AST:
+        k = self.rng.randrange(1, 1 << 32)
+        return self._parse_expr(f'(({n ^ k}) ^ {k})')
+
+    def _t_sum(self, n: int) -> ast.AST:
+        a = self.rng.randint(-(1 << 16), 1 << 16)
+        return self._parse_expr(f'(({n - a}) + ({a}))')
+
+    def _t_shift(self, n: int) -> ast.AST:
+        if abs(n) >= 1 << 60:
+            return self._t_xor(n)
+        s = self.rng.randint(1, 8)
+        return self._parse_expr(f'(({n} << {s}) >> {s})')
+
+    def _t_frombytes(self, n: int) -> ast.AST:
+        signed = n < 0
+        bit_len = max(1, n.bit_length() + (1 if signed else 0))
+        byte_len = (bit_len + 7) // 8
+        try:
+            data = n.to_bytes(byte_len, 'big', signed=signed)
+        except OverflowError:
+            return self._t_xor(n)
+        signed_arg = 'True' if signed else 'False'
+        return self._parse_expr(f"int.from_bytes({data!r}, 'big', signed={signed_arg})")
+
+    def _t_affine(self, n: int) -> ast.AST:
+        a = self.rng.randint(2, 64)
+        b = self.rng.randint(-(1 << 12), 1 << 12)
+        return self._parse_expr(f'((({n * a + b}) - ({b})) // {a})')
+
+    def _t_split(self, n: int) -> ast.AST:
+        p = self.rng.randint(-(1 << 16), 1 << 16)
+        q = n - p
+        return self._parse_expr(f'({p} + ({q}))')
+
+    def _t_invert(self, n: int) -> ast.AST:
+        return self._parse_expr(f'(~({~n}))')
+
+    def _t_bitmask(self, n: int) -> ast.AST:
+        mask = self.rng.randrange(1, 1 << 32)
+        return self._parse_expr(f'(({n} & {mask}) | ({n} & ~{mask}))')
+
+    def _t_divmod(self, n: int) -> ast.AST:
+        d = self.rng.randint(2, 97)
+        q, r = divmod(n, d)
+        return self._parse_expr(f'(({q} * {d}) + {r})')
+
+    def _t_table(self, n: int) -> ast.AST:
+        values = [
+            self.rng.randint(-(1 << 20), 1 << 20),
+            self.rng.randint(-(1 << 20), 1 << 20),
+            n,
+            self.rng.randint(-(1 << 20), 1 << 20),
+        ]
+        self.rng.shuffle(values)
+        idx = values.index(n)
+        return self._parse_expr(f'({values!r}[{idx}])')
+
+    def _parse_expr(self, src: str) -> ast.AST:
+        return ast.parse(src, mode='eval').body

patchwork/transforms/strings.py

@@ -5,10 +5,12 @@ from collections.abc import Sequence
 
 class StringEncryptor(ast.NodeTransformer):
 
-    def __init__(self, rng: random.Random, decrypt_str_name: str, decrypt_bytes_name: str):
+    def __init__(self, rng: random.Random, decrypt_str_name: str, decrypt_bytes_name: str, pool_name: str):
         self.rng = rng
         self.dec_s = decrypt_str_name
         self.dec_b = decrypt_bytes_name
+        self.pool_name = pool_name
+        self.entries: list[tuple[list[bytes], list[bytes], list[int], list[int], int, int]] = []
 
     def _strip_docstring(self, body: list[ast.stmt]) -> list[ast.stmt]:
         if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant) and isinstance(body[0].value.value, str):
@@ -99,19 +101,29 @@ class StringEncryptor(ast.NodeTransformer):
     def _int_tuple(self, values: Sequence[int]) -> ast.Tuple:
         return ast.Tuple(elts=[ast.Constant(value=value) for value in values], ctx=ast.Load())
 
-    def _encoded_args(self, data: bytes, key: bytes) -> list[ast.AST]:
+    def _encoded_entry(self, data: bytes, key: bytes) -> tuple[list[bytes], list[bytes], list[int], list[int], int, int]:
         enc = self._xor(data, key)
         rotation = self.rng.randrange(len(enc)) if enc else 0
         rotated = self._rotate_left(enc, rotation)
         data_chunks, data_order = self._shuffle_chunks(self._split_chunks(rotated))
         key_chunks, key_order = self._shuffle_chunks(self._split_chunks(key))
-        return [
-            self._bytes_tuple(data_chunks),
-            self._bytes_tuple(key_chunks),
-            self._int_tuple(data_order),
-            self._int_tuple(key_order),
-            ast.Constant(value=rotation),
-        ]
+        mode = self.rng.randrange(3)
+        return data_chunks, key_chunks, data_order, key_order, rotation, mode
+
+    def _store_entry(self, data: bytes, key: bytes) -> int:
+        idx = len(self.entries)
+        self.entries.append(self._encoded_entry(data, key))
+        return idx
+
+    def _pool_ref(self, idx: int, decoder_name: str) -> ast.Call:
+        return ast.Call(
+            func=ast.Name(id=decoder_name, ctx=ast.Load()),
+            args=[
+                ast.Name(id=self.pool_name, ctx=ast.Load()),
+                ast.Constant(value=idx),
+            ],
+            keywords=[],
+        )
 
     def _encrypt_str(self, s: str) -> ast.AST:
         try:
@@ -119,11 +131,34 @@ class StringEncryptor(ast.NodeTransformer):
         except UnicodeEncodeError:
             return ast.Constant(value=s)
         key = self._make_key()
-        return ast.Call(func=ast.Name(id=self.dec_s, ctx=ast.Load()), args=self._encoded_args(data, key), keywords=[])
+        return self._pool_ref(self._store_entry(data, key), self.dec_s)
 
     def _encrypt_bytes(self, b: bytes) -> ast.AST:
         key = self._make_key()
-        return ast.Call(func=ast.Name(id=self.dec_b, ctx=ast.Load()), args=self._encoded_args(b, key), keywords=[])
+        return self._pool_ref(self._store_entry(b, key), self.dec_b)
+
+    def build_pool_stmt(self) -> ast.stmt | None:
+        if not self.entries:
+            return None
+        entry_nodes = []
+        for data_chunks, key_chunks, data_order, key_order, rotation, mode in self.entries:
+            entry_nodes.append(
+                ast.Tuple(
+                    elts=[
+                        self._bytes_tuple(data_chunks),
+                        self._bytes_tuple(key_chunks),
+                        self._int_tuple(data_order),
+                        self._int_tuple(key_order),
+                        ast.Constant(value=rotation),
+                        ast.Constant(value=mode),
+                    ],
+                    ctx=ast.Load(),
+                )
+            )
+        return ast.Assign(
+            targets=[ast.Name(id=self.pool_name, ctx=ast.Store())],
+            value=ast.Tuple(elts=entry_nodes, ctx=ast.Load()),
+        )
 
 def build_decrypt_helper(decrypt_str_name: str, decrypt_bytes_name: str) -> list[ast.stmt]:
     src = f'''
@@ -136,12 +171,32 @@ def _pw_unrotate(_data, _amount):
     _amount %= len(_data)
     return _data[-_amount:] + _data[:-_amount] if _amount else _data
 
-def {decrypt_bytes_name}(_d, _k, _do, _ko, _rot):
-    _data = _pw_unrotate(_pw_join_chunks(_d, _do), _rot)
-    _key = _pw_join_chunks(_k, _ko)
+def _pw_xor_a(_data, _key):
     return bytes(_b ^ _key[_i % len(_key)] for _i, _b in enumerate(_data))
 
-def {decrypt_str_name}(_d, _k, _do, _ko, _rot):
-    return {decrypt_bytes_name}(_d, _k, _do, _ko, _rot).decode()
+def _pw_xor_b(_data, _key):
+    _out = bytearray(len(_data))
+    for _i, _b in enumerate(_data):
+        _out[_i] = _b ^ _key[_i % len(_key)]
+    return bytes(_out)
+
+def _pw_xor_c(_data, _key):
+    return bytes(map(lambda _p: _p[1] ^ _key[_p[0] % len(_key)], enumerate(_data)))
+
+def _pw_decode_entry(_pool, _idx):
+    _d, _k, _do, _ko, _rot, _mode = _pool[_idx]
+    _data = _pw_unrotate(_pw_join_chunks(_d, _do), _rot)
+    _key = _pw_join_chunks(_k, _ko)
+    if _mode == 1:
+        return _pw_xor_b(_data, _key)
+    if _mode == 2:
+        return _pw_xor_c(_data, _key)
+    return _pw_xor_a(_data, _key)
+
+def {decrypt_bytes_name}(_pool, _idx):
+    return _pw_decode_entry(_pool, _idx)
+
+def {decrypt_str_name}(_pool, _idx):
+    return _pw_decode_entry(_pool, _idx).decode()
 '''
     return ast.parse(src).body