#!/usr/bin/env python3
"""Recover payloads from the two reviewed J@IL-GPT loaders as analysis files.

Python 3.9+; standard library only. This tool accepts only the exact advertised
EXE or the historical MSI's already-extracted Binary.p PE, identified by SHA-256.
It does not extract MSI/ZIP containers, execute sample code, contact a network,
or emulate Windows. Recovered .sample files are still malware.

Usage:
    python3 extract-jailgpt.py INPUT_FILE NEW_OUTPUT_DIRECTORY

The output directory must not exist. Both recovered hashes are verified before
any output is written. Files are created exclusively with mode 0600; the new
directory uses mode 0700 (permission enforcement depends on the operating system).
"""

import argparse
import base64
import hashlib
import json
import os
from pathlib import Path
import re
import struct
import sys


MAX_INPUT_BYTES = 7_202_239
INNER_BYTES = 2_662_400
DLL_BYTES = 1_140_224
DLL_CIPHER_OFFSET = 0xEA374
DLL_CIPHER_BYTES = 1_520_300
STRING_KEY = b"xnasff3wcedj"

SAMPLES = {
    "056f5bdc2cb87c8d1290311cc9735e9104d45c6148021f5bd33857f5c14f523a": {
        "lineage": "exe",
        "bytes": 7_202_239,
        "footer": (0x6D512115, 519_590, 6_682_637),
        "key_offset": 0x6210D,
        "inner_sha256": "e395fd44a750ccc4da5740d83645fda7b2cdc0bb9de3fd6e8cbfba7ac90c3feb",
        "dll_sha256": "28b7159e0f0702d595711318c34531e553e8e4910ba215912a9d4cfcf6cc0431",
    },
    "485f69d2d9515ad7d64382d8e986964815a6528f0c61ff62dce307690900c5d0": {
        "lineage": "msi",
        "bytes": 7_135_320,
        "footer": (0x28EEC550, 482_878, 6_652_430),
        "key_offset": 0x3AFC8,
        "inner_sha256": "f8dfd1e44ad3acd3e7fd0cff6cb0f02f0c2e88ebbc363990c22517f78a0c0011",
        "dll_sha256": "1f9ee6cb61f77176208b5cf3cc59f739b543520536de4d1ca141de3a4b03de4e",
    },
}


def require(condition, message):
    if not condition:
        raise ValueError(message)


def sha256(data):
    return hashlib.sha256(data).hexdigest()


def bounded_slice(data, offset, length):
    require(0 <= offset <= len(data), "Invalid data offset")
    require(0 <= length <= len(data) - offset, "Data range exceeds file")
    return data[offset : offset + length]


def validate_pe(data):
    """Read a few bounded PE fields as data; never load the image."""
    require(len(data) >= 64 and data[:2] == b"MZ", "Missing DOS header")
    pe_offset = struct.unpack_from("<I", data, 0x3C)[0]
    header = bounded_slice(data, pe_offset, 24)
    require(header[:4] == b"PE\0\0", "Missing PE signature")
    machine, section_count = struct.unpack_from("<HH", header, 4)
    optional_size = struct.unpack_from("<H", header, 20)[0]
    require(machine == 0x8664, "Expected AMD64 PE")
    require(1 <= section_count <= 96, "Invalid section count")
    require(optional_size >= 112, "Truncated optional header")
    optional = bounded_slice(data, pe_offset + 24, optional_size)
    require(struct.unpack_from("<H", optional)[0] == 0x20B, "Expected PE32+")
    sections = bounded_slice(data, pe_offset + 24 + optional_size, section_count * 40)
    for index in range(section_count):
        raw_size, raw_offset = struct.unpack_from("<II", sections, index * 40 + 16)
        bounded_slice(data, raw_offset, raw_size)


def rotate_right(value, shift):
    return ((value >> shift) | (value << (8 - shift))) & 255


def rotate_left(value, shift):
    return ((value << shift) | (value >> (8 - shift))) & 255


def decode_exe(ciphertext, key):
    """Arithmetic transcribed from the EXE decoder at VA 0x140007c70."""
    output = bytearray()
    previous = 100
    for index, current in enumerate(ciphertext):
        mix = (((index * 0x515A) >> 8) & 255) ^ key[index & 31]
        descending = (-index - 0x10) & 255
        value = ((mix + 0x5A) ^ ((rotate_right(current, 3) ^ 0xAF) + 0x4B)) & 255
        intermediate = (
            ((index & 255) + 0x26)
            ^ (descending + 0x97)
            ^ (mix - 1)
            ^ rotate_left(value, 2)
            ^ 0x9D
        )
        value = (
            (mix + 0x91)
            ^ (
                (
                    (mix + 0x51)
                    ^ (
                        (previous ^ 0x47)
                        + (previous ^ 4)
                        + ((intermediate - 5) ^ 0xF6)
                        - 0x2E
                    )
                )
                + 0x74
            )
        ) & 255
        value = ((previous ^ 0xEE) + (previous ^ 0x38) + rotate_right(value, 3) + 2) & 255
        value = (
            ((index & 255) - 0x1D)
            ^ ((index & 255) + 0x95)
            ^ ((previous ^ 0xFE) + (rotate_right(value, 3) ^ descending) + 0x4A)
            ^ 0xAC
        ) & 255
        output.append(rotate_left(value, 2) ^ mix ^ 0x98)
        previous = current
    return bytes(output)


