PythonUnpackLLM.py

import argparse
import os
import sys
from file_extractors import PycHandler
import time
from file_extractors.decompiler import PyInstallerUnpacker
import logging
import re
from tqdm import tqdm
from pathlib import Path
import marshal
import dis
import types
import io

# setup logging to print to console
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

BASE_DIRECTORY = os.path.abspath(os.getcwd())

def disassemble_pyc(pyc_path: str) -> str:
    """
    Return recursive disassembly of a .pyc file as text.
    """
    with open(pyc_path, "rb") as f:
        f.read(16)  # skip pyc header (works for 3.7+)
        code = marshal.load(f)

    output = io.StringIO()

    def walk(code_obj: types.CodeType, indent=0):
        prefix = " " * indent
        print(f"{prefix}==> Function: {code_obj.co_name}", file=output)
        dis.dis(code_obj, file=output)
        print("\n", file=output)

        for const in code_obj.co_consts:
            if isinstance(const, types.CodeType):
                walk(const, indent + 2)

    walk(code)
    return output.getvalue()

def detect_exe_type(exe_path: str) -> str:
    """
    Classify the executable packaging/obfuscation method.
    """
    data = Path(exe_path).read_bytes()

    # PyInstaller
    if b"MEI\x0c\x0b\x0a\x0b\x0e" in data or b"PYZ-00.pyz" in data:
        if b"pyi-archive_viewer" in data:
            return "PyInstaller"

        if b"pyi_rth_" in data and b"cryptography" in data:
            return "PyInstaller + AES encrypted archive"

        return "PyInstaller"

    # Nuitka
    if b"NUITKA" in data or b"__nuitka" in data:
        return "Nuitka"

    # PyArmor
    if b"pyarmor_runtime" in data or b"pytransform" in data:
        return "PyArmor Obfuscated"

    # PyInstaller but encrypted loader
    if b"PYZ" in data and b"AES" in data:
        return "PyInstaller Encrypted"

    # Native C/C++
    if b"MSVCRT" in data or b"GCC:" in data:
        return "Native C/C++ Binary"

    return "Unknown or Packed by Other Tool"

def validate_path(path):
    # Example: Allow only alphanumeric characters, underscores, hyphens, and slashes
    if re.match(r'^[\w\-/\\.]+$', path): 
        return path
    else:
        raise ValueError("Invalid path")


def is_safe_path(base_dir, path, follow_symlinks=True):
    # Resolve the absolute path
    if follow_symlinks:
        resolved_path = os.path.realpath(path)
    else:
        resolved_path = os.path.abspath(path)
    # Check if the resolved path starts with the base directory
    return resolved_path.startswith(os.path.realpath(base_dir))

def main():
    start = time.time()

    llm_options = ['auto']
    if os.path.exists('llm'):
        llm_options += os.listdir('llm')

    parser = argparse.ArgumentParser(description='PythonUnpackLLM')
    parser.add_argument('--path', type=str, required=True)
    parser.add_argument(
    "--asm",
    action="store_true",
    help="Output raw Python bytecode disassembly instead of LLM reconstruction"
)
    parser.add_argument("--output", default=None)
    parser.add_argument('--type', choices=['exe', 'pyc', 'folder', 'py_bytecode'], default='pyc')
    parser.add_argument('-u', '--unpack', action='store_true',
                    help='Unpack EXE before processing')

    args = parser.parse_args()

    if not os.path.exists(args.path):
        raise ValueError(f"Path '{args.path}' does not exist.")

    try:
        validate_path(args.path)
        if not is_safe_path(BASE_DIRECTORY, args.path):
            raise ValueError(f"Unsafe path '{args.path}'.")
    except ValueError as e:
        logger.error(e)
        sys.exit(1)

    if args.type == 'folder':

        input_dir = Path(args.path).resolve()
        output_dir = Path(args.output).resolve() if args.output else input_dir.parent / (input_dir.name + "_decompiled")
        output_dir.mkdir(parents=True, exist_ok=True)

        pyc_files = list(input_dir.rglob("*.pyc"))
        print(f"📦 Found {len(pyc_files)} .pyc files\n")

        for pyc_file in pyc_files:
            rel_path = pyc_file.relative_to(input_dir)

            if args.asm:
                out_path = output_dir / rel_path.with_suffix(".asm.txt")
                res = disassemble_pyc(str(pyc_file))
            else:
                out_path = output_dir / rel_path.with_suffix(".py")
                res = PycHandler(str(pyc_file)).handle()

            out_path.parent.mkdir(parents=True, exist_ok=True)

            try:
                with open(out_path, "w", encoding="utf-8") as f:
                    f.write(res)
            except Exception as e:
                tqdm.write(f"❌ Failed {pyc_file.name}: {e}")

    elif args.type == 'pyc':

        if args.asm:
            res = disassemble_pyc(args.path)
            ext = ".asm"
        else:
            res = PycHandler(args.path).handle()
            ext = ".py"

        if args.output:
            if os.path.isdir(args.output):
                base_name = os.path.splitext(os.path.basename(args.path))[0] + ext
                output_path = os.path.join(args.output, base_name)
            else:
                output_path = args.output if args.output.endswith(ext) else args.output + ext
        else:
            input_path = os.path.abspath(args.path)
            output_path = os.path.splitext(input_path)[0] + ext

        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        with open(output_path, "w", encoding="utf-8") as f:
            f.write(res)

        print(f"[+] Saved to: {output_path}")

    elif args.unpack:
        if not args.output:
            raise ValueError("--unpack requires --output directory")

        exe_type = detect_exe_type(str(args.path))
        if exe_type in ["PyArmor Obfuscated", "Native C/C++ Binary", "Nuitka", "Unknown or Packed by Other Tool"]:
            raise ValueError(f"Format '{exe_type}' detected and not supported.")

        logger.info(f"[DETECT] EXE Type: {exe_type}. Attempt to extract..")
        unpacker = PyInstallerUnpacker(args.path, args.output)
        extracted_dir = unpacker.run()

        args.path = extracted_dir
        args.type = "folder"


    else:
        print("Type not implemented yet.")
        sys.exit(1)

    time_in_seconds = time.time() - start
    import datetime
    print(f"\n⏱ Time taken: {str(datetime.timedelta(seconds=time_in_seconds))}")

if __name__ == '__main__':
    main()