#!/usr/bin/env python3
"""
RED FOLDER - awaryjne odzyskiwanie (format v1), niezalezne od aplikacji HTML.

Wymaga: Python 3.8+ oraz biblioteki 'cryptography'  (pip install cryptography)

Uzycie:
  Rodzina (karty):   python3 odzyskaj_teczke.py czerwona_teczka.json Karta_Anna.pdf Karta_Piotr.pdf
                     (karty jako PDF z aplikacji albo pliki .txt z przepisanym kodem CT1:...)
  Wlasciciel:        python3 odzyskaj_teczke.py teczka.json --haslo
Odszyfrowane pliki i notatki trafiaja do folderu "teczka_odszyfrowana".

Opis formatu (gdyby trzeba bylo napisac program od nowa):
  * Klucz prywatny sejfu: ECDH P-256 (PKCS#8 DER).
  * vault.owner:    AES-256-GCM(klucz = PBKDF2-SHA256(haslo NFC, salt, iter), AAD = "owner:<vaultId>") -> klucz prywatny
  * vault.recovery: AES-256-GCM(klucz = R, AAD = "recovery:<vaultId>:<setId>") -> klucz prywatny
  * R (32 bajty) podzielony schematem Shamira w GF(256) (wielomian 0x11b), x = 1..n, prog k = floor(n/2)+1.
  * Karta (w PDF jako niewidoczny tekst, w QR i jako kod zapasowy): "CT1:" + base64url(JSON{vid,set,x,k,n,name,s=salt,i=iv,c=szyfrogram});
           y = AES-256-GCM(PBKDF2-SHA256(PIN, salt, 600000), AAD = "card:<vid>:<set>:<x>")
  * Wpis: efemeryczny klucz ECDH P-256 (eph, surowy 65 B); wspolny sekret -> HKDF-SHA256(salt = eph,
          info = "czerwona-teczka v1 entry") -> AES-256-GCM, AAD = "entry:<vaultId>:<id>"; tresc = JSON.
"""
import base64, getpass, json, os, re, sys, unicodedata
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec


def b64(s): return base64.b64decode(s)
def b64url(s): return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))


def pbkdf2(pw, salt, it):
    return PBKDF2HMAC(hashes.SHA256(), 32, salt, it).derive(unicodedata.normalize("NFC", pw).encode())


def gcm_open(key, box, aad):
    return AESGCM(key).decrypt(b64(box["iv"]), b64(box["ct"]), aad.encode())


# --- Shamir GF(256) ---
EXP, LOG = [0] * 512, [0] * 256
x = 1
for i in range(255):
    EXP[i] = x; LOG[x] = i
    x ^= ((x << 1) ^ (0x11b if x & 0x80 else 0)); x &= 0xff
for i in range(255, 512): EXP[i] = EXP[i - 255]
def gmul(a, b): return EXP[LOG[a] + LOG[b]] if a and b else 0
def gdiv(a, b): return EXP[(LOG[a] + 255 - LOG[b]) % 255] if a else 0


def combine(shares):
    out = bytearray(len(shares[0][1]))
    for i in range(len(out)):
        s = 0
        for a, (xa, ya) in enumerate(shares):
            num = den = 1
            for b, (xb, _) in enumerate(shares):
                if a != b: num = gmul(num, xb); den = gmul(den, xa ^ xb)
            s ^= gmul(ya[i], gdiv(num, den))
        out[i] = s
    return bytes(out)


def read_cards(path):
    """Karta: plik PDF z aplikacji (warstwa tekstowa) albo plik tekstowy z kodem CT1:...
    Jeden PDF moze zawierac kilka kart (wydruk zbiorczy)."""
    raw = open(path, "rb").read().decode("latin1")
    if not raw.startswith("%PDF"):
        raw = re.sub(r"\s+", "", raw)
    codes = re.findall(r"CT1:([A-Za-z0-9_-]{40,})", raw)
    if not codes: sys.exit(f"{path}: nie znaleziono kodu karty (CT1:...). Przepisz kod zapasowy z karty do pliku .txt.")
    return [json.loads(b64url(c)) for c in dict.fromkeys(codes)]


def main():
    a = sys.argv[1:]
    if not a: print(__doc__); sys.exit(1)
    v = json.load(open(a[0], encoding="utf-8"))
    vid = v["vaultId"]
    if "--haslo" in a:
        o = v["owner"]
        key = pbkdf2(getpass.getpass("Haslo wlasciciela: "), b64(o["salt"]), o["iter"])
        priv = gcm_open(key, o, f"owner:{vid}")
    else:
        need = v["recovery"].get("k") or len(v["meta"]["members"]) // 2 + 1
        print(f"Potrzebne karty: {need}. Pusty PIN = pomin te karte.")
        shares, seen = [], set()
        cards = [c for p in a[1:] for c in read_cards(p)]
        for c in cards:
            if len(shares) >= need: break
            if c["x"] in seen: continue
            seen.add(c["x"])
            if c["vid"] != vid: print(f"Karta {c['name']} nalezy do innej teczki - pomijam."); continue
            if c["set"] != v["recovery"]["setId"]: print(f"Karta {c['name']} jest uniewazniona - pomijam."); continue
            for proba in range(3):
                pin = getpass.getpass(f"PIN dla karty {c['name']}: ")
                if not pin: break
                try:
                    key = pbkdf2(pin, b64(c["s"]), 600000)
                    shares.append((c["x"], AESGCM(key).decrypt(b64(c["i"]), b64(c["c"]), f"card:{c['vid']}:{c['set']}:{c['x']}".encode())))
                    break
                except Exception:
                    print("Bledny PIN.")
        if len(shares) < need: sys.exit(f"Za malo kart: {len(shares)} z {need}.")
        R = combine(shares)
        r = v["recovery"]
        priv = gcm_open(R, r, f"recovery:{vid}:{r['setId']}")
    pk = serialization.load_der_private_key(priv, None)
    outdir = "teczka_odszyfrowana"; os.makedirs(outdir, exist_ok=True)
    notes = []
    for e in v["entries"]:
        eph_raw = b64(e["eph"])
        eph = ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), eph_raw)
        shared = pk.exchange(ec.ECDH(), eph)
        k = HKDF(hashes.SHA256(), 32, eph_raw, b"czerwona-teczka v1 entry").derive(shared)
        d = json.loads(gcm_open(k, e, f"entry:{vid}:{e['id']}"))
        if d["kind"] == "file":
            parts = [p for p in re.split(r"[\\/]+", d["name"]) if p not in ("", ".", "..")] or [e["id"]]
            path = os.path.join(outdir, *parts)
            os.makedirs(os.path.dirname(path), exist_ok=True)
            if os.path.exists(path):
                path = os.path.join(os.path.dirname(path), e["id"] + "_" + parts[-1])
            open(path, "wb").write(b64(d["data"]))
            print("plik:", path)
        else:
            notes.append(f"=== {d.get('title','')} ({e['added'][:10]}) ===\n{d.get('text','')}\n")
    if notes:
        open(os.path.join(outdir, "notatki.txt"), "w", encoding="utf-8").write("\n".join(notes))
        print("notatki:", os.path.join(outdir, "notatki.txt"))
    print("Gotowe.")


if __name__ == "__main__":
    main()