Implement Extended Masking#143
Draft
Quinntyx wants to merge 13 commits into
Draft
Conversation
Adds pylingual/masking/container_recovery/ package mirroring control_flow_reconstruction/ structure: - segment.py: Recovery + Segment dataclasses - stack_analysis.py: stack_effect (with 3.12+ NULL-push fixes), analyze_stack, parse_*_recursive - recovery.py: @register_recovery_strategy decorator + dispatch - utils.py: build-size and pattern helpers - strategies/: 12 recovery strategies split by type (dict, list, set, tuple, const), glob auto-imported Includes adjacent container parsing fix (cur_elem_id <= 0 break condition) and frozenset-to-set conversion in recover_const_container.
Rewrites container construction bytecode into single LOAD_CONST instructions. For each code object, analyzes the stack to identify container construction spans, recovers their values via strategy-based recovery, and replaces the bytecode span with a single LOAD_CONST referencing the materialized value in co_consts. Key behaviors: - Walks Segment tree recursively: if parent recovery is incomplete, tries children right-to-left. Only collapses children with container tags. - eval() with empty namespace (no frozenset builtin) - Co_consts deduplication: is (identity) then == - All exceptions propagate (no catch) - Skips empty containers after eval - preprocess() returns None
Add Preprocessor.preprocess() call in mask_bytecode() before create_global_masker(). Preprocessor exceptions are caught with graceful degradation — decompilation continues without container recovery if preprocessing fails. Separate log line for the preprocessing stage.
- Empty container skip uses len(value) == 0 instead of not value
(not value was truthy-check, incorrectly skipped non-empty falsy
containers like [0] or [""])
- co_consts dedup uses type-aware equality to prevent set/frozenset
collision (set({1,2}) == frozenset({1,2}) is True in Python)
- TypeSensitiveDict: use repr() for unhashable keys via try/except hash(),
eval() with full builtins to restore
- Remove replace_list() from get_model_view() — containers get single mask
token with type wrapper (e.g. list(<mask_42>))
- Remove KWARG padding logic (orphaned after replace_list removal)
- Add dict to container type check in get_model_view()
- Remove container flattening from create_global_masker() in model_disasm.py
(slice decomposition kept)
- Replace global_tab.update({k:v}) with global_tab[k]=v (dict literal fails
on unhashable keys)
Add 30 fixture files ported from recursive-decompile-demo and two test modules: - test_preprocessor.py: parametrized co_consts comparison test that extracts container literals from AST, runs the Preprocessor, and verifies each value appears in co_consts (with recursive search for nested containers). - test_edge_cases.py: 11 edge case tests covering empty containers, falsy non-empty containers, set/frozenset recovery, nested containers, deduplication, and type-aware dedup. 36 passed, 5 skipped (fixtures with no non-empty container literals). Closes syssec-utd#5
Replace try/except hash() with isinstance guard for container types
(list, set, dict, tuple, frozenset) to prevent custom __repr__ from
reaching repr(). Replace eval(value, {__builtins__: builtins}) with
ast.literal_eval in _key_restore.
…te eval Move the Preprocessor class and container_recovery package from pylingual/masking/container_recovery/ to pylingual/preprocessor/, presenting the Preprocessor as a first-class pass with container_recovery as a subpackage it uses. Refactor Recovery dataclass to carry actual Python objects (Recovery.value) instead of string expressions (Recovery.expr). All recovery strategies now build objects directly, and _collapse_segment uses the value directly instead of calling eval(). Remove INDENT and string formatting from strategies that no longer need them. Add comment noting xdis oppush/oppop tables are incorrect for null-pushing and variable-size opcodes, explaining why our custom stack_effect is needed.
…hashable code objects in TypeSensitiveDict - Remove custom stack_effect() in stack_analysis.py (~46 lines), replaced with xdis.cross_dis.xstack_effect which handles the same null-pushing opcodes (LOAD_GLOBAL, LOAD_ATTR, LOAD_SUPER_ATTR, CALL) and container opcodes (BUILD_MAP, BUILD_CONST_KEY_MAP, BUILD_LIST/SET/TUPLE) correctly. xdis also covers more opcodes (MAKE_FUNCTION, UNPACK_SEQUENCE, etc.) and fixes an off-by-one bug in our UNPACK_EX handling. - Add fallback for xstack_effect returning None (MAKE_FUNCTION with unusual opargs) by using the raw oppush/oppop tables. - Fix TypeSensitiveDict._key_transform to handle any unhashable type (not just the five container types) via try/hash/except fallback. This fixes a crash where code objects with dicts in co_consts (added by the Preprocessor) are unhashable on Python 3.12+. - Fix _key_restore to gracefully handle repr strings that can't be evaluated by ast.literal_eval (e.g. code object reprs).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR seeks to close #124 by implementing a new
Preprocessorpass prior to Masking, which rewrites bytecode into a valid version with statically-reconstructable container types (dict,set,list,tuple) statically folded.Preliminary testing indicates that it works; out of 2711 cases with container type related failures drawn from PyLingual prod MongoDB, 99.4% of them are solved by this fix.
There is also an included set of test cases that confirm that the Preprocessor correctly processes programs.
Model retraining shouldn't be required by this change, as the way the Preprocessor works rewrites bytecode into another valid (though unmarshallable) form before it is passed to the Masking pipeline, which is minimally changed and produces a compatible display format to the existing PyLingual models. I haven't seen any problems about this, the model seems to properly decompile everything with no retraining needed or performance regressions.
Remaining items: Full benchmark still needs to be run to evaluate any potential regressions