""" @file transformers/rename.py @brief Identifier renaming and string literal encryption transformer. @details Performs scoped variable/function renaming, consistent class/method/attribute remapping, and in-place string encryption producing runtime decryption code. """ import ast from utils.encryption import StringEncryptor from utils.name_gen import NameGenerator import logging # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger("RenameTransformer") class RenameTransformer(ast.NodeTransformer): """ @brief AST transformer for renaming and string encryption. @details Maintains scope stacks, global mappings, class method/attribute maps, and collects key-setup code for runtime decryption. Can emit debug telemetry. """ def __init__(self, name_generator, global_var_renames, class_attr_mapping, primary_key, secondary_key, salt, debug_mode=False): self.name_generator = name_generator self.global_var_renames = global_var_renames self.class_attr_mapping = class_attr_mapping # String encryption self.encryptor = StringEncryptor(primary_key, secondary_key, salt) # We'll collect any code needed for key setup (from encrypt_string calls). self.key_setup_code = [] # Each element is a dict that maps old_name -> new_name for the local scope self.scope_stack = [{}] # For class/attribute rename support self.in_class = False self.current_class_name = None # Add this to track methods in each class self.class_method_mapping = {} # Add a first-pass flag to help with method detection self.first_pass = True # Debug data if debug mode is enabled self.debug_mode = debug_mode if self.debug_mode: self.debug_data = { "variable_mappings": {}, "string_encryption": [], "renamed_nodes": 0, "issues": [] } def log_debug(self, category: str, data: any): """Log debugging information if debug mode is on.""" if self.debug_mode: if category not in self.debug_data: self.debug_data[category] = [] self.debug_data[category].append(data) def visit_Module(self, node: ast.Module) -> ast.Module: """ Enhanced module visitor to properly handle global variables AND top-level function definitions for consistent renaming. Now includes a two-pass strategy for classes to ensure method consistency. """ # First pass: detect methods in classes if self.first_pass: self.first_pass = False # First scan all classes and methods recursively for stmt in ast.walk(node): if isinstance(stmt, ast.ClassDef): self.scan_class_methods(stmt) # Rest of the code remains the same... # First pass: detect top-level assignments, globals, AND function definitions for stmt in node.body: # 1. If a global statement if isinstance(stmt, ast.Global): for name in stmt.names: if name not in self.global_var_renames: self.global_var_renames[name] = self.name_generator.generate_name() # 2. If a top-level assignment elif isinstance(stmt, ast.Assign): for target in stmt.targets: if isinstance(target, ast.Name): if target.id not in self.global_var_renames: self.global_var_renames[target.id] = self.name_generator.generate_name() # 3. If a top-level function definition elif isinstance(stmt, ast.FunctionDef): # Skip special (dunder) methods to avoid messing up e.g. __init__ if not (stmt.name.startswith('__') and stmt.name.endswith('__')): if stmt.name not in self.global_var_renames: self.global_var_renames[stmt.name] = self.name_generator.generate_name() # Create a new body that starts with any needed imports new_body = [ ast.parse("import base64").body[0], ast.parse("import random").body[0] ] # Transform the module body transformed_body = [] for item in node.body: visited = self.visit(item) if isinstance(visited, list): transformed_body.extend(visited) else: transformed_body.append(visited) # If there's any key setup code, parse & insert that if self.key_setup_code: setup_nodes = ast.parse('\n'.join(self.key_setup_code)).body new_body.extend(setup_nodes) new_body.extend(transformed_body) node.body = new_body return node def scan_class_methods(self, node): """ Pre-scan a class to identify and map all its methods before actual renaming. """ class_name = node.name # Create a mapping entry for this class if not exists if class_name not in self.class_method_mapping: self.class_method_mapping[class_name] = {} # Scan all method definitions in the class for item in node.body: if isinstance(item, ast.FunctionDef): method_name = item.name # Skip dunder methods if not (method_name.startswith('__') and method_name.endswith('__')): # Generate a consistent obfuscated name for this method new_name = self.name_generator.generate_name() self.class_method_mapping[class_name][method_name] = new_name def _push_scope(self): self.scope_stack.append({}) def _pop_scope(self): self.scope_stack.pop() def visit_Global(self, node: ast.Global) -> ast.Global: """ Handle global statement declarations by: 1. Adding the variable names to global_var_renames if not already there 2. Adding them to the current scope to mark them as global """ for name in node.names: # If this global name hasn't been seen before, generate a new obfuscated name if name not in self.global_var_renames: self.global_var_renames[name] = self.name_generator.generate_name() # Mark this name as global in the current scope self.scope_stack[-1][name] = self.global_var_renames[name] # Update the global statement with obfuscated names node.names = [self.global_var_renames[name] for name in node.names] return node def visit_ListComp(self, node: ast.ListComp) -> ast.ListComp: self._push_scope() for gen in node.generators: gen = self.visit(gen) node.elt = self.visit(node.elt) self._pop_scope() return node def visit_SetComp(self, node: ast.SetComp) -> ast.SetComp: self._push_scope() for gen in node.generators: gen = self.visit(gen) node.elt = self.visit(node.elt) self._pop_scope() return node def visit_DictComp(self, node: ast.DictComp) -> ast.DictComp: self._push_scope() for gen in node.generators: gen = self.visit(gen) node.key = self.visit(node.key) node.value = self.visit(node.value) self._pop_scope() return node def visit_GeneratorExp(self, node: ast.GeneratorExp) -> ast.GeneratorExp: self._push_scope() for gen in node.generators: gen = self.visit(gen) node.elt = self.visit(node.elt) self._pop_scope() return node def visit_comprehension(self, node: ast.comprehension) -> ast.comprehension: if isinstance(node.target, ast.Name): if node.target.id not in self.global_var_renames: new_name = self.name_generator.generate_name() self.scope_stack[-1][node.target.id] = new_name node.target.id = new_name else: node.target = self.visit(node.target) node.iter = self.visit(node.iter) node.ifs = [self.visit(i) for i in node.ifs] return node def visit_Name(self, node: ast.Name) -> ast.AST: """ Handle variable names and function names in calls. 1. If it's a known global/ top-level function, use global_var_renames. 2. Otherwise, handle locally within scope_stack. """ if node.id in self.global_var_renames: node.id = self.global_var_renames[node.id] return node # Check if this name is marked as global in any scope for scope in self.scope_stack: if node.id in scope and node.id in self.global_var_renames: node.id = self.global_var_renames[node.id] return node # Otherwise, handle local variables if isinstance(node.ctx, ast.Store): if self.in_class and not isinstance(node.ctx, ast.Param): # Class attribute if node.id not in self.scope_stack[-1]: new_name = self.name_generator.generate_name() self.scope_stack[-1][node.id] = new_name if self.current_class_name: self.class_attr_mapping[self.current_class_name][node.id] = new_name node.id = self.scope_stack[-1][node.id] else: # Regular variable assignment if node.id not in self.scope_stack[-1]: self.scope_stack[-1][node.id] = self.name_generator.generate_name() node.id = self.scope_stack[-1][node.id] else: # Load context for scope in reversed(self.scope_stack): if node.id in scope: node.id = scope[node.id] break return node def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: has_bases = len(node.bases) > 0 self.scope_stack[-1]['_has_bases'] = has_bases prev_in_class = self.in_class prev_class_name = self.current_class_name self.in_class = True new_class_name = self.name_generator.generate_name() self.scope_stack[-1][node.name] = new_class_name self.current_class_name = new_class_name self.class_attr_mapping[new_class_name] = {} # Transfer method mappings from the original class name to the renamed one if node.name in self.class_method_mapping: self.class_method_mapping[new_class_name] = self.class_method_mapping[node.name] del self.class_method_mapping[node.name] self.scope_stack.append({}) node.bases = [self.visit(base) for base in node.bases] node.body = [self.visit(b) for b in node.body] self.scope_stack[-2][node.name] = new_class_name node.name = new_class_name self.in_class = prev_in_class self.current_class_name = prev_class_name self.scope_stack.pop() return node def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef: """ Make sure top-level function definitions also get renamed if they're in global_var_renames. """ is_special = node.name.startswith('__') and node.name.endswith('__') has_bases = any(('_has_bases' in scope and scope['_has_bases']) for scope in self.scope_stack) # Save original name for later orig_name = node.name self._push_scope() # Handle arguments (skip renaming for self if method): if self.in_class and len(node.args.args) > 0: self.scope_stack[-1][node.args.args[0].arg] = node.args.args[0].arg for arg in node.args.args[1:]: if arg.arg not in self.global_var_renames: new_arg_name = self.name_generator.generate_name() self.scope_stack[-1][arg.arg] = new_arg_name arg.arg = new_arg_name else: for arg in node.args.args: if arg.arg not in self.global_var_renames: new_arg_name = self.name_generator.generate_name() self.scope_stack[-1][arg.arg] = new_arg_name arg.arg = new_arg_name # Visit body node.body = [self.visit(n) for n in node.body] # -------------------------------- # NEW LOGIC: Actually rename the func if it's top-level # or if we want to rename it anyway (and not dunder). # -------------------------------- if not is_special: # If the function was declared top-level and recognized in global_var_renames, # then rename using that. Otherwise generate a brand new obfuscated name. if node.name in self.global_var_renames: node.name = self.global_var_renames[node.name] elif self.in_class and not has_bases: # For class methods, use the name we saved earlier if self.current_class_name and orig_name in self.class_method_mapping.get(self.current_class_name, {}): node.name = self.class_method_mapping[self.current_class_name][orig_name] else: new_fn_name = self.name_generator.generate_name() self.scope_stack[-2][node.name] = new_fn_name node.name = new_fn_name self._pop_scope() return node def visit_Attribute(self, node: ast.Attribute) -> ast.AST: """ Handle attribute access for renaming: 1. Map method calls on self to the corresponding renamed methods 2. Preserve external/inherited method names 3. Rename self.attribute for class-defined attributes """ # First visit any nested expressions node = self.generic_visit(node) # Check if this is a self.something attribute access if isinstance(node.value, ast.Name) and node.value.id == 'self': # Check all class methods across all classes (more robust) for class_name, methods in self.class_method_mapping.items(): if node.attr in methods: # We found a match in our method mapping node.attr = methods[node.attr] return node # If we're in a class context, apply class-specific logic if self.current_class_name: # Case 1: Is it a call to one of our renamed methods? class_methods = self.class_method_mapping.get(self.current_class_name, {}) if node.attr in class_methods: node.attr = class_methods[node.attr] return node # Case 2: Is it an external method call with Qt-style naming? is_external_method = ( node.attr[0].islower() and any(c.isupper() for c in node.attr) and not node.attr.startswith('__') ) if is_external_method: return node # Case 3: Handle normal class attributes attr_map = self.class_attr_mapping.get(self.current_class_name, {}) if node.attr not in attr_map: attr_map[node.attr] = self.name_generator.generate_name() node.attr = attr_map[node.attr] return node def visit_Subscript(self, node: ast.Subscript) -> ast.AST: node.value = self.visit(node.value) if isinstance(node.slice, ast.AST): self.visit(node.slice) return node def visit_Constant(self, node: ast.Constant) -> ast.AST: """ Encrypt string literals into a multi-step XOR, then base85 decode at runtime. """ if isinstance(node.value, str): encoded, key_setup, modifier = self.encryptor.encrypt_string(node.value) if key_setup not in self.key_setup_code: self.key_setup_code.append(key_setup) decrypt_str = ( f"bytes((" f"k2^k1^m for k1,k2,m in zip(" f"bytes(c^k for c,k in zip(base64.b85decode('{encoded}')," f"_sk*((len(base64.b85decode('{encoded}'))//8)+1)))," f"_pk*((len(base64.b85decode('{encoded}'))//8)+1)," f"bytes.fromhex('{modifier}')*((len(base64.b85decode('{encoded}'))//8)+1)))" f").decode()" ) return ast.parse(decrypt_str).body[0].value return node def visit_Import(self, node: ast.Import) -> ast.AST: return self.generic_visit(node) def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST: return self.generic_visit(node) def visit_Call(self, node: ast.Call) -> ast.Call: """ Handle super() calls specifically to ensure class names are properly updated. """ # First visit all arguments and the function itself node = self.generic_visit(node) # Check if this is a super() call if isinstance(node.func, ast.Name) and node.func.id == 'super': # For super() with no args in Python 3 if not node.args: return node # For super(Class, self) style calls if len(node.args) >= 1 and isinstance(node.args[0], ast.Name): class_name = node.args[0].id # Look for the renamed class in all scopes for scope in reversed(self.scope_stack): if class_name in scope: node.args[0].id = scope[class_name] break return node