#!/usr/bin/env python3
"""Passphrase-based encryption for PIC backup archives.

A backup archive contains key material (WireGuard keys, the vault Fernet key,
the internal CA, admin credentials). When the operator supplies a passphrase we
encrypt the archive at rest.

The repo's only available crypto primitive is `cryptography` (Fernet, scrypt) —
PyNaCl / the age binary are not installed in the API image. We therefore derive
a Fernet key from the passphrase with scrypt and wrap the archive bytes. The
encrypted file keeps the `.age` extension expected by the UI/restore detection;
the embedded MAGIC distinguishes our format from a real age file.
"""

import os
import struct
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
import base64

# File layout: MAGIC | salt(16) | n(4) | r(4) | p(4) | fernet_token
MAGIC = b'PICBKP1\n'
_SALT_LEN = 16
# scrypt cost parameters (interactive-strong; ~tens of ms)
_N = 2 ** 15
_R = 8
_P = 1


class BackupDecryptError(Exception):
    """Raised when an encrypted backup cannot be decrypted (wrong passphrase)."""


def _derive_key(passphrase: str, salt: bytes, n: int, r: int, p: int) -> bytes:
    kdf = Scrypt(salt=salt, length=32, n=n, r=r, p=p)
    raw = kdf.derive(passphrase.encode('utf-8'))
    return base64.urlsafe_b64encode(raw)


def encrypt_bytes(plaintext: bytes, passphrase: str) -> bytes:
    """Encrypt archive bytes with a passphrase. Returns the on-disk blob."""
    if not passphrase:
        raise ValueError('passphrase required for encryption')
    salt = os.urandom(_SALT_LEN)
    key = _derive_key(passphrase, salt, _N, _R, _P)
    token = Fernet(key).encrypt(plaintext)
    header = MAGIC + salt + struct.pack('>III', _N, _R, _P)
    return header + token


def is_encrypted(blob: bytes) -> bool:
    return blob[:len(MAGIC)] == MAGIC


def decrypt_bytes(blob: bytes, passphrase: str) -> bytes:
    """Decrypt a blob produced by encrypt_bytes. Raises BackupDecryptError."""
    if not is_encrypted(blob):
        raise BackupDecryptError('not a PIC encrypted backup')
    if not passphrase:
        raise BackupDecryptError('passphrase required')
    off = len(MAGIC)
    salt = blob[off:off + _SALT_LEN]
    off += _SALT_LEN
    n, r, p = struct.unpack('>III', blob[off:off + 12])
    off += 12
    token = blob[off:]
    key = _derive_key(passphrase, salt, n, r, p)
    try:
        return Fernet(key).decrypt(token)
    except (InvalidToken, ValueError) as e:
        raise BackupDecryptError('invalid passphrase or corrupt archive') from e