diff --git a/pyjsclear/__init__.py b/pyjsclear/__init__.py index 52fa0af..e03cbb3 100644 --- a/pyjsclear/__init__.py +++ b/pyjsclear/__init__.py @@ -12,26 +12,38 @@ __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. @@ -39,6 +51,9 @@ def deobfuscate_file( 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. @@ -46,7 +61,7 @@ def deobfuscate_file( 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 diff --git a/pyjsclear/deobfuscator.py b/pyjsclear/deobfuscator.py index fb59b50..74041b1 100644 --- a/pyjsclear/deobfuscator.py +++ b/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), diff --git a/pyjsclear/transforms/base.py b/pyjsclear/transforms/base.py index 9dbd46c..739375c 100644 --- a/pyjsclear/transforms/base.py +++ b/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() diff --git a/pyjsclear/transforms/class_static_resolver.py b/pyjsclear/transforms/class_static_resolver.py index d90c4ef..8f505f6 100644 --- a/pyjsclear/transforms/class_static_resolver.py +++ b/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) diff --git a/pyjsclear/transforms/object_simplifier.py b/pyjsclear/transforms/object_simplifier.py index 1049ff6..bd4561c 100644 --- a/pyjsclear/transforms/object_simplifier.py +++ b/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: diff --git a/pyjsclear/transforms/string_revealer.py b/pyjsclear/transforms/string_revealer.py index 5433f4c..c354fa7 100644 --- a/pyjsclear/transforms/string_revealer.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 6c25ca8..7eb8c91 100644 --- a/pyproject.toml +++ b/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"}, diff --git a/requirements.txt b/requirements.txt index 843a8fa..a8d1c1b 100644 --- a/requirements.txt +++ b/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 diff --git a/tests/unit/deobfuscate_budget_test.py b/tests/unit/deobfuscate_budget_test.py new file mode 100644 index 0000000..a542f0c --- /dev/null +++ b/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;' diff --git a/tests/unit/incremental_parent_map_test.py b/tests/unit/incremental_parent_map_test.py new file mode 100644 index 0000000..179426c --- /dev/null +++ b/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 diff --git a/tests/unit/transforms/base_test.py b/tests/unit/transforms/base_test.py index 5752f3b..78d5f6d 100644 --- a/tests/unit/transforms/base_test.py +++ b/tests/unit/transforms/base_test.py @@ -1,6 +1,8 @@ import pytest +from pyjsclear.parser import parse from pyjsclear.transforms.base import Transform +from pyjsclear.traverser import build_parent_map class TestTransformInit: @@ -52,6 +54,60 @@ def test_set_changed_is_idempotent(self): assert transform.has_changed() is True +class TestRecordReplacement: + def test_find_parent_returns_recorded_entry_without_rebuild(self): + ast = parse('f(1);') + transform = Transform(ast) + cached_map = transform.get_parent_map() + + call = ast['body'][0]['expression'] + replacement = {'type': 'Literal', 'value': 2, 'raw': '2'} + call['arguments'][0] = replacement + transform.record_replacement(replacement, call, 'arguments', 0) + + parent, key, index = transform.find_parent(replacement) + assert parent is call + assert key == 'arguments' + assert index == 0 + # The lookup was served by the patched cache, not a full rebuild. + assert transform._parent_map is cached_map + + def test_noop_when_map_not_built(self): + ast = parse('f(1);') + transform = Transform(ast) + + call = ast['body'][0]['expression'] + replacement = {'type': 'Literal', 'value': 2, 'raw': '2'} + call['arguments'][0] = replacement + transform.record_replacement(replacement, call, 'arguments', 0) + + assert transform._parent_map is None + + def test_recorded_entry_matches_rebuild_for_list_child(self): + ast = parse('f(1);') + transform = Transform(ast) + transform.get_parent_map() + + call = ast['body'][0]['expression'] + replacement = {'type': 'Literal', 'value': 2, 'raw': '2'} + call['arguments'][0] = replacement + transform.record_replacement(replacement, call, 'arguments', 0) + + assert transform.find_parent(replacement) == build_parent_map(ast)[id(replacement)] + + def test_recorded_entry_matches_rebuild_for_dict_child(self): + ast = parse('x + 1;') + transform = Transform(ast) + transform.get_parent_map() + + statement = ast['body'][0] + replacement = {'type': 'Literal', 'value': 2, 'raw': '2'} + statement['expression'] = replacement + transform.record_replacement(replacement, statement, 'expression', None) + + assert transform.find_parent(replacement) == build_parent_map(ast)[id(replacement)] + + class TestTransformRebuildScope: def test_class_default_is_false(self): assert Transform.rebuild_scope is False