obfuscator.py

import ast
import base64
import zlib
import os
import random
import sys
import hashlib
import marshal
import keyword
import builtins

if not hasattr(ast, "unparse"):
    print("Requires Python 3.9+ for ast.unparse", file=sys.stderr)
    sys.exit(1)


RESERVED_NAMES = (
    set(keyword.kwlist)
    | set(dir(builtins))
    | {
        "__name__", "__file__", "__doc__", "__import__",
        "__spec__", "__loader__", "__builtins__", "__package__",
        "self", "cls",
    }
)
CONFUSABLE = "Il1O0o_"


def rand_name(length=None, used=None):
    length = length or random.randint(8, 16)
    used = used if used is not None else set()
    while True:
        n = "".join(random.choice(CONFUSABLE) for _ in range(length))
        if n not in used and not n[0].isdigit():
            used.add(n)
            return n


class DeadCodeInjector(ast.NodeTransformer):
    def __init__(self, n_funcs=4):
        self.n_funcs = n_funcs

    def _make_dead_func(self):
        fname = rand_name()
        v1, v2, v3 = rand_name(), rand_name(), rand_name()
        a = random.randint(1, 99999)
        b = random.randint(1, 99999)
        c = random.randint(1, 99999)
        body = [
            ast.Assign([ast.Name(v1, ast.Store())], ast.Constant(a)),
            ast.Assign([ast.Name(v2, ast.Store())], ast.Constant(b)),
            ast.Assign([ast.Name(v3, ast.Store())], ast.Constant(c)),
            ast.AugAssign(ast.Name(v1, ast.Store()), ast.Add(), ast.Name(v3, ast.Load())),
            ast.While(
                ast.Compare(ast.Name(v2, ast.Load()), [ast.Gt()], [ast.Constant(0)]),
                [ast.AugAssign(ast.Name(v2, ast.Store()), ast.Sub(), ast.Constant(1))],
                [],
            ),
            ast.Return(
                ast.BinOp(ast.Name(v1, ast.Load()), ast.Mult(), ast.Name(v3, ast.Load()))
            ),
        ]
        fn = ast.FunctionDef(
            name=fname,
            args=ast.arguments(
                posonlyargs=[],
                args=[ast.arg(arg=rand_name())],
                kwonlyargs=[],
                kw_defaults=[],
                defaults=[],
            ),
            body=body,
            decorator_list=[],
            returns=None,
        )
        ast.fix_missing_locations(fn)
        return fn

    def visit_Module(self, node):
        self.generic_visit(node)
        new_body = []
        for _ in range(self.n_funcs):
            new_body.append(self._make_dead_func())
        for stmt in node.body:
            if random.random() < 0.3:
                new_body.append(self._make_dead_func())
            new_body.append(stmt)
        node.body = new_body
        return node


class IdentifierRenamer(ast.NodeTransformer):
    def __init__(self):
        self.mapping = {}
        self.used = set()
        self.imported_names = set()

    def _rename(self, name):
        if name in RESERVED_NAMES or name.startswith("__") or name.endswith("__"):
            return name
        if name not in self.mapping:
            self.mapping[name] = rand_name(used=self.used)
        return self.mapping[name]

    def visit_Import(self, node):
        for alias in node.names:
            if alias.asname:
                self.imported_names.add(alias.asname)
            else:
                self.imported_names.add(alias.name.split('.')[0])
        return node

    def visit_ImportFrom(self, node):
        for alias in node.names:
            if alias.asname:
                self.imported_names.add(alias.asname)
            else:
                self.imported_names.add(alias.name)
        return node

    def visit_Name(self, node):
        self.generic_visit(node)
        if node.id in RESERVED_NAMES or node.id.startswith("__") or node.id in self.imported_names:
            return node
        node.id = self._rename(node.id)
        return node

    def visit_arg(self, node):
        self.generic_visit(node)
        if node.arg in RESERVED_NAMES or node.arg.startswith("__") or node.arg in self.imported_names:
            return node
        node.arg = self._rename(node.arg)
        return node

    def _visit_def(self, node):
        self.generic_visit(node)
        if not (node.name in RESERVED_NAMES
                or node.name in self.imported_names
                or node.name.startswith("__")
                or node.name.endswith("__")):
            node.name = self._rename(node.name)
        return node

    visit_FunctionDef = _visit_def
    visit_AsyncFunctionDef = _visit_def

    def visit_ClassDef(self, node):
        self.generic_visit(node)
        if not (node.name in RESERVED_NAMES or node.name in self.imported_names or node.name.startswith("__")):
            node.name = self._rename(node.name)
        return node

    def visit_Attribute(self, node):
        node.value = self.visit(node.value)
        return node

    def visit_keyword(self, node):
        self.generic_visit(node)
        return node


