perf: incremental parent-map maintenance

pyjsclear/__init__.py

@@ -12,41 +12,56 @@
 
 __all__ = ['Deobfuscator', 'deobfuscate', 'deobfuscate_file']
 
-__version__ = '0.1.5'
+__version__ = '0.1.6'
 
 
-def deobfuscate(code: str, max_iterations: int = 50) -> str:
+def deobfuscate(
+    code: str,
+    max_iterations: int = 50,
+    time_budget_seconds: float | None = None,
+) -> str:
     """Deobfuscate JavaScript code and return cleaned source.
 
     Args:
         code: JavaScript source code string.
         max_iterations: Maximum transform passes (default 50).
+        time_budget_seconds: Optional coarse wall-clock budget checked between transform
+            cycles (a stuck transform is not interrupted — enforce a hard bound externally);
+            restarts per nested decode layer, returning the best result so far on expiry. None disables it.
 
     Returns:
         Deobfuscated JavaScript source code.
     """
-    return Deobfuscator(code, max_iterations=max_iterations).execute()
+    return Deobfuscator(
+        code,
+        max_iterations=max_iterations,
+        time_budget_seconds=time_budget_seconds,
+    ).execute()
 
 
 def deobfuscate_file(
     input_path: str | Path,
     output_path: str | Path | None = None,
     max_iterations: int = 50,
+    time_budget_seconds: float | None = None,
 ) -> str | bool:
     """Deobfuscate a JavaScript file.
 
     Args:
         input_path: Path to input JS file.
         output_path: Path to write output (if None, returns string).
         max_iterations: Maximum transform passes.
+        time_budget_seconds: Optional coarse wall-clock budget checked between transform
+            cycles (a stuck transform is not interrupted — enforce a hard bound externally);
+            restarts per nested decode layer, returning the best result so far on expiry. None disables it.
 
     Returns:
         True if content changed (when output_path given), or the deobfuscated string.
     """
     with open(input_path, 'r', errors='replace') as input_file:
         code = input_file.read()
 
-    result = deobfuscate(code, max_iterations=max_iterations)
+    result = deobfuscate(code, max_iterations=max_iterations, time_budget_seconds=time_budget_seconds)
 
     if not output_path:
         return result

pyjsclear/deobfuscator.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import time
+
 from .generator import generate
 from .parser import parse
 from .scope import build_scope_tree
@@ -171,9 +173,15 @@ class Deobfuscator:
 
     _MAX_OUTER_CYCLES: int = 5
 
-    def __init__(self, code: str, max_iterations: int = 50) -> None:
+    def __init__(
+        self,
+        code: str,
+        max_iterations: int = 50,
+        time_budget_seconds: float | None = None,
+    ) -> None:
         self.original_code: str = code
         self.max_iterations: int = max_iterations
+        self.time_budget_seconds: float | None = time_budget_seconds
 
     def _run_pre_passes(self, code: str) -> str | None:
         """Detect whole-file encodings (JSFuck, AAEncode, etc.) and decode them.
@@ -209,7 +217,11 @@ def execute(self) -> str:
 
         decoded = self._run_pre_passes(code)
         if decoded:
-            recursive_deobfuscator = Deobfuscator(decoded, max_iterations=self.max_iterations)
+            recursive_deobfuscator = Deobfuscator(
+                decoded,
+                max_iterations=self.max_iterations,
+                time_budget_seconds=self.time_budget_seconds,
+            )
             return recursive_deobfuscator.execute()
 
         syntax_tree = self._try_parse_or_fallback(code)
@@ -235,13 +247,22 @@ def _try_parse_or_fallback(self, code: str) -> dict | str:
     def _transform_loop(self, syntax_tree: dict, code: str) -> str:
         """Run the outer generate-reparse convergence loop and post-passes.
 
