#!/usr/bin/env python3
"""Recupera un secreto repartido con ShamirVault, sin ShamirVault.

Esto existe por una razón concreta: si la aplicación desapareciera, las partes
que alguien tenga guardadas en un papel seguirían sirviendo. Con SLIP-39 eso lo
resuelve el estándar; con el formato propio lo resuelve este fichero.

Usa solo la biblioteca estándar de Python. Descifrar un secreto protegido con
contraseña necesita además el paquete `cryptography`, y el script lo dice si
falta en vez de fallar de mala manera.

    python3 recuperar.py parte1 parte2 ...
    python3 recuperar.py --contrasena "mi clave" parte1 parte2
    python3 recuperar.py --autoprueba

El formato está descrito en `docs/formato-partes.md`.
"""
import argparse
import hashlib
import json
import sys

# ── Aritmética en GF(256) ────────────────────────────────────────────
#
# Generador 3 y polinomio 0x11B, que es el de AES. El generador 2 NO vale con
# ese polinomio: solo tiene orden 51, así que no genera el campo entero.

_EXP = [0] * 512
_LOG = [0] * 256


def _tablas():
    x = 1
    for i in range(255):
        _EXP[i] = x
        _LOG[x] = i
        x = _mul3(x)
    for i in range(255, 512):
        _EXP[i] = _EXP[i - 255]


def _mul2(x):
    return ((x << 1) ^ 0x1B) & 0xFF if x & 0x80 else (x << 1) & 0xFF


def _mul3(x):
    return _mul2(x) ^ x


def _mul(a, b):
    if a == 0 or b == 0:
        return 0
    return _EXP[_LOG[a] + _LOG[b]]


def _div(a, b):
    if b == 0:
        raise ZeroDivisionError('division por cero en GF(256)')
    if a == 0:
        return 0
    return _EXP[(_LOG[a] + 255 - _LOG[b]) % 255]


_tablas()


# ── Codificación de las partes ───────────────────────────────────────

_B58 = '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'


