swap techniques to ast-based

Author: davidteather

pyproject.toml

@@ -81,7 +81,7 @@ reportMissingTypeStubs = false
 [tool.coverage.run]
 branch = true
 source = ["python_obfuscator"]
-omit = ["tests/*"]
+omit = ["tests/*", "python_obfuscator/cli/*"]
 
 [tool.coverage.report]
 show_missing = true

python_obfuscator/__init__.py

@@ -1,4 +1,5 @@
-from .obfuscator import obfuscator
+from .config import ObfuscationConfig
+from .obfuscator import Obfuscator, obfuscate
 from .version import __version__
 
-__all__ = ["obfuscator", "__version__"]
+__all__ = ["Obfuscator", "ObfuscationConfig", "obfuscate", "__version__"]

python_obfuscator/cli/__init__.py

@@ -5,7 +5,9 @@
 import typer
 
 import python_obfuscator
-from python_obfuscator.techniques import one_liner as one_liner_technique
+from python_obfuscator.config import ObfuscationConfig
+from python_obfuscator.obfuscator import Obfuscator
+from python_obfuscator.techniques import all_technique_names
 from python_obfuscator.version import __version__
 
 DEFAULT_OUTPUT_DIR = "obfuscated"
@@ -43,24 +45,25 @@ def main(
             help="Print obfuscated code to stdout instead of writing a file.",
         ),
     ] = False,
-    include_one_liner: Annotated[
-        bool,
+    disable: Annotated[
+        list[str],
         typer.Option(
-            "--one-liner",
-            "-ol",
-            help="Include the one-liner obfuscation technique.",
+            "--disable",
+            "-d",
+            help=(
+                "Disable a technique by name.  May be repeated.  "
+                f"Available: {', '.join(sorted(all_technique_names()))}"
+            ),
         ),
-    ] = False,
+    ] = [],
 ) -> None:
     resolved = input_path.expanduser().resolve()
 
-    obfuscate = python_obfuscator.obfuscator()
-    remove = []
-    if not include_one_liner:
-        remove.append(one_liner_technique)
+    config = ObfuscationConfig.all_enabled().without(*disable)
+    obfuscator = Obfuscator(config)
 
     data = resolved.read_text()
-    obfuscated_data = obfuscate.obfuscate(data, remove_techniques=remove)
+    obfuscated_data = obfuscator.obfuscate(data)
 
     if stdout:
         typer.echo(obfuscated_data)

python_obfuscator/config.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ObfuscationConfig:
+    """Immutable configuration describing which techniques are enabled.
+
+    Prefer the factory classmethods over constructing directly::
+
+        # Enable everything registered
+        cfg = ObfuscationConfig.all_enabled()
+
+        # Enable a specific subset
+        cfg = ObfuscationConfig.only("variable_renamer", "string_hex_encoder")
+
+        # Start from all-enabled and exclude one
+        cfg = ObfuscationConfig.all_enabled().without("string_hex_encoder")
+    """
+
+    enabled_techniques: frozenset[str]
+
+    @classmethod
+    def all_enabled(cls) -> ObfuscationConfig:
+        from .techniques.registry import all_technique_names
+
+        return cls(enabled_techniques=all_technique_names())
+
+    @classmethod
+    def only(cls, *names: str) -> ObfuscationConfig:
+        return cls(enabled_techniques=frozenset(names))
+
+    def without(self, *names: str) -> ObfuscationConfig:
+        return ObfuscationConfig(
+            enabled_techniques=self.enabled_techniques - frozenset(names)
+        )
+
+    def with_added(self, *names: str) -> ObfuscationConfig:
+        return ObfuscationConfig(
+            enabled_techniques=self.enabled_techniques | frozenset(names)
+        )

python_obfuscator/helpers/random_datatype.py

@@ -1,23 +1,34 @@
+from __future__ import annotations
+
 import random
 import string
-import time
+from typing import Callable
 
 
 class RandomDataTypeGenerator:
-    def __init__(self):
-        self.generator_options = [self.random_string, self.random_int]
+    """Generates random Python literal values (str or int).
+
+    Pass a seeded :class:`random.Random` instance for reproducible output::
+
+        gen = RandomDataTypeGenerator(rng=random.Random(42))
+    """
+
+    def __init__(self, rng: random.Random | None = None) -> None:
+        self._rng = rng or random.Random()
+        self._generator_options: list[Callable[[], str | int]] = [
+            self.random_string,
+            self.random_int,
+        ]
 
-    def get_random(self):
-        return random.choice(self.generator_options)()
+    def get_random(self) -> str | int:
+        return self._rng.choice(self._generator_options)()
 
-    def random_string(self, length=79):
-        # Why is it 79 by default?
-        # See: https://stackoverflow.com/a/16920876/11472374
-        # As kirelagin commented readability is very important
+    def random_string(self, length: int = 79) -> str:
+        # 79 chars: see https://stackoverflow.com/a/16920876/11472374
         return "".join(
-            random.choice(string.ascii_lowercase + string.ascii_uppercase)
-            for i in range(length)
+            self._rng.choice(string.ascii_lowercase + string.ascii_uppercase)
+            for _ in range(length)
         )
 
