"""Run 6: fully independent blind recovery of real files encrypted by the live sample.

The rules below were fixed BEFORE the detonation (see 04-dynamic/07-run6-independent-blind.md) and are
not changed after it. The script works on any folder of .akira files, so it can be dry-run on run 5.

Inputs a responder would have, and nothing else:
  * the .akira files (<run-dir>/akira/), pulled from the VM; each is hash-checked against the VM listing;
  * each file's LastWriteTime and the QPC/system-time pairs measured on the victim after the attack
    (post listing; same boot, no reboot);
  * known plaintext at offset 0 derived from the file TYPE, or from the file's own unencrypted part
    (KNOWN_RULES below). Files whose type has no rule are reported as out of scope, not as failures.
Per file, at most two passes:
  pass 1  key2 window = mtime - offset, -8/+1 ms; gap 3,800..8,000 ticks          (the recover defaults)
  pass 2  only if pass 1 misses: --anchor iv2, iv2 in mtime -2.5/+0.1 ms; gap 8,001..30,000 ticks
No seeds, keys, breakpoints or debugger captures are used. The pre-encryption listing is read only after
every file has been processed, to compare SHA-256 of the decrypted output with the original.

usage: run6_blind_recover.py [--run-dir DIR] [--pre pre_run6.txt] [--post post_run6.txt] [--only NAME ...]
"""
import argparse, glob, hashlib, json, os, re, 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

PASSES = [("pass 1 (defaults)", ["--mtime-window", "8:1", "--gap", "3800:8000"]),
          ("pass 2 (--anchor iv2)", ["--anchor", "iv2", "--mtime-window", "2.5:0.1", "--gap", "8001:30000"])]
