# SPDX-License-Identifier: LGPL-3.0-or-later
"""
The keys a member derives from a passphrase, and the identities they open.

Byte for byte what `static/keyderive.js` does (docs/MESHBAY_DESIGN.md §3.1,
§3.7); `test_bundle_kdf_parity.py` holds this module to it. A mismatch does not
look like an error — it looks like an account nobody can open.
"""

import base64
import hashlib
import json
from dataclasses import dataclass

from argon2.low_level import Type, hash_secret_raw
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

PBKDF2_ITERATIONS = 600_000
ARGON2_MEM_KIB = 131072   # 128 MB
ARGON2_TIME = 3
ARGON2_LANES = 1

BUNDLE_MAGIC = b"MBK3"
LEGACY_MAGIC = b"MBK2"    # TRANSITIONAL — read, never written


def b64(data: bytes) -> str:
    return base64.b64encode(data).decode()


def unb64(text: str) -> bytes:
    return base64.b64decode(text)


def hkdf(key: bytes, info: str) -> bytes:
    return HKDF(algorithm=hashes.SHA256(), length=32, salt=None,
                info=info.encode()).derive(key)


def auth_key(username: str, passphrase: str) -> str:
    """What the hub checks at sign-in. The passphrase itself never leaves here."""
    salt = hashlib.sha256(f"meshbay:auth:v1:{username}".encode()).digest()
    return b64(PBKDF2HMAC(algorithm=hashes.SHA512(), length=32, salt=salt,
                          iterations=PBKDF2_ITERATIONS).derive(passphrase.encode()))


def passphrase_half(username: str, passphrase: str) -> bytes:
    """`A`: Argon2id of the passphrase, salted by the account name."""
    salt = hashlib.sha256(f"meshbay:bundle:v2:{username}".encode()).digest()[:16]
    return hash_secret_raw(passphrase.encode(), salt, time_cost=ARGON2_TIME,
                           memory_cost=ARGON2_MEM_KIB, parallelism=ARGON2_LANES,
                           hash_len=32, type=Type.ID)


def bundle_master_key(a: bytes, pepper: bytes, user_id: str) -> bytes:
    """`M`, from `A` and the pepper the hub hands out at sign-in."""
    return hkdf(a + pepper, f"meshbay:bundle-master:v3|{user_id}")


def node_bundle_key(master: bytes, node_pk: str) -> bytes:
    return hkdf(master, f"meshbay:bundle:v3|node|{node_pk}")


class BundleError(Exception):
    pass


def open_bundle(bundle_b64: str, *, master: bytes | None, legacy: bytes | None,
                user_id: str, node_pk: str) -> dict:
    """
    The `{skEd, skX}` sealed in a keypair bundle.

    `MBK3` opens under the node's key derived from `M`, with the account and the
    node as associated data. `MBK2` opens under `A` alone; it is read so a member
    is not locked out of a node still holding one, and never rewritten here —
    the application or a browser replaces it. Anything older is refused by name.
    """
    try:
        raw = unb64(bundle_b64)
    except Exception as e:
        raise BundleError("the node's copy of your keys is unreadable") from e
    try:
        if raw[:4] == BUNDLE_MAGIC:
            if master is None:
                raise BundleError("no bundle key in this session")
            aad = f"meshbay:bundle:v3|{user_id}|{node_pk}".encode()
            plain = AESGCM(node_bundle_key(master, node_pk)).decrypt(raw[5:17], raw[17:], aad)
        elif raw[:4] == LEGACY_MAGIC:
            if legacy is None:
                raise BundleError("this node holds an older copy of your keys; sign in "
                                  "again with your passphrase to open it")
            plain = AESGCM(legacy).decrypt(raw[4:16], raw[16:], None)
        else:
            raise BundleError("this node holds your keys in a format no longer read; "
                              "ask its operator to unpin you and send a fresh code")
    except BundleError:
        raise
    except Exception as e:
        raise BundleError("the node's copy of your keys does not open with this "
                          "passphrase") from e
    return json.loads(plain)


@dataclass
class Identity:
    """One account's keys on one node."""

    sk_ed: Ed25519PrivateKey
    sk_x: X25519PrivateKey

    @classmethod
    def from_pkcs8(cls, keys: dict) -> "Identity":
        sk_ed = serialization.load_der_private_key(unb64(keys["skEd"]), password=None)
        sk_x = serialization.load_der_private_key(unb64(keys["skX"]), password=None)
        if not isinstance(sk_ed, Ed25519PrivateKey) or not isinstance(sk_x, X25519PrivateKey):
            raise BundleError("the node's copy of your keys holds the wrong kind of key")
        return cls(sk_ed, sk_x)

    def to_pkcs8(self) -> dict:
        def der(k) -> str:
            return b64(k.private_bytes(serialization.Encoding.DER,
                                       serialization.PrivateFormat.PKCS8,
                                       serialization.NoEncryption()))
        return {"skEd": der(self.sk_ed), "skX": der(self.sk_x)}

    @property
    def pk_ed(self) -> str:
        return b64(self.sk_ed.public_key().public_bytes_raw())

    @property
    def pk_x(self) -> str:
        return b64(self.sk_x.public_key().public_bytes_raw())