def decode_msi_loader(ciphertext, key):
    """Arithmetic transcribed from the MSI's loader at VA 0x140001284."""
    output = bytearray()
    previous = 0xC0
    for index, current in enumerate(ciphertext):
        value = (current - (previous ^ 0xF3)) & 255
        value ^= ((index - 0x75) & 255) ^ 0x75
        value = (value + (previous ^ 0xAF) + 0xBF) & 255
        value ^= (index + 0x5B) & 255
        value = (value + (previous ^ 0xB4) + 0x29) & 255
        mix = (((index * 0xB05) >> 8) & 255) ^ key[index & 31]
        value ^= ((mix + 0x35) & 255) ^ 0x48
        value = (value + (previous ^ 0x4B) + 0x83) & 255
        value ^= (0x8F - mix) & 255
        value = (value + (previous ^ 0xAB)) & 255
        output.append(rotate_left(value, 1) ^ mix)
        previous = current
    return bytes(output)


def extract(data, sample):
    validate_pe(data)
    require(len(data) == sample["bytes"], "Unexpected input length")
    footer = struct.unpack("<III", bounded_slice(data, len(data) - 12, 12))
    require(footer == sample["footer"], "Footer does not match the preserved sample")
    _, offset, length = footer
    require(offset + length + 12 == len(data), "Footer boundary mismatch")
    blob = bounded_slice(data, offset, length)
    require(blob[0] == 0, "Unexpected container version")
    expected = struct.unpack_from("<I", blob, 1)[0]
    require(expected == INNER_BYTES + 4, "Unexpected encoded-stream length")

    encoded = bytearray()
    for line in blob[5:].splitlines():
        longest = max(re.findall(rb"[0-9a-f]+", line), key=len, default=b"")
        if len(longest) >= 32 and len(longest) % 2 == 0:
            fragment = bytes.fromhex(longest.decode("ascii"))
            encoded.extend(fragment[: expected - len(encoded)])
        if len(encoded) == expected:
            break
    require(len(encoded) == expected, "Incomplete encoded stream")
    key = bounded_slice(data, sample["key_offset"], 32)
    payload_length = struct.unpack("<I", bytes(a ^ b for a, b in zip(encoded[:4], key[:4])))[0]
    require(payload_length == INNER_BYTES, "Unexpected decoded-length prefix")
    decoder = decode_exe if sample["lineage"] == "exe" else decode_msi_loader
    inner = decoder(encoded[4:], key)
    require(len(inner) == INNER_BYTES and sha256(inner) == sample["inner_sha256"], "Inner SHA-256 mismatch")
    validate_pe(inner)

    ciphertext = bounded_slice(inner, DLL_CIPHER_OFFSET, DLL_CIPHER_BYTES)
    base64_text = bytes(value ^ STRING_KEY[index % len(STRING_KEY)] for index, value in enumerate(ciphertext))
    dll = base64.b64decode(base64_text, validate=True)
    require(len(dll) == DLL_BYTES and sha256(dll) == sample["dll_sha256"], "DLL SHA-256 mismatch")
    validate_pe(dll)
    return inner, dll


def write_exclusive(path, data):
    descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(descriptor, "wb") as output:
        output.write(data)
    path.chmod(0o600)


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("input_file", type=Path, help="Verified outer EXE or already-extracted historical Binary.p")
    parser.add_argument("output_directory", type=Path, help="New directory; an existing directory is refused")
    args = parser.parse_args()
    try:
        require(args.input_file.is_file(), "Input must be a regular file")
        with args.input_file.open("rb") as source:
            data = source.read(MAX_INPUT_BYTES + 1)
        require(len(data) <= MAX_INPUT_BYTES, "Input exceeds the largest accepted sample")
        digest = sha256(data)
        require(digest in SAMPLES, "Unknown input SHA-256; no sample parsing or extraction performed")
        sample = SAMPLES[digest]
        require(not os.path.lexists(args.output_directory), "Output directory already exists; refusing overwrite")
        inner, dll = extract(data, sample)
        lineage = sample["lineage"]
        inner_name = f"inner-from-{lineage}.sample"
        dll_name = f"browser-dll-from-{lineage}.sample"
        manifest = {
            "method": "Offline data extraction only; no sample execution or network access",
            "lineage": lineage,
            "input_sha256": digest,
            "input_bytes": len(data),
            "artifacts": [
                {"file": inner_name, "bytes": len(inner), "sha256": sha256(inner), "parent_sha256": digest},
                {"file": dll_name, "bytes": len(dll), "sha256": sha256(dll), "parent_sha256": sha256(inner)},
            ],
        }
        manifest_bytes = (json.dumps(manifest, indent=2) + "\n").encode("utf-8")
        args.output_directory.mkdir(mode=0o700)
        write_exclusive(args.output_directory / inner_name, inner)
        write_exclusive(args.output_directory / dll_name, dll)
        write_exclusive(args.output_directory / "extraction.json", manifest_bytes)
        print(manifest_bytes.decode("utf-8"), end="")
        return 0
    except (OSError, ValueError, struct.error) as error:
        print(f"Extraction refused or failed: {error}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())