"""
control/crypto.py
-----------------
Symmetric encryption for sensitive control-plane fields (tenant DB passwords).

Fernet (AES-128-CBC + HMAC). MT-23 adds key rotation via ``MultiFernet``.

Key configuration
-----------------
``CONTROL_FERNET_KEY``
    The **primary** key. Everything is encrypted with this one.

``CONTROL_FERNET_KEYS_OLD``
    Comma-separated retired keys, newest first. Decrypt-only: they open
    ciphertext written before the last rotation, and are never used to encrypt.
    Empty or unset in steady state.

``CONTROL_FERNET_KEY_VERSION``
    Integer, default 1. Stamped onto every row as it is (re-)encrypted, so
    rotation progress is a SQL query rather than a guess:

        SELECT COUNT(*) FROM tenants WHERE key_version < <current>;

Generate a key:

    python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())"

Rotation procedure
------------------
1. Generate a new key.
2. Move the current ``CONTROL_FERNET_KEY`` value into ``CONTROL_FERNET_KEYS_OLD``
   (prepending it if that variable already holds keys), set the new key as
   ``CONTROL_FERNET_KEY``, and bump ``CONTROL_FERNET_KEY_VERSION``.
3. ``python -m control.cli verify-keys`` — confirms every stored credential is
   still readable under the new key set **before** anything is rewritten.
4. ``python -m control.cli rotate-keys`` — re-encrypts every row under the
   primary key and stamps the new version.
5. ``python -m control.cli verify-keys`` again, then drop the retired key from
   ``CONTROL_FERNET_KEYS_OLD``.

Steps 3 and 5 are the point of the design: at no moment is there a window
where a key has been retired but rows still depend on it.
"""

import hashlib
import os

from cryptography.fernet import Fernet, MultiFernet

_KEY_ENV     = 'CONTROL_FERNET_KEY'
_OLD_KEY_ENV = 'CONTROL_FERNET_KEYS_OLD'
_VERSION_ENV = 'CONTROL_FERNET_KEY_VERSION'


def _as_bytes(key):
    return key.encode() if isinstance(key, str) else key


def primary_key():
    """The single key used for encryption. Raises if unset."""
    key = os.environ.get(_KEY_ENV)
    if not key:
        raise RuntimeError(
            f"{_KEY_ENV} is not set. Generate one with: "
            "python -c \"from cryptography.fernet import Fernet; "
            "print(Fernet.generate_key().decode())\""
        )
    return _as_bytes(key)


def old_keys():
    """Retired decrypt-only keys, in the order they are tried."""
    raw = os.environ.get(_OLD_KEY_ENV) or ''
    return [_as_bytes(k.strip()) for k in raw.split(',') if k.strip()]


def current_key_version():
    """Integer stamped onto rows encrypted under the current primary key."""
    try:
        return int(os.environ.get(_VERSION_ENV, 1))
    except (TypeError, ValueError):
        return 1


def key_fingerprint(key=None):
    """Short, non-reversible identifier for a key — safe to log.

    Lets an operator confirm which key a process actually loaded without ever
    putting key material in a log line or a support ticket.
    """
    material = _as_bytes(key) if key is not None else primary_key()
    return hashlib.sha256(material).hexdigest()[:12]


def _multi():
    """MultiFernet: primary first (it encrypts), retired keys after it.

    Ordering matters — MultiFernet encrypts with the first key and tries each
    in turn when decrypting.
    """
    keys = [Fernet(primary_key())]
    seen = {primary_key()}
    for k in old_keys():
        if k in seen:
            continue          # tolerate a stale copy of the primary in the old list
        seen.add(k)
        keys.append(Fernet(k))
    return MultiFernet(keys)


# Kept for backward compatibility: earlier code imported _fernet() directly.
def _fernet():
    return _multi()


def encrypt(plaintext):
    """Encrypt a string to a Fernet token (str) under the primary key.

    Passes None through.
    """
    if plaintext is None:
        return None
    return _multi().encrypt(plaintext.encode()).decode()


def decrypt(token):
    """Decrypt a Fernet token, trying the primary key then each retired key.

    Passes None through. Raises ``cryptography.fernet.InvalidToken`` when no
    configured key opens the token.
    """
    if token is None:
        return None
    return _multi().decrypt(token.encode()).decode()


def rotate(token):
    """Re-encrypt an existing token under the primary key.

    ``MultiFernet.rotate`` decrypts with whichever key works and re-encrypts
    with the first, preserving the plaintext without it ever being handled
    here. Passes None through.
    """
    if token is None:
        return None
    return _multi().rotate(token.encode()).decode()


def can_decrypt(token):
    """True when some configured key opens this token. Never raises.

    Used by ``control.cli verify-keys`` to audit every stored credential
    before a key is retired.
    """
    if token is None:
        return True           # nothing stored is not a failure
    try:
        decrypt(token)
        return True
    except Exception:
        return False