+        When ``time_budget_seconds`` is set, elapsed wall-clock is polled at the top of each
+        outer cycle (not raised as an exception, which the pipeline's broad ``except Exception``
+        handlers would swallow); on expiry the loop stops and the best result so far is returned.
+        The budget restarts at each nested decode layer, so it is not a global deadline.
+
         Returns the best deobfuscated source produced.
         """
         previous_code = code
         last_changed_tree: dict | None = None
+        start_time = time.monotonic()
 
         try:
             for _cycle in range(self._MAX_OUTER_CYCLES):
+                if self.time_budget_seconds is not None and time.monotonic() - start_time >= self.time_budget_seconds:
+                    break
+
                 changed = self._run_ast_transforms(
                     syntax_tree,
                     code_size=len(previous_code),

pyjsclear/transforms/base.py

@@ -52,9 +52,24 @@ def get_parent_map(self) -> dict[int, tuple[dict, str, int | None]]:
         return self._parent_map
 
     def invalidate_parent_map(self) -> None:
-        """Invalidate the cached parent map after AST modifications."""
+        """Drop the cached parent map so the next lookup rebuilds it.
+
+        Prefer record_replacement() after an in-place swap: it keeps the map valid in O(1)
+        instead of forcing an O(N) rebuild on the next find_parent (quadratic over many swaps).
+        """
         self._parent_map = None
 
+    def record_replacement(self, replacement: dict, parent: dict, key: str, index: int | None) -> None:
+        """Patch the cached parent map after an in-place node swap.
+
+        Valid only for swaps that leave list indices unchanged; for insertions/removals that
+        shift indices, call invalidate_parent_map() instead. The detached original subtree's
+        entries go stale (don't look them up), and descendants of the replacement aren't
+        registered, so find_parent returns None on them until a full rebuild.
+        """
+        if self._parent_map is not None:
+            self._parent_map[id(replacement)] = (parent, key, index)
+
     def find_parent(self, target_node: dict) -> tuple[dict, str, int | None] | None:
         """Find the parent of a node using the parent map."""
         parent_map = self.get_parent_map()

pyjsclear/transforms/class_static_resolver.py

@@ -227,3 +227,4 @@ def _replace_in_parent(
             parent[key][index] = replacement
         else:
             parent[key] = replacement
+        self.record_replacement(replacement, parent, key, index)

pyjsclear/transforms/object_simplifier.py

@@ -180,7 +180,7 @@ def _replace_node(self, target: dict, replacement: dict) -> bool:
             parent[key][index] = replacement
         else:
             parent[key] = replacement
-        self.invalidate_parent_map()
+        self.record_replacement(replacement, parent, key, index)
         return True
 
     def _inline_function(self, function_node: dict, arguments: list[dict]) -> dict | None:

pyjsclear/transforms/string_revealer.py

@@ -1297,7 +1297,7 @@ def _replace_node_in_ast(self, target: dict, replacement: dict) -> None:
                 parent[key][index] = replacement
             else:
                 parent[key] = replacement
-            self.invalidate_parent_map()
+            self.record_replacement(replacement, parent, key, index)
 
     # ================================================================
     # Strategy 3: Simple static array unpacking

pyproject.toml

@@ -9,7 +9,9 @@ description = "Pure Python JavaScript deobfuscator"
 readme = "README.md"
 license = "Apache-2.0"
 requires-python = ">=3.11"
-dependencies = ["esprima2>=5.0.1"]
+# esprima2 capped below 5.0.2: it switched optional-chaining to the ChainExpression AST wrapper
+# the generator/transforms don't yet support (tracked as a follow-up).
+dependencies = ["esprima2>=5.0.1,<5.0.2"]
 keywords = ["javascript", "deobfuscator", "deobfuscation", "security", "malware-analysis", "ast"]
 authors = [
     {name = "Intezer Labs", email = "info@intezer.com"},

requirements.txt

@@ -1 +1,3 @@
-esprima2>=5.0.0
+# Cap below 5.0.2: it changed optional-chaining (a?.b) parsing to the ESTree ChainExpression
+# wrapper, which the generator and transforms don't yet support (tracked as a follow-up).
+esprima2>=5.0.0,<5.0.2

tests/unit/deobfuscate_budget_test.py

@@ -0,0 +1,88 @@
+"""Tests for the optional wall-clock budget on deobfuscation (TKT-16478).
+
+The budget is deliberately coarse: it is checked inline between outer transform
+cycles (never via exceptions, which the pipeline's broad ``except Exception``
+handlers would swallow). On expiry the best result so far is returned. The
+budget restarts at each nested decode layer (JSFuck/eval-packed recursion).
+"""
+
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+import pyjsclear
+from pyjsclear.deobfuscator import Deobfuscator
+
+
+class TestDeobfuscateBudget:
+    def test_deobfuscate_accepts_a_time_budget_and_returns_a_string(self):
+        # Budget is honored between outer cycles; a tiny budget still returns valid source.
+        out = pyjsclear.deobfuscate('var a=1;var b=2;console.log(a+b);', time_budget_seconds=0.001)
+        assert isinstance(out, str)
+        assert out
+
+    def test_deobfuscate_file_accepts_a_time_budget(self, tmp_path):
+        input_file = tmp_path / 'input.js'
+        input_file.write_text('var x = 1;')
+
+        result = pyjsclear.deobfuscate_file(str(input_file), time_budget_seconds=0.001)
+        assert isinstance(result, str)
+
+
+class TestDeobfuscatorBudgetInternals:
+    def test_budget_defaults_to_none_and_is_stored_when_given(self):
+        assert Deobfuscator('var x = 1;').time_budget_seconds is None
+        assert Deobfuscator('var x = 1;', time_budget_seconds=2.5).time_budget_seconds == 2.5
+
+    def test_budget_expiry_mid_loop_stops_after_one_cycle_and_returns_source(self):
+        # Deterministic mid-loop expiry: fake the clock so the budget check
+        # passes at cycle 0 and trips at cycle 1. Expected monotonic() calls on
+        # this path (verified against deobfuscator.py, the only module using
+        # time): start capture -> 0.0, cycle-0 check -> 0.0, cycle-1 check ->
+        # 10.0. The `time` attribute is patched in the deobfuscator module
+        # namespace only, so nothing else consumes the side_effect values.
+        fake_time = MagicMock()
+        fake_time.monotonic.side_effect = [0.0, 0.0, 10.0]
+        deobfuscator = Deobfuscator('var a = 1; var b = a; console.log(b);', time_budget_seconds=5.0)
+
+        with (
+            patch.object(Deobfuscator, '_run_ast_transforms', return_value=True) as run_transforms_mock,
+            patch('pyjsclear.deobfuscator.time', fake_time),
+        ):
+            result = deobfuscator.execute()
+
+        # Exactly one transform cycle ran before the budget expired.
+        assert run_transforms_mock.call_count == 1
+        # All three clock reads happened: the cycle-1 check executed and
+        # tripped (i.e. the loop did not end early for another reason).
+        assert fake_time.monotonic.call_count == 3
+        assert isinstance(result, str)
+        assert result
+
+    def test_exhausted_budget_skips_transform_cycles_entirely(self):
+        # The check sits at the top of each outer cycle: a zero budget is
+        # already expired at cycle 0, so no transform cycle should run at all,
+        # yet execute() must still return valid source (best-so-far path).
+        deobfuscator = Deobfuscator('var a = 1; var b = a; console.log(b);', time_budget_seconds=0.0)
+
+        with patch.object(Deobfuscator, '_run_ast_transforms', return_value=False) as run_transforms_mock:
+            result = deobfuscator.execute()
+
+        run_transforms_mock.assert_not_called()
+        assert isinstance(result, str)
+        assert result
+
+    def test_recursive_pre_pass_deobfuscator_inherits_the_budget(self):
+        # When a pre-pass decodes a nested layer (JSFuck/eval-packed), the
+        # recursively constructed Deobfuscator must receive the same budget.
+        outer = Deobfuscator('outer code', time_budget_seconds=1.5)
+        nested = MagicMock()
+        nested.execute.return_value = 'const x = 1;'
+
+        with (
+            patch.object(Deobfuscator, '_run_pre_passes', return_value='var x = 1;'),
+            patch('pyjsclear.deobfuscator.Deobfuscator', return_value=nested) as constructor_mock,
+        ):
+            result = outer.execute()
+
+        constructor_mock.assert_called_once_with('var x = 1;', max_iterations=50, time_budget_seconds=1.5)
+        assert result == 'const x = 1;'

tests/unit/incremental_parent_map_test.py

@@ -0,0 +1,73 @@
+"""Tests for incremental parent-map maintenance across transform call sites (TKT-16478).
+
+Transforms replace nodes in place while iterating; each replacement calls
+record_replacement() to patch the cached parent map in O(1) instead of
+invalidate_parent_map(), which would force a full O(N) rebuild on the next
+find_parent lookup -- quadratic overall when many nodes are replaced in one
+pass. These tests assert the observable effect: build_parent_map() is called
+at most once per transform execution, no matter how many nodes it replaces.
+"""
+
+from unittest.mock import patch
+
+from pyjsclear.parser import parse
+from pyjsclear.transforms.class_static_resolver import ClassStaticResolver
+from pyjsclear.transforms.object_simplifier import ObjectSimplifier
+from pyjsclear.transforms.string_revealer import StringRevealer
+from pyjsclear.traverser import build_parent_map
+
+
+_REFERENCE_COUNT = 200
+
+
+class TestObjectSimplifierIncrementalParentMap:
+    def test_build_parent_map_called_once_when_many_properties_are_inlined(self):
+        # Arrange
+        properties = ', '.join(f'k{i}: {i}' for i in range(_REFERENCE_COUNT))
+        accesses = '; '.join(f'console.log(o.k{i})' for i in range(_REFERENCE_COUNT))
+        ast = parse(f'const o = {{{properties}}}; {accesses};')
+        transform = ObjectSimplifier(ast)
+
+        # Act
+        with patch('pyjsclear.transforms.base.build_parent_map', wraps=build_parent_map) as mock_build_parent_map:
+            changed = transform.execute()
+
+        # Assert
+        assert changed is True
+        assert mock_build_parent_map.call_count == 1
+
+
+class TestStringRevealerIncrementalParentMap:
+    def test_build_parent_map_called_once_when_many_array_accesses_are_replaced(self):
+        # Arrange
+        elements = ', '.join(f'"s{i}"' for i in range(_REFERENCE_COUNT))
+        accesses = '; '.join(f'f(arr[{i}])' for i in range(_REFERENCE_COUNT))
+        ast = parse(f'var arr = [{elements}]; {accesses};')
+        transform = StringRevealer(ast)
+
+        # Act
+        with patch('pyjsclear.transforms.base.build_parent_map', wraps=build_parent_map) as mock_build_parent_map:
+            changed = transform.execute()
+
+        # Assert
+        assert changed is True
+        assert mock_build_parent_map.call_count <= 1
+
+
+class TestClassStaticResolverIncrementalParentMap:
+    def test_build_parent_map_called_once_when_many_static_properties_are_inlined(self):
+        # Arrange
+        accesses = '; '.join(f'console.log(C.X + {i})' for i in range(_REFERENCE_COUNT))
+        ast = parse(f'var C = class {{}}; C.X = 100; {accesses};')
+        transform = ClassStaticResolver(ast)
+
+        # Act
+        with patch('pyjsclear.transforms.base.build_parent_map', wraps=build_parent_map) as mock_build_parent_map:
+            changed = transform.execute()
+
+        # Assert
+        assert changed is True
+        assert mock_build_parent_map.call_count == 1
+        # The transform invalidates the cache once the traversal completes,
+        # so the next lookup rebuilds fresh rather than reusing stale entries.
+        assert transform._parent_map is None