mirror of
https://github.com/bikini/patchwork
synced 2026-06-27 08:08:41 +00:00
Expand compatibility stress coverage
This commit is contained in:
ashton
+98
-79
@@ -1,79 +1,98 @@
|
||||
from __future__ import annotations
|
||||
import ast
|
||||
from pathlib import Path
|
||||
from .lazy import encrypt_user_functions
|
||||
from .loader import build_loader
|
||||
from .packer import pack
|
||||
from .transforms import IdentifierRenamer, NumberObfuscator, OpaquePredicateInjector, StringEncryptor, build_decrypt_helper, collect_skip_names
|
||||
from .transforms.junk import JunkBranchInjector
|
||||
from .transforms.mba import MBATransformer
|
||||
from .transforms.opaque import build_seed_stmt
|
||||
from .util import make_rng
|
||||
_SEED_NAME = '_pw_seed_value'
|
||||
_DEC_S_NAME = '_pw_dec_str'
|
||||
_DEC_B_NAME = '_pw_dec_bytes'
|
||||
_RESOLVE_NAME = '__pw_resolve_lazy__'
|
||||
|
||||
class Obfuscator:
|
||||
|
||||
def __init__(self, *, seed: int | None=None, rename: bool=True, encrypt_strings: bool=True, obfuscate_numbers: bool=True, opaque_predicates: bool=True, mba: bool=True, junk_branches: bool=True, lazy_funcs: bool=True, anti_debug: bool=True, layers: int=3, stage2_layers: int=3, keep: set[str] | None=None):
|
||||
self.rng, self.seed = make_rng(seed)
|
||||
self.rename = rename
|
||||
self.encrypt_strings = encrypt_strings
|
||||
self.obfuscate_numbers = obfuscate_numbers
|
||||
self.opaque_predicates = opaque_predicates
|
||||
self.mba = mba
|
||||
self.junk_branches = junk_branches
|
||||
self.lazy_funcs = lazy_funcs
|
||||
self.anti_debug = anti_debug
|
||||
self.layers = max(1, int(layers))
|
||||
self.stage2_layers = max(1, int(stage2_layers))
|
||||
self.keep_extra = set(keep or ())
|
||||
|
||||
def obfuscate(self, source: str) -> str:
|
||||
tree = ast.parse(source)
|
||||
skip = collect_skip_names(tree)
|
||||
seed_used = False
|
||||
if self.opaque_predicates:
|
||||
opi = OpaquePredicateInjector(self.rng, _SEED_NAME)
|
||||
tree = opi.visit(tree)
|
||||
seed_used = seed_used or opi.injected
|
||||
if self.junk_branches:
|
||||
jbi = JunkBranchInjector(self.rng, _SEED_NAME)
|
||||
tree = jbi.visit(tree)
|
||||
seed_used = seed_used or jbi.injected
|
||||
if seed_used:
|
||||
tree.body.insert(0, build_seed_stmt(_SEED_NAME))
|
||||
if self.mba:
|
||||
tree = MBATransformer(self.rng).visit(tree)
|
||||
if self.obfuscate_numbers:
|
||||
tree = NumberObfuscator(self.rng).visit(tree)
|
||||
if self.encrypt_strings:
|
||||
enc = StringEncryptor(self.rng, _DEC_S_NAME, _DEC_B_NAME)
|
||||
tree = enc.visit(tree)
|
||||
tree.body[:0] = build_decrypt_helper(_DEC_S_NAME, _DEC_B_NAME)
|
||||
if self.rename:
|
||||
renamer = IdentifierRenamer(self.rng, keep=skip | self.keep_extra | {_RESOLVE_NAME})
|
||||
tree = renamer.visit(tree)
|
||||
ast.fix_missing_locations(tree)
|
||||
user_code = compile(tree, '<patchwork>', 'exec')
|
||||
if self.lazy_funcs:
|
||||
user_code, lazy_blobs = encrypt_user_functions(user_code, self.rng, _RESOLVE_NAME)
|
||||
else:
|
||||
lazy_blobs = []
|
||||
user_payload, user_keys = pack(user_code, self.rng, layers=self.layers)
|
||||
return build_loader(user_payload, user_keys, lazy_blobs, resolver_name=_RESOLVE_NAME, rng=self.rng, anti_debug=self.anti_debug, stage2_layers=self.stage2_layers)
|
||||
|
||||
def obfuscate(source: str, **kwargs) -> str:
|
||||
return Obfuscator(**kwargs).obfuscate(source)
|
||||
|
||||
def obfuscate_file(input_path: str | Path, output_path: str | Path | None=None, **kwargs) -> Path:
|
||||
inp = Path(input_path)
|
||||
src = inp.read_text(encoding='utf-8')
|
||||
out = obfuscate(src, **kwargs)
|
||||
if output_path is None:
|
||||
outp = inp.with_name(inp.stem + '_obf.py')
|
||||
else:
|
||||
outp = Path(output_path)
|
||||
outp.write_text(out, encoding='utf-8')
|
||||
return outp
|
||||
from __future__ import annotations
|
||||
import ast
|
||||
from pathlib import Path
|
||||
from .lazy import encrypt_user_functions
|
||||
from .loader import build_loader
|
||||
from .packer import pack
|
||||
from .transforms import IdentifierRenamer, NumberObfuscator, OpaquePredicateInjector, StringEncryptor, build_decrypt_helper, collect_skip_names
|
||||
from .transforms.junk import JunkBranchInjector
|
||||
from .transforms.mba import MBATransformer
|
||||
from .transforms.opaque import build_seed_stmt
|
||||
from .util import make_rng
|
||||
_SEED_NAME = '_pw_seed_value'
|
||||
_DEC_S_NAME = '_pw_dec_str'
|
||||
_DEC_B_NAME = '_pw_dec_bytes'
|
||||
_RESOLVE_NAME = '__pw_resolve_lazy__'
|
||||
|
||||
def _insert_runtime_stmts(tree: ast.Module, stmts: list[ast.stmt]) -> None:
|
||||
if not stmts:
|
||||
return
|
||||
insert_at = 0
|
||||
if (
|
||||
tree.body
|
||||
and isinstance(tree.body[0], ast.Expr)
|
||||
and isinstance(tree.body[0].value, ast.Constant)
|
||||
and isinstance(tree.body[0].value.value, str)
|
||||
):
|
||||
insert_at = 1
|
||||
while (
|
||||
insert_at < len(tree.body)
|
||||
and isinstance(tree.body[insert_at], ast.ImportFrom)
|
||||
and tree.body[insert_at].module == '__future__'
|
||||
):
|
||||
insert_at += 1
|
||||
tree.body[insert_at:insert_at] = stmts
|
||||
|
||||
class Obfuscator:
|
||||
|
||||
def __init__(self, *, seed: int | None=None, rename: bool=True, encrypt_strings: bool=True, obfuscate_numbers: bool=True, opaque_predicates: bool=True, mba: bool=True, junk_branches: bool=True, lazy_funcs: bool=True, anti_debug: bool=True, layers: int=3, stage2_layers: int=3, keep: set[str] | None=None):
|
||||
self.rng, self.seed = make_rng(seed)
|
||||
self.rename = rename
|
||||
self.encrypt_strings = encrypt_strings
|
||||
self.obfuscate_numbers = obfuscate_numbers
|
||||
self.opaque_predicates = opaque_predicates
|
||||
self.mba = mba
|
||||
self.junk_branches = junk_branches
|
||||
self.lazy_funcs = lazy_funcs
|
||||
self.anti_debug = anti_debug
|
||||
self.layers = max(1, int(layers))
|
||||
self.stage2_layers = max(1, int(stage2_layers))
|
||||
self.keep_extra = set(keep or ())
|
||||
|
||||
def obfuscate(self, source: str) -> str:
|
||||
tree = ast.parse(source)
|
||||
skip = collect_skip_names(tree)
|
||||
seed_used = False
|
||||
if self.opaque_predicates:
|
||||
opi = OpaquePredicateInjector(self.rng, _SEED_NAME)
|
||||
tree = opi.visit(tree)
|
||||
seed_used = seed_used or opi.injected
|
||||
if self.junk_branches:
|
||||
jbi = JunkBranchInjector(self.rng, _SEED_NAME)
|
||||
tree = jbi.visit(tree)
|
||||
seed_used = seed_used or jbi.injected
|
||||
if seed_used:
|
||||
_insert_runtime_stmts(tree, [build_seed_stmt(_SEED_NAME)])
|
||||
if self.mba:
|
||||
tree = MBATransformer(self.rng).visit(tree)
|
||||
if self.obfuscate_numbers:
|
||||
tree = NumberObfuscator(self.rng).visit(tree)
|
||||
if self.encrypt_strings:
|
||||
enc = StringEncryptor(self.rng, _DEC_S_NAME, _DEC_B_NAME)
|
||||
tree = enc.visit(tree)
|
||||
_insert_runtime_stmts(tree, build_decrypt_helper(_DEC_S_NAME, _DEC_B_NAME))
|
||||
if self.rename:
|
||||
renamer = IdentifierRenamer(self.rng, keep=skip | self.keep_extra | {_RESOLVE_NAME})
|
||||
tree = renamer.visit(tree)
|
||||
ast.fix_missing_locations(tree)
|
||||
user_code = compile(tree, '<patchwork>', 'exec')
|
||||
if self.lazy_funcs:
|
||||
user_code, lazy_blobs = encrypt_user_functions(user_code, self.rng, _RESOLVE_NAME)
|
||||
else:
|
||||
lazy_blobs = []
|
||||
user_payload, user_keys = pack(user_code, self.rng, layers=self.layers)
|
||||
return build_loader(user_payload, user_keys, lazy_blobs, resolver_name=_RESOLVE_NAME, rng=self.rng, anti_debug=self.anti_debug, stage2_layers=self.stage2_layers)
|
||||
|
||||
def obfuscate(source: str, **kwargs) -> str:
|
||||
return Obfuscator(**kwargs).obfuscate(source)
|
||||
|
||||
def obfuscate_file(input_path: str | Path, output_path: str | Path | None=None, **kwargs) -> Path:
|
||||
inp = Path(input_path)
|
||||
src = inp.read_text(encoding='utf-8')
|
||||
out = obfuscate(src, **kwargs)
|
||||
if output_path is None:
|
||||
outp = inp.with_name(inp.stem + '_obf.py')
|
||||
else:
|
||||
outp = Path(output_path)
|
||||
outp.write_text(out, encoding='utf-8')
|
||||
return outp
|
||||
|
||||
+218
-167
@@ -1,167 +1,218 @@
|
||||
from __future__ import annotations
|
||||
import ast
|
||||
import builtins
|
||||
import keyword
|
||||
import random
|
||||
from ..util import gen_name
|
||||
_DUNDERS = {'__name__', '__main__', '__file__', '__doc__', '__builtins__', '__package__', '__path__', '__loader__', '__spec__', '__cached__', '__all__', '__init__', '__new__', '__del__', '__init_subclass__', '__subclasshook__', '__set_name__', '__class_getitem__', '__enter__', '__exit__', '__aenter__', '__aexit__', '__await__', '__aiter__', '__anext__', '__call__', '__repr__', '__str__', '__bytes__', '__format__', '__hash__', '__bool__', '__len__', '__length_hint__', '__sizeof__', '__dir__', '__eq__', '__ne__', '__lt__', '__le__', '__gt__', '__ge__', '__add__', '__sub__', '__mul__', '__truediv__', '__floordiv__', '__mod__', '__divmod__', '__pow__', '__lshift__', '__rshift__', '__and__', '__or__', '__xor__', '__radd__', '__rsub__', '__rmul__', '__rtruediv__', '__rfloordiv__', '__rmod__', '__rdivmod__', '__rpow__', '__rlshift__', '__rrshift__', '__rand__', '__ror__', '__rxor__', '__iadd__', '__isub__', '__imul__', '__itruediv__', '__ifloordiv__', '__imod__', '__ipow__', '__ilshift__', '__irshift__', '__iand__', '__ior__', '__ixor__', '__matmul__', '__rmatmul__', '__imatmul__', '__neg__', '__pos__', '__abs__', '__invert__', '__int__', '__float__', '__complex__', '__round__', '__index__', '__trunc__', '__floor__', '__ceil__', '__getattr__', '__setattr__', '__delattr__', '__getattribute__', '__getitem__', '__setitem__', '__delitem__', '__missing__', '__iter__', '__next__', '__reversed__', '__contains__', '__copy__', '__deepcopy__', '__reduce__', '__reduce_ex__', '__getstate__', '__setstate__', '__getnewargs__', '__getnewargs_ex__', '__class__', '__instancecheck__', '__subclasscheck__', '__slots__', '__dict__', '__weakref__', '__module__', '__qualname__', '__defaults__', '__kwdefaults__', '__annotations__', '__wrapped__', '__signature__', '__match_args__', '__post_init__', '__fspath__', '__index__', '__buffer__', '__release_buffer__'}
|
||||
RESERVED: set[str] = set(dir(builtins)) | set(keyword.kwlist) | _DUNDERS | {'self', 'cls', 'metaclass'}
|
||||
|
||||
def collect_skip_names(tree: ast.AST) -> set[str]:
|
||||
skip: set[str] = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Assign):
|
||||
for target in node.targets:
|
||||
if isinstance(target, ast.Name) and target.id == '__all__':
|
||||
if isinstance(node.value, (ast.List, ast.Tuple, ast.Set)):
|
||||
for elt in node.value.elts:
|
||||
if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
|
||||
skip.add(elt.value)
|
||||
elif isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and (node.func.id in ('getattr', 'setattr', 'hasattr', 'delattr')):
|
||||
if len(node.args) >= 2:
|
||||
arg = node.args[1]
|
||||
if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
|
||||
skip.add(arg.value)
|
||||
elif isinstance(node, ast.ClassDef):
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
skip.add(stmt.name)
|
||||
elif isinstance(stmt, ast.Assign):
|
||||
for tgt in stmt.targets:
|
||||
if isinstance(tgt, ast.Name):
|
||||
skip.add(tgt.id)
|
||||
elif isinstance(tgt, (ast.Tuple, ast.List)):
|
||||
for elt in tgt.elts:
|
||||
if isinstance(elt, ast.Name):
|
||||
skip.add(elt.id)
|
||||
elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
|
||||
skip.add(stmt.target.id)
|
||||
elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
|
||||
args = node.args
|
||||
for a in args.args + args.posonlyargs + args.kwonlyargs:
|
||||
skip.add(a.arg)
|
||||
if args.vararg:
|
||||
skip.add(args.vararg.arg)
|
||||
if args.kwarg:
|
||||
skip.add(args.kwarg.arg)
|
||||
return skip
|
||||
|
||||
class IdentifierRenamer(ast.NodeTransformer):
|
||||
|
||||
def __init__(self, rng: random.Random, keep: set[str] | None=None):
|
||||
self.rng = rng
|
||||
self.mapping: dict[str, str] = {}
|
||||
self.reserved = RESERVED | (keep or set())
|
||||
self._used_outputs: set[str] = set()
|
||||
|
||||
def _new(self) -> str:
|
||||
n = gen_name(self.rng, used=self._used_outputs)
|
||||
self._used_outputs.add(n)
|
||||
return n
|
||||
|
||||
def _rename(self, name: str) -> str:
|
||||
if not name or name in self.reserved:
|
||||
return name
|
||||
if name.startswith('__') and name.endswith('__') and (len(name) >= 4):
|
||||
return name
|
||||
if name not in self.mapping:
|
||||
self.mapping[name] = self._new()
|
||||
return self.mapping[name]
|
||||
|
||||
def visit_Name(self, node: ast.Name) -> ast.AST:
|
||||
node.id = self._rename(node.id)
|
||||
return node
|
||||
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
|
||||
node.name = self._rename(node.name)
|
||||
node.decorator_list = [self.visit(d) for d in node.decorator_list]
|
||||
if node.returns is not None:
|
||||
node.returns = self.visit(node.returns)
|
||||
node.args = self.visit(node.args)
|
||||
node.body = [self.visit(stmt) for stmt in node.body]
|
||||
return node
|
||||
visit_AsyncFunctionDef = visit_FunctionDef
|
||||
|
||||
def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST:
|
||||
node.name = self._rename(node.name)
|
||||
node.bases = [self.visit(b) for b in node.bases]
|
||||
for kw in node.keywords:
|
||||
kw.value = self.visit(kw.value)
|
||||
node.decorator_list = [self.visit(d) for d in node.decorator_list]
|
||||
node.body = [self.visit(stmt) for stmt in node.body]
|
||||
return node
|
||||
|
||||
def visit_Lambda(self, node: ast.Lambda) -> ast.AST:
|
||||
node.args = self.visit(node.args)
|
||||
node.body = self.visit(node.body)
|
||||
return node
|
||||
|
||||
def visit_arg(self, node: ast.arg) -> ast.AST:
|
||||
node.arg = self._rename(node.arg)
|
||||
if node.annotation is not None:
|
||||
node.annotation = self.visit(node.annotation)
|
||||
return node
|
||||
|
||||
def visit_arguments(self, node: ast.arguments) -> ast.AST:
|
||||
node.args = [self.visit(a) for a in node.args]
|
||||
node.posonlyargs = [self.visit(a) for a in node.posonlyargs]
|
||||
node.kwonlyargs = [self.visit(a) for a in node.kwonlyargs]
|
||||
if node.vararg is not None:
|
||||
node.vararg = self.visit(node.vararg)
|
||||
if node.kwarg is not None:
|
||||
node.kwarg = self.visit(node.kwarg)
|
||||
node.defaults = [self.visit(d) for d in node.defaults]
|
||||
node.kw_defaults = [self.visit(d) if d is not None else None for d in node.kw_defaults]
|
||||
return node
|
||||
|
||||
def visit_Attribute(self, node: ast.Attribute) -> ast.AST:
|
||||
node.value = self.visit(node.value)
|
||||
return node
|
||||
|
||||
def visit_keyword(self, node: ast.keyword) -> ast.AST:
|
||||
if node.value is not None:
|
||||
node.value = self.visit(node.value)
|
||||
return node
|
||||
|
||||
def visit_Global(self, node: ast.Global) -> ast.AST:
|
||||
node.names = [self._rename(n) for n in node.names]
|
||||
return node
|
||||
|
||||
def visit_Nonlocal(self, node: ast.Nonlocal) -> ast.AST:
|
||||
node.names = [self._rename(n) for n in node.names]
|
||||
return node
|
||||
|
||||
def visit_ExceptHandler(self, node: ast.ExceptHandler) -> ast.AST:
|
||||
if node.type is not None:
|
||||
node.type = self.visit(node.type)
|
||||
if node.name:
|
||||
node.name = self._rename(node.name)
|
||||
node.body = [self.visit(s) for s in node.body]
|
||||
return node
|
||||
|
||||
def visit_Import(self, node: ast.Import) -> ast.AST:
|
||||
for alias in node.names:
|
||||
if '.' in alias.name:
|
||||
if alias.asname:
|
||||
new = self._rename(alias.asname)
|
||||
alias.asname = new
|
||||
continue
|
||||
if alias.asname:
|
||||
alias.asname = self._rename(alias.asname)
|
||||
else:
|
||||
if alias.name not in self.mapping:
|
||||
self.mapping[alias.name] = self._new()
|
||||
alias.asname = self.mapping[alias.name]
|
||||
return node
|
||||
|
||||
def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST:
|
||||
for alias in node.names:
|
||||
if alias.name == '*':
|
||||
continue
|
||||
if alias.asname:
|
||||
alias.asname = self._rename(alias.asname)
|
||||
else:
|
||||
if alias.name not in self.mapping:
|
||||
self.mapping[alias.name] = self._new()
|
||||
alias.asname = self.mapping[alias.name]
|
||||
return node
|
||||
from __future__ import annotations
|
||||
import ast
|
||||
import builtins
|
||||
import keyword
|
||||
import random
|
||||
from ..util import gen_name
|
||||
_DUNDERS = {'__name__', '__main__', '__file__', '__doc__', '__builtins__', '__package__', '__path__', '__loader__', '__spec__', '__cached__', '__all__', '__init__', '__new__', '__del__', '__init_subclass__', '__subclasshook__', '__set_name__', '__class_getitem__', '__enter__', '__exit__', '__aenter__', '__aexit__', '__await__', '__aiter__', '__anext__', '__call__', '__repr__', '__str__', '__bytes__', '__format__', '__hash__', '__bool__', '__len__', '__length_hint__', '__sizeof__', '__dir__', '__eq__', '__ne__', '__lt__', '__le__', '__gt__', '__ge__', '__add__', '__sub__', '__mul__', '__truediv__', '__floordiv__', '__mod__', '__divmod__', '__pow__', '__lshift__', '__rshift__', '__and__', '__or__', '__xor__', '__radd__', '__rsub__', '__rmul__', '__rtruediv__', '__rfloordiv__', '__rmod__', '__rdivmod__', '__rpow__', '__rlshift__', '__rrshift__', '__rand__', '__ror__', '__rxor__', '__iadd__', '__isub__', '__imul__', '__itruediv__', '__ifloordiv__', '__imod__', '__ipow__', '__ilshift__', '__irshift__', '__iand__', '__ior__', '__ixor__', '__matmul__', '__rmatmul__', '__imatmul__', '__neg__', '__pos__', '__abs__', '__invert__', '__int__', '__float__', '__complex__', '__round__', '__index__', '__trunc__', '__floor__', '__ceil__', '__getattr__', '__setattr__', '__delattr__', '__getattribute__', '__getitem__', '__setitem__', '__delitem__', '__missing__', '__iter__', '__next__', '__reversed__', '__contains__', '__copy__', '__deepcopy__', '__reduce__', '__reduce_ex__', '__getstate__', '__setstate__', '__getnewargs__', '__getnewargs_ex__', '__class__', '__instancecheck__', '__subclasscheck__', '__slots__', '__dict__', '__weakref__', '__module__', '__qualname__', '__defaults__', '__kwdefaults__', '__annotations__', '__wrapped__', '__signature__', '__match_args__', '__post_init__', '__fspath__', '__index__', '__buffer__', '__release_buffer__'}
|
||||
RESERVED: set[str] = set(dir(builtins)) | set(keyword.kwlist) | _DUNDERS | {'self', 'cls', 'metaclass'}
|
||||
|
||||
def _has_future_annotations(tree: ast.AST) -> bool:
|
||||
if not isinstance(tree, ast.Module):
|
||||
return False
|
||||
for stmt in tree.body:
|
||||
if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant) and isinstance(stmt.value.value, str):
|
||||
continue
|
||||
if isinstance(stmt, ast.ImportFrom) and stmt.module == '__future__':
|
||||
if any(alias.name == 'annotations' for alias in stmt.names):
|
||||
return True
|
||||
continue
|
||||
break
|
||||
return False
|
||||
|
||||
def _names_in_annotation(node: ast.AST | None) -> set[str]:
|
||||
if node is None:
|
||||
return set()
|
||||
return {child.id for child in ast.walk(node) if isinstance(child, ast.Name)}
|
||||
|
||||
def collect_skip_names(tree: ast.AST) -> set[str]:
|
||||
skip: set[str] = set()
|
||||
preserve_annotation_names = _has_future_annotations(tree)
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Assign):
|
||||
for target in node.targets:
|
||||
if isinstance(target, ast.Name) and target.id == '__all__':
|
||||
if isinstance(node.value, (ast.List, ast.Tuple, ast.Set)):
|
||||
for elt in node.value.elts:
|
||||
if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
|
||||
skip.add(elt.value)
|
||||
elif isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and (node.func.id in ('getattr', 'setattr', 'hasattr', 'delattr')):
|
||||
if len(node.args) >= 2:
|
||||
arg = node.args[1]
|
||||
if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
|
||||
skip.add(arg.value)
|
||||
elif isinstance(node, ast.ClassDef):
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
skip.add(stmt.name)
|
||||
elif isinstance(stmt, ast.Assign):
|
||||
for tgt in stmt.targets:
|
||||
if isinstance(tgt, ast.Name):
|
||||
skip.add(tgt.id)
|
||||
elif isinstance(tgt, (ast.Tuple, ast.List)):
|
||||
for elt in tgt.elts:
|
||||
if isinstance(elt, ast.Name):
|
||||
skip.add(elt.id)
|
||||
elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
|
||||
skip.add(stmt.target.id)
|
||||
elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
|
||||
args = node.args
|
||||
if preserve_annotation_names:
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
skip.update(_names_in_annotation(node.returns))
|
||||
for a in args.args + args.posonlyargs + args.kwonlyargs:
|
||||
skip.update(_names_in_annotation(a.annotation))
|
||||
if args.vararg:
|
||||
skip.update(_names_in_annotation(args.vararg.annotation))
|
||||
if args.kwarg:
|
||||
skip.update(_names_in_annotation(args.kwarg.annotation))
|
||||
for a in args.args + args.posonlyargs + args.kwonlyargs:
|
||||
skip.add(a.arg)
|
||||
if args.vararg:
|
||||
skip.add(args.vararg.arg)
|
||||
if args.kwarg:
|
||||
skip.add(args.kwarg.arg)
|
||||
elif preserve_annotation_names and isinstance(node, ast.AnnAssign):
|
||||
skip.update(_names_in_annotation(node.annotation))
|
||||
return skip
|
||||
|
||||
class IdentifierRenamer(ast.NodeTransformer):
|
||||
|
||||
def __init__(self, rng: random.Random, keep: set[str] | None=None):
|
||||
self.rng = rng
|
||||
self.mapping: dict[str, str] = {}
|
||||
self.reserved = RESERVED | (keep or set())
|
||||
self._used_outputs: set[str] = set()
|
||||
|
||||
def _new(self) -> str:
|
||||
n = gen_name(self.rng, used=self._used_outputs)
|
||||
self._used_outputs.add(n)
|
||||
return n
|
||||
|
||||
def _rename(self, name: str) -> str:
|
||||
if not name or name in self.reserved:
|
||||
return name
|
||||
if name.startswith('__') and name.endswith('__') and (len(name) >= 4):
|
||||
return name
|
||||
if name not in self.mapping:
|
||||
self.mapping[name] = self._new()
|
||||
return self.mapping[name]
|
||||
|
||||
def visit_Name(self, node: ast.Name) -> ast.AST:
|
||||
node.id = self._rename(node.id)
|
||||
return node
|
||||
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
|
||||
node.name = self._rename(node.name)
|
||||
node.decorator_list = [self.visit(d) for d in node.decorator_list]
|
||||
if node.returns is not None:
|
||||
node.returns = self.visit(node.returns)
|
||||
node.args = self.visit(node.args)
|
||||
node.body = [self.visit(stmt) for stmt in node.body]
|
||||
return node
|
||||
visit_AsyncFunctionDef = visit_FunctionDef
|
||||
|
||||
def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST:
|
||||
node.name = self._rename(node.name)
|
||||
node.bases = [self.visit(b) for b in node.bases]
|
||||
for kw in node.keywords:
|
||||
kw.value = self.visit(kw.value)
|
||||
node.decorator_list = [self.visit(d) for d in node.decorator_list]
|
||||
node.body = [self.visit(stmt) for stmt in node.body]
|
||||
return node
|
||||
|
||||
def visit_Lambda(self, node: ast.Lambda) -> ast.AST:
|
||||
node.args = self.visit(node.args)
|
||||
node.body = self.visit(node.body)
|
||||
return node
|
||||
|
||||
def visit_arg(self, node: ast.arg) -> ast.AST:
|
||||
node.arg = self._rename(node.arg)
|
||||
if node.annotation is not None:
|
||||
node.annotation = self.visit(node.annotation)
|
||||
return node
|
||||
|
||||
def visit_arguments(self, node: ast.arguments) -> ast.AST:
|
||||
node.args = [self.visit(a) for a in node.args]
|
||||
node.posonlyargs = [self.visit(a) for a in node.posonlyargs]
|
||||
node.kwonlyargs = [self.visit(a) for a in node.kwonlyargs]
|
||||
if node.vararg is not None:
|
||||
node.vararg = self.visit(node.vararg)
|
||||
if node.kwarg is not None:
|
||||
node.kwarg = self.visit(node.kwarg)
|
||||
node.defaults = [self.visit(d) for d in node.defaults]
|
||||
node.kw_defaults = [self.visit(d) if d is not None else None for d in node.kw_defaults]
|
||||
return node
|
||||
|
||||
def visit_Attribute(self, node: ast.Attribute) -> ast.AST:
|
||||
node.value = self.visit(node.value)
|
||||
return node
|
||||
|
||||
def visit_keyword(self, node: ast.keyword) -> ast.AST:
|
||||
if node.value is not None:
|
||||
node.value = self.visit(node.value)
|
||||
return node
|
||||
|
||||
def visit_Global(self, node: ast.Global) -> ast.AST:
|
||||
node.names = [self._rename(n) for n in node.names]
|
||||
return node
|
||||
|
||||
def visit_Nonlocal(self, node: ast.Nonlocal) -> ast.AST:
|
||||
node.names = [self._rename(n) for n in node.names]
|
||||
return node
|
||||
|
||||
def visit_ExceptHandler(self, node: ast.ExceptHandler) -> ast.AST:
|
||||
if node.type is not None:
|
||||
node.type = self.visit(node.type)
|
||||
if node.name:
|
||||
node.name = self._rename(node.name)
|
||||
node.body = [self.visit(s) for s in node.body]
|
||||
return node
|
||||
|
||||
def visit_Import(self, node: ast.Import) -> ast.AST:
|
||||
for alias in node.names:
|
||||
if '.' in alias.name:
|
||||
if alias.asname:
|
||||
new = self._rename(alias.asname)
|
||||
alias.asname = new
|
||||
continue
|
||||
if alias.asname:
|
||||
alias.asname = self._rename(alias.asname)
|
||||
else:
|
||||
if alias.name not in self.mapping:
|
||||
self.mapping[alias.name] = self._new()
|
||||
alias.asname = self.mapping[alias.name]
|
||||
return node
|
||||
|
||||
def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST:
|
||||
if node.module == '__future__':
|
||||
return node
|
||||
for alias in node.names:
|
||||
if alias.name == '*':
|
||||
continue
|
||||
if alias.asname:
|
||||
alias.asname = self._rename(alias.asname)
|
||||
else:
|
||||
if alias.name not in self.mapping:
|
||||
self.mapping[alias.name] = self._new()
|
||||
alias.asname = self.mapping[alias.name]
|
||||
return node
|
||||
|
||||
def visit_MatchAs(self, node: ast.MatchAs) -> ast.AST:
|
||||
if node.pattern is not None:
|
||||
node.pattern = self.visit(node.pattern)
|
||||
if node.name is not None:
|
||||
node.name = self._rename(node.name)
|
||||
return node
|
||||
|
||||
def visit_MatchStar(self, node: ast.MatchStar) -> ast.AST:
|
||||
if node.name is not None:
|
||||
node.name = self._rename(node.name)
|
||||
return node
|
||||
|
||||
def visit_MatchMapping(self, node: ast.MatchMapping) -> ast.AST:
|
||||
node.keys = [self.visit(key) for key in node.keys]
|
||||
node.patterns = [self.visit(pattern) for pattern in node.patterns]
|
||||
if node.rest is not None:
|
||||
node.rest = self._rename(node.rest)
|
||||
return node
|
||||
|
||||
Reference in New Issue
Block a user