dis

import argparse
import marshal
import os
import types
import dis
import json
import sys
def load_pyc(file_path):
    with open(file_path, "rb") as f:
        header_size = 16  # Python 3.13
        f.read(header_size)
        return marshal.load(f)
def extract_functions(code_obj, parent=""):
    functions = {}
    for const in code_obj.co_consts:
        if isinstance(const, types.CodeType):
            name = const.co_name
            full_name = f"{parent}.{name}" if parent else name
            if full_name not in functions:
                functions[full_name] = const
            functions.update(extract_functions(const, parent=full_name))
    return functions
def save_functions(functions, output_dir):
    for name, code in functions.items():
        safe_name = name.replace(".", "_")
        file_path = os.path.join(output_dir, f"{safe_name}.dis")
        with open(file_path, "w", encoding="utf-8") as f:
            try:
                f.write(dis.code_info(code) + "\n\n")
                dis.dis(code, file=f)
            except Exception as e:
                f.write(f"Disassembly error: {e}")
def extract_imports_vars_classes(code_obj):
    imports = set()
    variables = []
    classes = []
    for instr in dis.get_instructions(code_obj):
        if instr.opname in ("IMPORT_NAME", "IMPORT_FROM"):
            imports.add(instr.argval)
        elif instr.opname == "STORE_NAME":
            variables.append({"name": instr.argval, "line": instr.starts_line or -1, "value": "<?>"} )
        elif instr.opname == "LOAD_BUILD_CLASS":
            idx = instr.offset
            next_instrs = list(dis.get_instructions(code_obj))[idx:idx+5]
            for ni in next_instrs:
                if ni.opname == "LOAD_CONST" and isinstance(ni.argval, str):
                    classes.append({"name": ni.argval, "line": ni.starts_line or -1})
                    break
    return list(imports), variables, classes
def extract_function_defs_with_consts(code_obj):
    functions = []
    for const in code_obj.co_consts:
        if isinstance(const, types.CodeType):
            if const.co_name != "<module>":
                argcount = const.co_argcount
                argnames = const.co_varnames[:argcount]
                docstring = ""
                if const.co_consts and isinstance(const.co_consts[0], str):
                    docstring = const.co_consts[0]
                literals = [c for c in const.co_consts if isinstance(c, (str, int, float, bool))]
                functions.append({
                    "name": const.co_name,
                    "args": list(argnames),
                    "line": const.co_firstlineno,
                    "doc": docstring,
                    "literals": list(set(literals) - {docstring}),
                })
            functions += extract_function_defs_with_consts(const)
    return functions
def detect_if_main(code_obj):
    for instr in dis.get_instructions(code_obj):
        if instr.opname == "LOAD_NAME" and instr.argval == "__name__":
            return True
    return False
def save_json(imports, variables, classes, output_dir):
    with open(os.path.join(output_dir, "imports.json"), "w", encoding="utf-8") as f:
        json.dump(sorted(imports), f, indent=2)
    with open(os.path.join(output_dir, "vars.json"), "w", encoding="utf-8") as f:
        json.dump(variables, f, indent=2)
    with open(os.path.join(output_dir, "classes.json"), "w", encoding="utf-8") as f:
        json.dump(classes, f, indent=2)
def generate_stub(imports, variables, classes, functions, has_main, output_path, include_lines=True):
    with open(output_path, "w", encoding="utf-8") as f:
        for imp in sorted(imports):
            f.write(f"import {imp}\n")
        f.write("\n")
        for var in variables:
            line = f"# line {var['line']} " if include_lines and var['line'] > 0 else ""
            f.write(f"{line}{var['name']} = {var['value']}\n")
        f.write("\n")
        for cls in classes:
            line = f"# line {cls['line']} " if include_lines and cls['line'] > 0 else ""
            f.write(f"{line}class {cls['name']}:\n    \"\"\"TODO: Add class documentation\"\"\"\n    pass\n\n")
        for func in functions:
            args = ", ".join(func["args"])
            line = f"# line {func['line']} " if include_lines and func['line'] > 0 else ""
            f.write(f"{line}def {func['name']}({args}):\n")
            if func["doc"]:
                f.write(f"    \"\"\"{func['doc']}\"\"\"\n")
            else:
                f.write(f"    \"\"\"TODO: Add function documentation\"\"\"\n")
            f.write(f"    pass\n\n")
        if has_main:
            f.write('if __name__ == "__main__":\n    pass\n')
def generate_type_stub(functions, classes, output_path):
    with open(output_path, "w", encoding="utf-8") as f:
        for cls in classes:
            f.write(f"class {cls['name']}:\n    ...\n\n")
        for func in functions:
            args = ", ".join(func["args"])
            f.write(f"def {func['name']}({args}) -> None: ...\n")
def parse_args():
    parser = argparse.ArgumentParser(description="Disassemble a .pyc file and extract components.")
    parser.add_argument("pyc_path", help="Path to the .pyc file")
    parser.add_argument("--stub", action="store_true", help="Generate a Python stub file")
    parser.add_argument("--json", action="store_true", help="Export imports/vars/classes as JSON")
    parser.add_argument("--filter", help="Only include functions containing this string")
    return parser.parse_args()
def main():
    args = parse_args()
    # Use same folder as the .pyc file
    output_dir = os.path.dirname(os.path.abspath(args.pyc_path))
    os.makedirs(output_dir, exist_ok=True)
    code = load_pyc(args.pyc_path)
    functions = extract_functions(code)
    if args.filter:
        functions = {name: fn for name, fn in functions.items() if args.filter in name}
    save_functions(functions, output_dir=output_dir)
    imports, variables, classes = extract_imports_vars_classes(code)
    function_defs = extract_function_defs_with_consts(code)
    has_main = detect_if_main(code)
    if args.json:
        save_json(imports, variables, classes, output_dir=output_dir)
    if args.stub:
        generate_stub(imports, variables, classes, function_defs, has_main,
                      output_path=os.path.join(output_dir, "decompiled_stub.py"))
        generate_type_stub(function_defs, classes,
                           output_path=os.path.join(output_dir, "decompiled_stub.pyi"))
if __name__ == "__main__":
    main()