#!/usr/bin/env python3
"""Open a KEEP/01 backup file on a computer, using your recovery code.

    python recover.py keep01-backup.k01
    python recover.py keep01-backup.k01 --json > vault.json     (careful: plain text!)
    python recover.py keep01-backup.zip          (the .zip saved by the phone scanner works too)

Needs only standard Python 3. Run it on a computer you trust, ideally offline, because it
shows your passwords on screen. The recovery code is asked for without echoing it.
"""
import argparse, getpass, hashlib, hmac, json, struct, sys

KDF_ITER = 20000                     # must match vault.py on the device
BMAGIC = b"K01B"
ALPH = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"


def h(*parts):
    s = hashlib.sha256()
    for p in parts:
        s.update(p)
    return s.digest()


def mac(key, msg):
    return hmac.new(key, msg, hashlib.sha256).digest()


def xcrypt(key, iv, data):
    out = bytearray(len(data))
    for blk in range(0, len(data), 32):
        ks = mac(key, iv + struct.pack("<I", blk // 32))
        for i in range(min(32, len(data) - blk)):
            out[blk + i] = data[blk + i] ^ ks[i]
    return bytes(out)


def kdf(secret, salt, iters):
    k = h(salt, secret)
    for _ in range(iters):
        k = h(k, salt, secret)
    return h(k, b"enc"), h(k, b"mac")


def normalize(s):
    s = s.upper().replace("O", "0").replace("I", "1").replace("L", "1")
    return "".join(c for c in s if c in ALPH)


def code_ok(code):
    if len(code) != 24:
        return False
    d = h(b"keep01-code", code[:22].encode())
    v = (d[0] << 8) | d[1]
    return code[22:] == ALPH[(v >> 11) & 31] + ALPH[(v >> 6) & 31]


def open_backup(blob, code, iters=KDF_ITER):
    if blob[:4] != BMAGIC:
        raise SystemExit("This isn't a KEEP/01 backup file.")
    salt, wrapped, data = blob[4:20], blob[20:100], blob[100:]
    enc, mk_mac = kdf(code.encode(), salt, iters)
    iv, ct, tag = wrapped[:16], wrapped[16:48], wrapped[48:80]
    if not hmac.compare_digest(mac(mk_mac, iv + ct), tag):
        raise SystemExit("Wrong recovery code for this backup.")
    mk = xcrypt(enc, iv, ct)
    denc, dmac = h(mk, b"data-enc"), h(mk, b"data-mac")
    iv, tag, ct = data[:16], data[16:48], data[48:]
    if not hmac.compare_digest(mac(dmac, iv + ct), tag):
        raise SystemExit("The backup file is damaged or was changed.")
    return json.loads(xcrypt(denc, iv, ct).decode())


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("file")
    ap.add_argument("--json", action="store_true", help="print everything as JSON")
    ap.add_argument("--code", help="recovery code (otherwise asked for privately)")
    ap.add_argument("--iter", type=int, default=KDF_ITER, help=argparse.SUPPRESS)
    a = ap.parse_args()
    blob = open(a.file, "rb").read()
    if blob[:2] == b"PK":                        # .zip from the phone scanner page
        import io, zipfile
        with zipfile.ZipFile(io.BytesIO(blob)) as z:
            names = [n for n in z.namelist() if n.endswith(".k01")]
            if not names:
                raise SystemExit("No .k01 backup inside this zip.")
            blob = z.read(names[0])
    code = normalize(a.code or getpass.getpass("Recovery code: "))
    if not code_ok(code):
        raise SystemExit("That recovery code has a typo (check each group of 4).")
    data = open_backup(blob, code, a.iter)
    items = data["items"]
    if a.json:
        print(json.dumps(items, indent=2, ensure_ascii=False))
        return
    print(f"\n{len(items)} item(s)\n")
    for it in items:
        if it.get("kind") == "seed":
            print(f"[SEED] {it['label']}  ({it.get('date', '')})")
            ws = it["words"]
            for i in range(0, len(ws), 4):
                print("   " + "  ".join(f"{j + 1:2d}. {ws[j]:<9}" for j in range(i, min(i + 4, len(ws)))))
        else:
            print(f"{it['label']:<22} {it['pw']}")
            if it.get("old"):
                print(f"{'':<22} (old: {it['old']})")
    print()


if __name__ == "__main__":
    main()