"""Run 5: blind recovery of real files encrypted by the live sample (2026-10-12).

Inputs a responder would have, and nothing else:
  * the .akira files (akira/), pulled from the VM and hash-checked against the VM's listing;
  * each file's LastWriteTime and a QPC/system-time pair measured on the victim after the attack
    (post_run5.txt; same boot, no reboot);
  * known plaintext derived from the file TYPE or from the file's own unencrypted part:
      PNG  -> the fixed 16-byte signature + IHDR chunk header shared by every PNG
      PDF  -> "%PDF-1." (7 bytes, version-agnostic)
      DOCX -> the first local file header rebuilt from the ZIP central directory, which Spot mode
              leaves unencrypted at the end of the file.
No seeds, keys, breakpoints or debugger captures are used. The pre-encryption listing (inv_pre.txt)
is read only at the very end, to compare SHA-256 of the decrypted output with the original.
"""
import hashlib, os, statistics, struct, subprocess, sys, time

HERE = os.path.dirname(os.path.abspath(__file__))
CASE = os.path.dirname(os.path.dirname(HERE))
TOOL = os.path.join(CASE, "06-decryptor", "akira_decrypt.py")
sys.path.insert(0, os.path.dirname(TOOL))
import akira_decrypt as ad

PNG_HDR = bytes.fromhex("89504e470d0a1a0a0000000d49484452")

def zip_first_local_header(enc, size):
    """Rebuild the local header at offset 0 from the central directory in the plaintext tail."""
    body = enc[:size]
    e = body.rfind(b"PK\x05\x06")
    cd_size, cd_off = struct.unpack("<II", body[e + 12:e + 20])
    p = cd_off
    while p < cd_off + cd_size:
        (sig, vmade, vneed, flags, meth, mtime, mdate, crc, csz, usz, nlen, xlen, clen,
         disk, iattr, eattr, lho) = struct.unpack("<IHHHHHHIIIHHHHHII", body[p:p + 46])
        name = body[p + 46:p + 46 + nlen]
        if lho == 0:
            # local header: signature, version needed, flags, method, time, date, crc, sizes, name length;
            # the extra-field length is not repeated reliably, so stop before it
            return struct.pack("<IHHHHHIIIH", 0x04034B50, vneed, flags, meth, mtime, mdate, crc, csz, usz, nlen), name
        p += 46 + nlen + xlen + clen
    raise ValueError("no central-directory entry for offset 0")

def rows(path):
    return [l.rstrip("\r\n").split("|") for l in open(path, encoding="ascii")]

post = rows(os.path.join(HERE, "post_run5.txt"))
pairs = [list(map(int, r[1:])) for r in post if r[0] == "PAIR"]
offs = [ft - (q0 + q1) // 2 for q0, ft, q1 in pairs[1:]]          # first call is slowed by JIT
offset = int(statistics.median(offs))
mtimes = {os.path.basename(r[1]): int(r[5]) for r in post if r[0] == "FILE" and r[1].endswith(".akira")}
vm_sha = {os.path.basename(r[1]): r[3] for r in post if r[0] == "FILE"}

print("# Run 5 blind recovery - %s UTC" % time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()))
print("# decryptor SHA-256 %s" % hashlib.sha256(open(TOOL, "rb").read()).hexdigest())
print("# QPC offset from %d post-run pairs: %d (spread %d ticks), QPC freq %s" % (
    len(offs), offset, max(offs) - min(offs), [r[1] for r in post if r[0] == "QPCFREQ"][0]))

order = ["Screenshot 2026-09-29 125745.png.akira", "1.pdf.akira",
         "QC-STC-T-45-26_Sample_IR_Deliverables_v1.0.docx.akira", "Screenshot 2026-10-05 102304.png.akira"]
only = sys.argv[1:] or order
outdir = os.path.join(HERE, "decrypted"); os.makedirs(outdir, exist_ok=True)
results = {}
for name in order:
    if name not in only:
        continue
    path = os.path.join(HERE, "akira", name)
    enc = open(path, "rb").read()
    size = len(enc) - 512
    ext = os.path.splitext(name[:-6])[1].lower()
    if ext == ".png":
        known, why = PNG_HDR, "PNG signature + IHDR header (identical in every PNG)"
    elif ext == ".pdf":
        known, why = b"%PDF-1.", "'%PDF-1.' (every PDF, version digit left unknown)"
    else:
        known, nm = zip_first_local_header(enc, size)
        why = "ZIP local header rebuilt from the unencrypted central directory (first entry %r)" % nm.decode()
    print("\n## %s" % name)
    print("  .akira SHA-256 %s (VM listing %s: %s)" % (hashlib.sha256(enc).hexdigest(), vm_sha[name][:16],
          hashlib.sha256(enc).hexdigest() == vm_sha[name]))
    print("  mtime %d -> key2 window from mtime - offset, default -8/+1 ms" % mtimes[name])
    print("  known plaintext: %d bytes = %s: %s" % (len(known), why, known.hex()))
    out = os.path.join(outdir, name[:-6])
    t0 = time.time()
    r = subprocess.run([sys.executable, "-I", TOOL, "recover", path, "--qpc-offset", str(offset),
                        "--mtime", str(mtimes[name]), "--known-hex", known.hex(), "--out", out, "--force"],
                       capture_output=True, text=True)
    for line in r.stdout.splitlines():
        if "derived" not in line:
            print("  | " + line)
    if r.returncode or not os.path.exists(out):
        print("  RECOVER FAILED (exit %d) %s" % (r.returncode, r.stderr[-500:]))
        results[name] = None
        continue
    results[name] = hashlib.sha256(open(out, "rb").read()).hexdigest()
    print("  elapsed %.0f s; decrypted SHA-256 %s" % (time.time() - t0, results[name]))

# only now: compare with the pre-encryption listing
pre = {os.path.basename(r[1]): (int(r[2]), r[3]) for r in rows(os.path.join(HERE, "inv_pre.txt")) if r[0] == "FILE"}
print("\n## Comparison with the pre-encryption SHA-256 (inv_pre.txt, taken on the VM before the run)")
fails = 0
for name, got in results.items():
    size, want = pre[name[:-6]]
    ok = got == want
    fails += not ok
    print("  %-52s original %s  decrypted %s  %s" % (name[:-6], want[:16], (got or "-")[:16], "MATCH" if ok else "NO MATCH"))
print("FAILURES: %d of %d" % (fails, len(results)))
sys.exit(1 if fails else 0)