def _de_base58(s):
    ceros = len(s) - len(s.lstrip('1'))
    n = 0
    for c in s:
        i = _B58.find(c)
        if i < 0:
            raise ValueError('caracter no valido en base58: %r' % c)
        n = n * 58 + i
    cuerpo = n.to_bytes((n.bit_length() + 7) // 8, 'big') if n else b''
    return b'\x00' * ceros + cuerpo


def _de_base64url(s):
    import base64
    return base64.urlsafe_b64decode(s + '=' * (-len(s) % 4))


def _checksum(cuerpo):
    return hashlib.sha256(cuerpo).digest()[:4]


def decodifica_parte(texto):
    """Devuelve (indice, total, umbral, datos, version).

    No se pregunta en qué formato está escrita: se prueban los tres alfabetos
    y el checksum decide. Sin él —partes de la versión 1— hay que caer en una
    heurística, que es justamente el motivo de que la versión 2 lo lleve.
    """
    t = texto.strip()
    intentos = []
    for decodificador in (_de_base58, bytes.fromhex, _de_base64url):
        try:
            intentos.append(decodificador(t))
        except Exception:
            pass

    # Primero v2: el checksum no deja ambigüedad posible.
    for blob in intentos:
        if len(blob) >= 9 and blob[0] == 0x02:
            if _checksum(blob[:-4]) == blob[-4:]:
                return blob[1], blob[2], blob[3], blob[4:-4], 2
    # Después v1, que no lleva checksum.
    for blob in intentos:
        if len(blob) >= 5 and blob[0] == 0x01:
            return blob[1], blob[2], blob[3], blob[4:], 1
    raise ValueError('no se reconoce como una parte: %r' % texto[:40])


# ── Reconstrucción ───────────────────────────────────────────────────

def reconstruye(partes):
    """Junta las partes y devuelve el secreto en bytes."""
    leidas = [decodifica_parte(p) for p in partes]

    umbral = leidas[0][2]
    version = leidas[0][4]
    if len({(i[1], i[2], i[4]) for i in leidas}) != 1:
        raise ValueError('las partes no son del mismo reparto')
    if len({i[0] for i in leidas}) != len(leidas):
        raise ValueError('hay partes repetidas')
    if len(leidas) < umbral:
        raise ValueError('hacen falta %d partes y hay %d' % (umbral, len(leidas)))
    if len({len(i[3]) for i in leidas}) != 1:
        raise ValueError('las partes tienen longitudes distintas')

    usadas = leidas[:umbral]
    xs = [p[0] for p in usadas]

    # Pesos de Lagrange evaluados en x=0: solo dependen de las coordenadas, así
    # que se calculan una vez para todos los bytes.
    pesos = []
    for i, xi in enumerate(xs):
        num = den = 1
        for j, xj in enumerate(xs):
            if i == j:
                continue
            num = _mul(num, xj)
            den = _mul(den, xi ^ xj)
        pesos.append(_div(num, den))

    salida = bytearray(len(usadas[0][3]))
    for b in range(len(salida)):
        acc = 0
        for i, parte in enumerate(usadas):
            acc ^= _mul(parte[3][b], pesos[i])
        salida[b] = acc

    if version == 1:
        return bytes(salida)

    # La versión 2 antepone al secreto los 4 primeros bytes de su SHA-256.
    resumen, secreto = bytes(salida[:4]), bytes(salida[4:])
    if hashlib.sha256(secreto).digest()[:4] != resumen:
        raise ValueError(
            'las partes encajan entre si pero el resumen no cuadra: '
            'alguna no pertenece a este secreto')
    return secreto


# ── Sobres cifrados (opcional) ───────────────────────────────────────

def descifra(sobre, contrasena):
    """Abre un sobre `SVE2:`. Necesita el paquete `cryptography`."""
    import base64
    if not sobre.startswith('SVE2:'):
        raise ValueError('no es un sobre SVE2')
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    except ImportError:
        raise SystemExit(
            'Para descifrar hace falta el paquete `cryptography`:\n'
            '    pip install cryptography\n'
            'El secreto en si ya se ha recuperado; solo falta abrirlo.')

    datos = base64.b64decode(sobre[5:])
    if datos[0] != 0x01:
        raise ValueError('derivacion de clave desconocida')
    iteraciones = int.from_bytes(datos[1:5], 'big')
    salt, nonce = datos[5:37], datos[37:49]
    aad, sellado = datos[:49], datos[49:]

    clave = hashlib.pbkdf2_hmac('sha256', contrasena.encode(), salt,
                                iteraciones, 32)
    return AESGCM(clave).decrypt(nonce, sellado, aad).decode('utf-8')


# ── Autoprueba ───────────────────────────────────────────────────────

VECTORES = json.loads(r'''[
 {
  "nombre": "texto simple 2 de 3, hex",
  "secreto": "la clave del banco",
  "n": 3,
  "k": 2,
  "encoding": "hex",
  "partes": [
   "02010302eef68e43bb8d37b1bcf769f82891ea8ca3421217e686da4b10e0",
   "02020302172cd869d9a20edcd7564844309560b73d22879c72a6479110da",
   "020303024093ea860e4e190e07c057d93860ef57be02f4e5f74f5adcaa94"
  ]
 },
 {
  "nombre": "texto simple 3 de 5, base58",
  "secreto": "abandon ability able about",
  "n": 5,
  "k": 3,
  "encoding": "base58",
  "partes": [
   "JCF6a5Jz2yP59NyKrpuNUKaaJ4rhWQV3ZPfybPRjABEsCrJoAfY",
   "JEBqmtFaXSkz1UsnsFa9inDXxfyNXER4YYNDA9qvsYyBmWY2HHC",
   "JG8auiNFpoXdPMMjQ7LHjdnoeeZBSdo4bnfoJK7yJ9iiJ84Yuz2",
   "JJ5L5mKzXjCNrYZCnjUi7nA8HTivEvXnhJXFAreV48to65kfBze",
   "JL25KV3YEjXWKVj1hjnmZseP9WhJZBbmBJRC7jyDWJ7eGpqncGR"
  ]
 },
 {
  "nombre": "acentos y emoji, base64",
  "secreto": "contraseña ñ €uro 🔐",
  "n": 4,
  "k": 2,
  "encoding": "base64",
  "partes": [
   "AgEEAuW3YYdiSu6MY6oKRbxHNH13KpkkwAHgyt5gBKUv9JcjWJ8",
   "AgIEAv6PW0phJXWfUOyBJT1Gy5qwnEl1Bu1EGRagA-v5WEH9U44",
   "AgMEAvduTfhgAPVnQSf4BUKwnscEB_CzREDRoafg99FCPJC0RXQ",
   "AgQEAsj_L8tn-1i5NmCM5SRELk8l6_LXkS4XpJ07DXdOGwOoa6I"
  ]
 },
 {
  "nombre": "umbral igual al total",
  "secreto": "todas hacen falta",
  "n": 3,
  "k": 3,
  "encoding": "base58",
  "partes": [
   "6H31eHqV7u6upKPrDiFKLhQtvYiWYWDBMQvgHGc",
   "6HdcRg2YcmiRheUEH8PkYYNARZ6UU36uz49DfNV",
   "6JEDCnWUcRuETaBhXwcFw5BANMcqJQVE8ZxJrtJ"
  ]
 },
 {
  "nombre": "cifrado con contrasena, base58",
  "secreto": "secreto bajo contrasena",
  "passphrase": "clave 2026",
  "n": 3,
  "k": 2,
  "encoding": "base58",
  "partes": [
   "fkseeb4wv54KjWRrSaemWRhaqsJHdMK9ByWgxZQxeRpCWZkJKJL9hz6prQdsKYyMzxBKJqVQ46z8UjkTejzuoVPqx1omxWYDsocZUqTfRUuwr4r1vDNAJamtacsEwqbRcZZbo9rZYDerPZEuVrFG44EhheZu3x1QPyBC75arM4HXPuWPosPi3PXVuq",
   "fqFnnRQPZTCY2zdjJPrbK7MqoKRHdetyzbsXGoSWmjz7cFZabDeL4tm4TY4hUJX8x3LrP26KQVTRxFfiFtjUCCg7Ft8edhJZQKoXBbk48EoKuAN8frurNQCnDcDzHirxSLBpKJPxHHvWu4sY6G4vt7UGqEQ2WoL2NPaaPho9XrU8zghgKafiEPBrSe",
   "fudvvfrB5ABKhxYey9B3FPPmZpYVXJbBLwb3s1HHRzxcPbrHd4uw7qBP41wMnm12rHHAc7SzoFN1F8CxF5ZLJvovbZghGkhaiPjM7e27g3HpSYLVE1XcEnP16g4Kd71v9m3k5pBKnY3B4vHZpq6Kfeq2pEouZxaECTQf8uGqeCMxbcucDFVKZY8RA5"
  ]
 }
]''')


def autoprueba():
    fallos = 0
    for caso in VECTORES:
        try:
            recuperado = reconstruye(caso['partes'][:caso['k']]).decode('utf-8')
            if 'passphrase' in caso:
                recuperado = descifra(recuperado, caso['passphrase'])
            ok = recuperado == caso['secreto']
        except SystemExit as e:
            print('  OMITIDO  %-34s %s' % (caso['nombre'], e))
            continue
        except Exception as e:
            ok, recuperado = False, 'error: %s' % e
        print('  %s  %-34s %s' % ('OK    ' if ok else 'FALLO ', caso['nombre'],
                                  '' if ok else repr(recuperado)[:60]))
        fallos += 0 if ok else 1

    # Y que rechace lo que tiene que rechazar.
    for nombre, partes in (
        ('parte inventada', ['zzzz']),
        ('partes de repartos distintos',
         [VECTORES[0]['partes'][0], VECTORES[1]['partes'][0]]),
        ('partes insuficientes', VECTORES[1]['partes'][:1]),
    ):
        try:
            reconstruye(partes)
            print('  FALLO   %-34s deberia haber sido rechazado' % nombre)
            fallos += 1
        except Exception:
            print('  OK      %-34s rechazado' % nombre)

    print('\n%s' % ('Todo correcto.' if not fallos else '%d FALLOS' % fallos))
    return 1 if fallos else 0


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('partes', nargs='*', help='las partes, tal cual estan escritas')
    ap.add_argument('--contrasena', help='si el secreto se protegio con una')
    ap.add_argument('--autoprueba', action='store_true',
                    help='comprueba el script contra vectores conocidos')
    args = ap.parse_args()

    if args.autoprueba:
        return autoprueba()
    if not args.partes:
        ap.print_help()
        return 2

    # Los errores se cuentan, no se lanzan: quien usa esto esta recuperando
    # algo en un mal momento, y una traza de Python no le dice nada.
    try:
        secreto = reconstruye(args.partes)
    except ValueError as e:
        print('No se ha podido recuperar: %s' % e, file=sys.stderr)
        return 1

    try:
        texto = secreto.decode('utf-8')
    except UnicodeDecodeError:
        print('El secreto no es texto. En hexadecimal:')
        print(secreto.hex())
        return 0

    if texto.startswith('SVE2:') or texto.startswith('SVE1:'):
        if not args.contrasena:
            print('El secreto esta protegido con contrasena. '
                  'Vuelve a ejecutarlo con --contrasena.', file=sys.stderr)
            return 1
        try:
            texto = descifra(texto, args.contrasena)
        except Exception:
            # Con cifrado autenticado no hay forma de distinguir una
            # contrasena equivocada de un sobre manipulado, y decir que si la
            # habria seria mentir.
            print('No se ha podido abrir: o la contrasena no es esa, o el '
                  'sobre esta danado.', file=sys.stderr)
            return 1

    print(texto)
    return 0


if __name__ == '__main__':
    sys.exit(main())