class NumberObfuscator(ast.NodeTransformer):
    def __init__(self, intensity=0.7):
        self.intensity = intensity

    def _obfuscate_int(self, n):
        if n in (0, 1, 2):
            return ast.Constant(value=n)
        if random.random() > self.intensity:
            return ast.Constant(value=n)
        s = random.randint(0, 5)
        if s == 0:
            a = random.randint(-99999, 99999)
            return ast.BinOp(ast.Constant(a), ast.Add(), ast.Constant(n - a))
        elif s == 1:
            a = n + random.randint(1, 99999)
            return ast.BinOp(ast.Constant(a), ast.Sub(), ast.Constant(a - n))
        elif s == 2:
            if n != 0:
                f = random.choice([2, 3, 5, 7, 11, 13])
                if n % f == 0:
                    return ast.BinOp(ast.Constant(n // f), ast.Mult(), ast.Constant(f))
            return ast.BinOp(ast.Constant(n + 5), ast.Sub(), ast.Constant(5))
        elif s == 3:
            mask = random.randint(1, 65535)
            return ast.BinOp(ast.Constant(n ^ mask), ast.BitXor(), ast.Constant(mask))
        elif s == 4:
            a = random.randint(1, 1000)
            c = random.randint(1, 1000)
            d = n - a - c
            return ast.BinOp(
                ast.BinOp(ast.Constant(a), ast.Add(), ast.Constant(c)),
                ast.Add(),
                ast.Constant(d),
            )
        else:
            k = random.randint(1, 99999)
            return ast.BinOp(ast.Constant(n + k), ast.Sub(), ast.Constant(k))

    def visit_Constant(self, node):
        if isinstance(node.value, int) and not isinstance(node.value, bool):
            return self._obfuscate_int(node.value)
        return node


class StringEncryptor(ast.NodeTransformer):
    def __init__(self, key):
        self.key = key
        self.decrypt_fn_name = rand_name(12)

    def _build_decrypt_fn(self):
        key_bytes = list(self.key)
        src = (
            f"def {self.decrypt_fn_name}(_b):\n"
            f"    _k={key_bytes!r}\n"
            f"    return bytes(_c^_k[_i%len(_k)] for _i,_c in enumerate(_b))"
            f".decode('utf-8','replace')\n"
        )
        fn = ast.parse(src).body[0]
        ast.fix_missing_locations(fn)
        return fn

    def visit_Module(self, node):
        self.generic_visit(node)
        node.body = [self._build_decrypt_fn()] + node.body
        return node

    def _encrypt_str(self, s):
        b = s.encode("utf-8")
        return bytes(c ^ self.key[i % len(self.key)] for i, c in enumerate(b))

    def visit_Constant(self, node):
        if isinstance(node.value, str) and node.value:
            enc = self._encrypt_str(node.value)
            return ast.Call(
                func=ast.Name(id=self.decrypt_fn_name, ctx=ast.Load()),
                args=[
                    ast.List(
                        elts=[ast.Constant(value=b) for b in enc],
                        ctx=ast.Load(),
                    )
                ],
                keywords=[],
            )
        return node


class TopLevelFlattener(ast.NodeTransformer):
    def __init__(self):
        self.state_var = rand_name(8)

    def visit_Module(self, node):
        self.generic_visit(node)
        stmts = node.body
        if len(stmts) < 2:
            return node

        states = list(range(len(stmts)))
        random.shuffle(states)
        idx_to_state = {i: states[i] for i in range(len(stmts))}

        branches = []
        for i, stmt in enumerate(stmts):
            next_state = idx_to_state[i + 1] if i + 1 < len(stmts) else -1
            branch_body = [
                stmt,
                ast.Assign(
                    [ast.Name(self.state_var, ast.Store())],
                    ast.Constant(value=next_state),
                ),
            ]
            branches.append(
                ast.If(
                    test=ast.Compare(
                        ast.Name(self.state_var, ast.Load()),
                        [ast.Eq()],
                        [ast.Constant(value=idx_to_state[i])],
                    ),
                    body=branch_body,
                    orelse=[],
                )
            )
        branches.append(
            ast.If(
                test=ast.Compare(
                    ast.Name(self.state_var, ast.Load()),
                    [ast.Lt()],
                    [ast.Constant(value=0)],
                ),
                body=[ast.Break()],
                orelse=[],
            )
        )

        loop = ast.While(test=ast.Constant(True), body=branches, orelse=[])
        init = ast.Assign(
            [ast.Name(self.state_var, ast.Store())],
            ast.Constant(value=idx_to_state[0]),
        )
        node.body = [init, loop]
        ast.fix_missing_locations(node)
        return node


def make_key(seed):
    rng = random.Random(seed)
    return bytes(rng.randint(0, 255) for _ in range(32))


def xor_bytes(data, key):
    return bytes(b ^ key[i % len(key)] for i, b in enumerate(data))


def encode_payload(code_obj, key):
    marshalled = marshal.dumps(code_obj)
    compressed = zlib.compress(marshalled, 9)
    encrypted = xor_bytes(compressed, key)
    encoded = base64.b85encode(encrypted).decode("ascii")
    integrity = hashlib.sha256(encrypted).hexdigest()
    return encoded, integrity


def build_loader(encoded, key, integrity_hash):
    key_list = list(key)
    return f"""
import base64 as _b, zlib as _z, marshal as _m, sys as _s, types as _t, hashlib as _h
_K = {key_list!r}
_C = "{encoded}"
_I = "{integrity_hash}"
if _s.gettrace() is not None and 'bdb' in str(type(_s.gettrace())):
    _s.exit(1)
def _X(_d, _k):
    return bytes(_x ^ _k[_i % len(_k)] for _i, _x in enumerate(_d))
_E = _b.b85decode(_C.encode())
if _h.sha256(_E).hexdigest() != _I:
    _s.exit(1)
_D = _z.decompress(_X(_E, _K))
_C0 = _m.loads(_D)
_g = _t.ModuleType("__main__")
_g.__file__ = __file__ if '__file__' in globals() else ''
exec(_C0, _g.__dict__)
"""


def obfuscate(source, seed=None, flatten=True):
    if seed is None:
        seed = random.randint(0, 2**31)
    random.seed(seed)

    tree = ast.parse(source)

    tree = DeadCodeInjector(n_funcs=random.randint(3, 6)).visit(tree)
    ast.fix_missing_locations(tree)

    tree = IdentifierRenamer().visit(tree)
    ast.fix_missing_locations(tree)

    tree = NumberObfuscator(intensity=0.85).visit(tree)
    ast.fix_missing_locations(tree)

    str_key = make_key(seed ^ 0xDEADBEEF)
    tree = StringEncryptor(str_key).visit(tree)
    ast.fix_missing_locations(tree)

    tree = NumberObfuscator(intensity=0.9).visit(tree)
    ast.fix_missing_locations(tree)

    if flatten:
        tree = TopLevelFlattener().visit(tree)
        ast.fix_missing_locations(tree)

    obfuscated_source = ast.unparse(tree)
    ast.parse(obfuscated_source)

    code_obj = compile(obfuscated_source, "<obfuscated>", "exec")
    payload_key = make_key(seed ^ 0xCAFEBABE)
    encoded, integrity = encode_payload(code_obj, payload_key)

    loader = build_loader(encoded, payload_key, integrity)
    return loader + "\n"


def main():
    if len(sys.argv) < 3:
        print("Usage: python obfuscator.py <input.py> <output.py> [--no-flatten]")
        sys.exit(1)
    inp, outp = sys.argv[1], sys.argv[2]
    flatten = "--no-flatten" not in sys.argv[3:]
    with open(inp, "r") as f:
        src = f.read()
    result = obfuscate(src, flatten=flatten)
    with open(outp, "w") as f:
        f.write(result)
    print(f"[+] Obfuscated: {inp} -> {outp}")
    print(f"[+] Output size: {len(result)} bytes")


if __name__ == "__main__":
    main()