File size: 3,137 Bytes
6e9bd63 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | import ast
import hashlib
from typing import Dict, Set
class WatermarkStripper(ast.NodeTransformer):
"""
AST-based watermark / fingerprint stripper + light name obfuscator.
- Renames user-defined functions, classes, and arguments
- Removes leading docstrings from functions and classes
- Preserves builtins, dunders, and imports
"""
def __init__(self, prefix: str = "_w"):
self.prefix = prefix
self.name_map: Dict[str, str] = {}
self.counter = 0
self.protected: Set[str] = {
"print", "len", "range", "list", "dict", "set", "tuple",
"str", "int", "float", "bool", "None", "True", "False",
"self", "cls", "super", "object", "type", "Exception",
"__init__", "__name__", "__main__", "__file__", "__doc__",
"__all__", "__dict__", "__class__", "__module__",
}
def _get_obfuscated_name(self, original: str) -> str:
if original in self.protected or original.startswith("__"):
return original
if original not in self.name_map:
digest = hashlib.md5(original.encode()).hexdigest()[:6]
self.name_map[original] = f"{self.prefix}{self.counter}_{digest}"
self.counter += 1
return self.name_map[original]
def visit_FunctionDef(self, node: ast.FunctionDef):
node.name = self._get_obfuscated_name(node.name)
if (node.body and isinstance(node.body[0], ast.Expr)
and isinstance(node.body[0].value, ast.Constant)
and isinstance(node.body[0].value.value, str)):
node.body.pop(0)
return self.generic_visit(node)
def visit_ClassDef(self, node: ast.ClassDef):
node.name = self._get_obfuscated_name(node.name)
if (node.body and isinstance(node.body[0], ast.Expr)
and isinstance(node.body[0].value, ast.Constant)
and isinstance(node.body[0].value.value, str)):
node.body.pop(0)
return self.generic_visit(node)
def visit_arg(self, node: ast.arg):
node.arg = self._get_obfuscated_name(node.arg)
return self.generic_visit(node)
def visit_Name(self, node: ast.Name):
if isinstance(node.ctx, (ast.Load, ast.Store, ast.Del)):
if node.id in self.name_map:
node.id = self.name_map[node.id]
return node
def visit_Attribute(self, node: ast.Attribute):
self.generic_visit(node)
return node
def sanitize_code(source_code: str) -> str:
tree = ast.parse(source_code)
stripper = WatermarkStripper()
transformed = stripper.visit(tree)
ast.fix_missing_locations(transformed)
return ast.unparse(transformed)
if __name__ == "__main__":
sample = '''
def calculate_sum(a, b):
"""This is a watermark docstring."""
return a + b
class Calculator:
"""Class-level watermark."""
def multiply(self, x, y):
return x * y
result = calculate_sum(3, 4)
print(result)
'''
print("=== ORIGINAL ===")
print(sample)
print("\n=== STRIPPED / OBFUSCATED ===")
print(sanitize_code(sample))
|