-    def random_int(self):
-        return random.randint(random.randint(0, 300), random.randint(300, 999))
+    def random_int(self) -> int:
+        return self._rng.randint(self._rng.randint(0, 300), self._rng.randint(300, 999))

python_obfuscator/helpers/variable_name_generator.py

@@ -1,11 +1,21 @@
+from __future__ import annotations
+
 import random
 import string
-import time
+from typing import Callable
 
 
 class VariableNameGenerator:
-    def __init__(self):
-        self.generator_options = [
+    """Generates obfuscated-looking variable names.
+
+    Pass a seeded :class:`random.Random` instance for reproducible output::
+
+        gen = VariableNameGenerator(rng=random.Random(42))
+    """
+
+    def __init__(self, rng: random.Random | None = None) -> None:
+        self._rng = rng or random.Random()
+        self._generator_options: list[Callable[[int], str]] = [
             self.random_string,
             self.l_and_i,
             self.time_based,
@@ -14,33 +24,29 @@ def __init__(self):
             self.single_letter_a_lot,
         ]
 
-    def get_random(self, id):
-        return random.choice(self.generator_options)(id)
+    def get_random(self, id: int) -> str:
+        return self._rng.choice(self._generator_options)(id)
 
-    def random_string(self, id, length=79):
-        # Why is it 79 by default?
-        # See: https://stackoverflow.com/a/16920876/11472374
-        # As kirelagin commented readability is very important
+    def random_string(self, id: int, length: int = 79) -> str:
+        # 79 chars: see https://stackoverflow.com/a/16920876/11472374
         return "".join(
-            random.choice(string.ascii_letters) for i in range(length)
+            self._rng.choice(string.ascii_letters) for _ in range(length)
         ) + str(id)
 
-    def l_and_i(self, id):
-        return "".join(random.choice("Il") for i in range(id))
+    def l_and_i(self, id: int) -> str:
+        return "".join(self._rng.choice("Il") for _ in range(id))
 
-    def time_based(self, id):
-        return (
-            random.choice(string.ascii_letters)
-            + str(time.time()).replace(".", "")
-            + str(id)
-        )
+    def time_based(self, id: int) -> str:
+        # Use the rng to produce a large pseudo-time value so that this
+        # generator is fully deterministic when the rng is seeded.
+        pseudo_time = str(self._rng.randint(10**12, 10**13))
+        return self._rng.choice(string.ascii_letters) + pseudo_time + str(id)
 
-    def just_id(self, id):
-        # python doesn't work with numbers for variable names
-        return random.choice(string.ascii_letters) + str(id)
+    def just_id(self, id: int) -> str:
+        return self._rng.choice(string.ascii_letters) + str(id)
 
-    def scream(self, id):
-        return "".join(random.choice("Aa") for i in range(id))
+    def scream(self, id: int) -> str:
+        return "".join(self._rng.choice("Aa") for _ in range(id))
 
-    def single_letter_a_lot(self, id):
-        return random.choice(string.ascii_letters) * id
+    def single_letter_a_lot(self, id: int) -> str:
+        return self._rng.choice(string.ascii_letters) * id

python_obfuscator/obfuscator.py

@@ -1,10 +1,53 @@
-import logging
-from .techniques import obfuscate
+from __future__ import annotations
 
+import ast
 
-class obfuscator:
-    def __init__(self, logging_level=logging.error):
-        pass
+from .config import ObfuscationConfig
+from .techniques.registry import all_technique_names, get_transforms
 
-    def obfuscate(self, code, remove_techniques=[]):
-        return obfuscate(code, remove_techniques)
+
+def _validate_config(config: ObfuscationConfig) -> None:
+    unknown = config.enabled_techniques - all_technique_names()
+    if unknown:
+        raise ValueError(
+            f"Unknown technique(s): {sorted(unknown)}. "
+            f"Available: {sorted(all_technique_names())}"
+        )
+
+
+def obfuscate(source: str, config: ObfuscationConfig | None = None) -> str:
+    """Obfuscate *source* using the techniques described by *config*.
+
+    When *config* is ``None`` all registered techniques are applied.
+    """
+    resolved = config if config is not None else ObfuscationConfig.all_enabled()
+    _validate_config(resolved)
+
+    tree = ast.parse(source)
+    for transform_cls in get_transforms(resolved.enabled_techniques):
+        tree = transform_cls().apply(tree)
+    return ast.unparse(tree)
+
+
+class Obfuscator:
+    """Stateful wrapper that pre-builds and caches the transform pipeline.
+
+    Prefer this over the module-level :func:`obfuscate` when processing many
+    files with the same configuration, since the pipeline is validated and
+    sorted once at construction time.
+    """
+
+    def __init__(self, config: ObfuscationConfig | None = None) -> None:
+        self._config = config if config is not None else ObfuscationConfig.all_enabled()
+        _validate_config(self._config)
+        self._transforms = get_transforms(self._config.enabled_techniques)
+
+    @property
+    def config(self) -> ObfuscationConfig:
+        return self._config
+
+    def obfuscate(self, source: str) -> str:
+        tree = ast.parse(source)
+        for transform_cls in self._transforms:
+            tree = transform_cls().apply(tree)
+        return ast.unparse(tree)