inflate_aflplusplus.py

import os
import zipfile
import shutil


def reassemble_file(chunks_dir, output_zip):
    """
    Reassembles chunk files into the original file.
    """
    chunk_files = sorted([
        f for f in os.listdir(chunks_dir)
        if f.startswith("chunk_") and f.endswith(".part")
    ])

    if not chunk_files:
        raise ValueError("No chunk files found in directory.")

    print(f"Reassembling {len(chunk_files)} chunks into '{output_zip}'...")

    with open(output_zip, "wb") as output_file:
        for chunk_file in chunk_files:
            chunk_path = os.path.join(chunks_dir, chunk_file)

            with open(chunk_path, "rb") as f:
                output_file.write(f.read())

            print(f"Added: {chunk_file}")

    print("Reassembly complete.")


def extract_zip(zip_path, extract_to):
    """
    Extracts a ZIP archive.
    """
    print(f"Extracting '{zip_path}' to '{extract_to}'...")

    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(extract_to)

    print("Extraction complete.")


def flatten_subfolder(parent_dir, subfolder_name):
    """
    Moves contents of parent_dir/subfolder_name into parent_dir,
    then removes subfolder_name if empty.
    """
    subfolder_path = os.path.join(parent_dir, subfolder_name)

    if not os.path.isdir(subfolder_path):
        print(f"Subfolder '{subfolder_name}' not found in '{parent_dir}', skipping flatten.")
        return

    print(f"Moving contents of '{subfolder_path}' to '{parent_dir}'...")

    for item in os.listdir(subfolder_path):
        src = os.path.join(subfolder_path, item)
        dst = os.path.join(parent_dir, item)

        if os.path.exists(dst):
            print(f"  Skipping '{item}' — already exists in destination.")
            continue

        shutil.move(src, dst)
        print(f"  Moved: {item}")

    if not os.listdir(subfolder_path):
        os.rmdir(subfolder_path)
        print(f"Removed empty subfolder '{subfolder_path}'.")
    else:
        print(f"Subfolder '{subfolder_path}' is not empty, skipping removal.")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Reassemble chunks into ZIP and optionally extract.")
    parser.add_argument("chunks_dir", help="Directory containing chunk files")
    parser.add_argument("output_zip", nargs="?", default="aflplusplus_pre_built.zip",
                        help="Output ZIP file path (default: aflplusplus_pre_built.zip)")
    parser.add_argument("--extract", help="Optional extraction directory")
    parser.add_argument("--clean", action="store_true",
                        help="Remove the ZIP file after extraction (only applies when --extract is used)")

    args = parser.parse_args()

    # Step 1: Reassemble
    reassemble_file(args.chunks_dir, args.output_zip)

    # Step 2: Optional extraction
    if args.extract:
        extract_to = os.path.expanduser(args.extract)
        extract_zip(args.output_zip, extract_to)

        # Step 3: Flatten subfolder
        flatten_subfolder(extract_to, "AFLplusplus-patched")

        # Step 4: Optional cleanup
        if args.clean:
            os.remove(args.output_zip)
            print(f"Removed '{args.output_zip}'.")