ZIP_EXTS = {".docx", ".xlsx", ".pptx", ".docm", ".xlsm", ".pptm", ".zip", ".odt", ".ods", ".odp", ".jar", ".epub", ".vsdx"}
OLE_EXTS = {".doc", ".xls", ".ppt", ".msg", ".vsd", ".pub"}


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")
    if e < 0:
        raise ValueError("no end-of-central-directory record in the unencrypted part")
    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])
        if sig != 0x02014B50:
            raise ValueError("central directory entry at %d is not readable (encrypted?)" % p)
        name = body[p + 46:p + 46 + nlen]
        if lho == 0:
            # signature .. 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 known_plaintext(name, enc, size):
    """KNOWN_RULES: (bytes, reason) or (None, reason) when the type has no rule."""
    ext = os.path.splitext(name)[1].lower()
    if ext == ".png":
        return bytes.fromhex("89504e470d0a1a0a0000000d49484452"), "PNG signature + IHDR header (every PNG)"
    if ext == ".pdf":
        return b"%PDF-1.", "'%PDF-1.' (every PDF, version digit unknown)"
    if ext in ZIP_EXTS:
        try:
            hdr, nm = zip_first_local_header(enc, size)
            return hdr, "ZIP local header rebuilt from the unencrypted central directory (first entry %r)" % nm.decode("cp437")
        except Exception as e:
            return None, "ZIP type, but the local header could not be rebuilt: %s" % e
    if ext in OLE_EXTS:
        return bytes.fromhex("d0cf11e0a1b11ae1") + bytes(16), "OLE2 signature + zero CLSID (24 bytes)"
    if ext in (".jpg", ".jpeg"):
        return bytes.fromhex("ffd8ffe000104a46494600"), "JPEG with a JFIF APP0 header (11 bytes; wrong for EXIF-first JPEGs)"
    if ext in (".sqlite", ".sqlite3", ".db"):
        return b"SQLite format 3\x00", "SQLite header (16 bytes)"
    if ext == ".rtf":
        return b"{\\rtf1\\ansi", "RTF header '{\\rtf1\\ansi' (11 bytes)"
    return None, "no known-plaintext rule for '%s'" % ext


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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--run-dir", default=HERE)
    ap.add_argument("--pre", default="pre_run6.txt")
    ap.add_argument("--post", default="post_run6.txt")
    ap.add_argument("--only", nargs="*")
    ap.add_argument("--engine", default="auto")
    a = ap.parse_args()
    rd = a.run_dir

    post = rows(os.path.join(rd, a.post))
    pairs = [list(map(int, r[1:4])) 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))
    vm = {os.path.basename(r[1]): r for r in post if r[0] == "FILE"}

    print("# Run 6 blind recovery - %s UTC" % time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()))
    print("# driver SHA-256    %s" % hashlib.sha256(open(os.path.abspath(__file__), "rb").read()).hexdigest())
    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]))

    files = sorted(glob.glob(os.path.join(rd, "akira", "*.akira")))
    if a.only:
        files = [f for f in files if os.path.basename(f) in a.only]
    outdir = os.path.join(rd, "decrypted"); os.makedirs(outdir, exist_ok=True)
    res = {}
    for path in files:
        name = os.path.basename(path)
        enc = open(path, "rb").read()
        size = len(enc) - 512
        r = {"size": size, "passes": []}
        res[name] = r
        print("\n## %s (%d bytes original)" % (name, size))
        sha = hashlib.sha256(enc).hexdigest()
        r["vm_hash_ok"] = name in vm and vm[name][3] == sha
        print("  .akira SHA-256 %s, matches VM listing: %s" % (sha, r["vm_hash_ok"]))
        mode, pct = ad.infer_mode(name[:-6], size, 50)
        plan = list(ad.chunk_plan(mode, pct, size))
        r["mode"] = ad.MODE_NAMES[mode]
        r["chacha"] = any(c[2] == "chacha" for c in plan)
        print("  mode %s, %d chunks, ChaCha chunks: %s" % (r["mode"], len(plan), r["chacha"]))
        known, why = known_plaintext(name[:-6], enc, size)
        if known is None or name not in vm:
            r["status"] = "OUT OF SCOPE"
            r["why"] = why if known is None else "not in the VM listing"
            print("  OUT OF SCOPE: %s" % r["why"])
            continue
        mtime = int(vm[name][5])
        print("  mtime %d; known plaintext %d bytes = %s: %s" % (mtime, len(known), why, known.hex()))
        out = os.path.join(outdir, name[:-6])
        if os.path.exists(out):
            os.remove(out)
        for label, extra in PASSES:
            t0 = time.time()
            p = subprocess.run([sys.executable, "-I", TOOL, "recover", path, "--qpc-offset", str(offset),
                                "--mtime", str(mtime), "--known-hex", known.hex(), "--engine", a.engine,
                                "--out", out, "--force"] + extra, capture_output=True, text=True)
            el = time.time() - t0
            for line in p.stdout.splitlines():
                if "derived" not in line and "pairs tested" not in line:
                    print("  | " + line)
            m = re.search(r"FOUND key2 seed (\d+), iv2 seed (\d+) \(gap (\d+)", p.stdout)
            r["passes"].append({"pass": label, "elapsed_s": round(el, 1), "found": bool(m), "exit": p.returncode})
            print("  %s: %s in %.0f s" % (label, "FOUND" if m else "not found", el))
            if m:
                r.update(key2=int(m.group(1)), iv2=int(m.group(2)), gap=int(m.group(3)), found_by=label,
                         chacha_found=("FOUND key1 seed" in p.stdout) if r["chacha"] else None)
                break
            if p.returncode and "no KCipher-2 seed pair found" not in p.stdout:
                print("  recover error (exit %d): %s" % (p.returncode, p.stderr[-400:]))
        r["out_sha"] = hashlib.sha256(open(out, "rb").read()).hexdigest() if os.path.exists(out) else None

    # only now: compare with the pre-encryption listing
    pre = {os.path.basename(x[1]): (int(x[2]), x[3]) for x in rows(os.path.join(rd, a.pre)) if x[0] == "FILE"}
    print("\n## Comparison with the pre-encryption SHA-256 (%s, taken on the VM before the run)" % a.pre)
    tally = {}
    for name, r in res.items():
        if r.get("status") != "OUT OF SCOPE":
            want = pre.get(name[:-6], (None, None))[1]
            r["orig_sha"] = want
            if r.get("out_sha") and r["out_sha"] == want:
                r["status"] = "MATCH"
            elif "key2" in r:
                r["status"] = "KEYS FOUND, OUTPUT DIFFERS"
            else:
                r["status"] = "NOT FOUND"
        tally[r["status"]] = tally.get(r["status"], 0) + 1
        print("  %-60s %-26s %s" % (name[:-6][:60], r["status"], r.get("found_by", r.get("why", ""))))
    scored = len(res) - tally.get("OUT OF SCOPE", 0)
    print("\nRESULT: %d of %d in-scope files hash-exact; %s" % (tally.get("MATCH", 0), scored,
          ", ".join("%s %d" % kv for kv in sorted(tally.items()))))
    json.dump({"offset": offset, "files": res}, open(os.path.join(rd, "results.json"), "w"), indent=1)
    sys.exit(0 if tally.get("MATCH", 0) == scored else 1)


if __name__ == "__main__":
    main()