#!/usr/bin/env python3
"""pass.new recovery: rebuild a password from any 3 of its 5 blocks.

Works offline with plain Python 3 (no installs). Run it, then type or paste 3 blocks:
    python3 recover.py
or pass them as arguments:
    python3 recover.py "BLOCK ONE" "BLOCK TWO" "BLOCK THREE"
Public domain. Format details: see README.md in this kit.
"""
import hashlib, itertools, sys

# Block text is Crockford Base32: no I, L, O or U. Case, spaces and dashes don't matter.
ALPHABET = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
VALUE = {c: i for i, c in enumerate(ALPHABET)}
VALUE.update({"O": 0, "I": 1, "L": 1})

def read_block(text):
    """Block text -> (block number 1-5, set id, data bytes). Checks the block's checksum."""
    if "r=" in text:                      # a scanned QR link: .../split#r=BLOCK
        text = text.split("r=", 1)[1]
    buffer = bits = 0
    raw = bytearray()
    for ch in text.upper():
        if ch in " -_.\t\r\n":
            continue
        if ch not in VALUE:
            raise ValueError(f"'{ch}' is not a character used in blocks")
        buffer = (buffer << 5) | VALUE[ch]
        bits += 5
        if bits >= 8:
            bits -= 8
            raw.append((buffer >> bits) & 0xFF)
        buffer &= (1 << bits) - 1
    body, check = bytes(raw[:-4]), bytes(raw[-4:])
    if len(body) < 8 or hashlib.sha256(body).digest()[:4] != check:
        raise ValueError("this block doesn't check out (probably a typo or a missing character)")
    if body[0] != 1 or body[1] != 0x35:  # version 1, "any 3 of 5"
        raise ValueError("this isn't a pass.new version 1 block")
    return body[2], body[3:7].hex().upper(), body[7:]

# Arithmetic in GF(256) with the AES polynomial x^8 + x^4 + x^3 + x + 1 (0x11B).
EXP, LOG = [0] * 510, [0] * 256
x = 1
for i in range(255):
    EXP[i], LOG[x] = x, i
    x ^= ((x << 1) ^ (0x1B if x & 0x80 else 0)) & 0xFF   # multiply by 3
for i in range(255, 510):
    EXP[i] = EXP[i - 255]

def mul(a, b):
    return 0 if a == 0 or b == 0 else EXP[LOG[a] + LOG[b]]

def div(a, b):
    return 0 if a == 0 else EXP[(LOG[a] + 255 - LOG[b]) % 255]

def combine(blocks):
    """Lagrange interpolation at x = 0, byte by byte (Shamir's method, 3 of 5)."""
    out = bytearray(len(blocks[0][2]))
    for i, (xi, _, yi) in enumerate(blocks):
        weight = 1
        for j, (xj, _, _) in enumerate(blocks):
            if i != j:
                weight = mul(weight, div(xj, xj ^ xi))
        for k, byte in enumerate(yi):
            out[k] ^= mul(byte, weight)
    return bytes(out)

def unpack(data):
    """[2-byte length] + UTF-8 password + first 4 bytes of SHA-256(password) + zero padding."""
    n = (data[0] << 8) | data[1]
    password, fingerprint = data[2:2 + n], data[2 + n:6 + n]
    if n == 0 or len(password) != n or hashlib.sha256(password).digest()[:4] != fingerprint:
        return None
    return password.decode("utf-8")

def recover(texts):
    blocks = {}
    for text in texts:
        number, set_id, data = read_block(text)
        blocks[number] = (number, set_id, data)
    if len({b[1] for b in blocks.values()}) > 1:
        raise ValueError("these blocks come from different sets (the set IDs don't match)")
    if len(blocks) < 3:
        raise ValueError(f"you need 3 different blocks; you have {len(blocks)}")
    for three in itertools.combinations(blocks.values(), 3):
        password = unpack(combine(list(three)))
        if password is not None:
            return password
    raise ValueError("these blocks don't fit together; one may be damaged or from another set")

if __name__ == "__main__":
    texts = sys.argv[1:]
    while len(texts) < 3:
        try:
            line = input(f"Block {len(texts) + 1} of 3 (paste, then press Enter): ").strip()
        except EOFError:
            break
        if line:
            try:
                read_block(line)
                texts.append(line)
            except ValueError as err:
                print(f"  Problem: {err}. Please try that block again.")
    try:
        password = recover(texts)
    except ValueError as err:
        sys.exit(f"Could not recover the password: {err}")
    print("\nYour password is between the lines below (spaces at either end are part of it):")
    print("-" * 40)
    print(password)
    print("-" * 40)
    print(f"{len(password)} characters, verified as an exact match.")