WebChecker/utils/crypto.py

"""
utils/crypto.py — Fernet symmetric encryption for website credentials.

Key derivation:
  - A 32-byte random salt is generated on first use and stored in the
    app_settings table (key: 'crypto.salt') instead of config.ini.
  - The Fernet key is derived from the salt + a fixed application secret
    using PBKDF2-HMAC-SHA256 (100,000 iterations).
  - Because the salt lives in the shared MySQL database, any machine that
    connects to the same DB can decrypt credentials without needing a local
    config.ini — making the app fully portable across machines.

Migration:
  - On first start after this change, if config.ini still contains a
    [crypto]/salt entry it is automatically migrated to app_settings and
    removed from the file.
  - _decrypt() tries Fernet first; if that fails it returns the raw value
    unchanged so that plaintext legacy credentials are still readable.
  - Callers should re-encrypt on next write (update_website handles this).
"""

import base64
import logging
import os

from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes

logger = logging.getLogger("crypto")

_APP_SECRET  = b"WebsiteChecker-v1-CredentialKey"
_ITERATIONS  = 100_000
_SETTING_KEY = "crypto.salt"
_fernet: "Fernet | None" = None


# ─── Key bootstrap ────────────────────────────────────────────────────────────

def _ensure_app_settings_table() -> bool:
    """
    Guarantee the app_settings table exists before we try to read/write it.
    Returns True if the table is available, False if it could not be created
    (e.g. the DB pool itself is not yet ready).

    This guard is necessary because crypto.py can be called during login —
    before initialize_database() has had a chance to run on a fresh install
    or an upgraded database that doesn't yet have the app_settings table.
    """
    try:
        from config import get_connection
        conn = get_connection()
        cur  = conn.cursor()
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS app_settings (
                key_name   VARCHAR(100) NOT NULL PRIMARY KEY,
                value      TEXT         NULL,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
                           ON UPDATE CURRENT_TIMESTAMP
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """
        )
        conn.commit()
        cur.close()
        conn.close()
        return True
    except Exception as e:
        logger.warning(f"Crypto: could not ensure app_settings table: {e}")
        return False


def _get_or_create_salt() -> bytes:
    """
    Read the Fernet salt from app_settings, falling back to config.ini for
    legacy installs, and generating a fresh salt for brand-new installs.

    Order of precedence:
      1. app_settings table  (primary — shared across machines via the DB)
      2. config.ini [crypto]/salt  (legacy migration path)
      3. Generate a new random salt and persist it to app_settings

    The table is created here if it doesn't yet exist, so this function is
    safe to call before initialize_database() has run.
    """
    # ── Step 1: try config.ini first (fastest, no DB needed yet) ─────────────
    # Reading config.ini for the salt is always safe — it doesn't require the
    # app_settings table to exist. If found, we attempt to also persist it to
    # the DB (best-effort), but we use the value regardless.
    legacy_b64 = None
    try:
        import configparser, os as _os
        _cfg_file = "config.ini"
        if _os.path.exists(_cfg_file):
            cfg = configparser.ConfigParser()
            cfg.read(_cfg_file, encoding="utf-8")
            if cfg.has_option("crypto", "salt"):
                legacy_b64 = cfg.get("crypto", "salt")
    except Exception as e:
        logger.warning(f"Crypto: could not read config.ini for salt: {e}")

    # ── Step 2: ensure app_settings table exists ──────────────────────────────
    table_ok = _ensure_app_settings_table()

    # ── Step 3: try to read salt from DB ──────────────────────────────────────
    if table_ok:
        try:
            from config import get_setting
            raw = get_setting(_SETTING_KEY, "")
            if raw:
                return base64.b64decode(raw)
        except Exception as e:
            logger.warning(f"Crypto: could not read salt from app_settings: {e}")

    # ── Step 4: migrate legacy salt from config.ini → DB ─────────────────────
    if legacy_b64:
        if table_ok:
            try:
                from config import set_setting
                set_setting(_SETTING_KEY, legacy_b64)
                logger.info("Crypto: migrated salt from config.ini to app_settings.")
            except Exception as e:
                logger.warning(f"Crypto: could not persist migrated salt to DB: {e}")
        else:
            logger.info("Crypto: using salt from config.ini (app_settings not available yet).")
        return base64.b64decode(legacy_b64)

    # ── Step 5: generate a fresh salt ────────────────────────────────────────
    salt     = os.urandom(32)
    b64_salt = base64.b64encode(salt).decode("ascii")
    if table_ok:
        try:
            from config import set_setting
            set_setting(_SETTING_KEY, b64_salt)
            logger.info("Crypto: generated and persisted new salt to app_settings.")
        except Exception as e:
            logger.warning(f"Crypto: could not persist new salt to DB: {e}")
    else:
        logger.warning(
            "Crypto: generated a new salt but app_settings is not available — "
            "salt will NOT persist across restarts until the table is created."
        )
    return salt


def _build_fernet() -> Fernet:
    salt = _get_or_create_salt()
    kdf  = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=_ITERATIONS,
    )
    key = base64.urlsafe_b64encode(kdf.derive(_APP_SECRET))
    return Fernet(key)


def _get_fernet() -> Fernet:
    global _fernet
    if _fernet is None:
        _fernet = _build_fernet()
    return _fernet


def reset_fernet():
    """Force key reload — call if the salt is ever rotated."""
    global _fernet
    _fernet = None


# ─── Public API ───────────────────────────────────────────────────────────────

def encrypt(plaintext: str) -> str:
    """
    Encrypt a plaintext string. Returns a UTF-8-safe ciphertext string
    prefixed with 'enc:' so we can detect encrypted values reliably.
    Returns the original string unchanged if plaintext is empty.
    """
    if not plaintext:
        return plaintext
    try:
        token = _get_fernet().encrypt(plaintext.encode("utf-8"))
        return "enc:" + token.decode("ascii")
    except Exception as e:
        logger.error(f"Credential encryption failed: {e}")
        return plaintext   # safe fallback — don't lose data


def decrypt(ciphertext: str) -> str:
    """
    Decrypt a ciphertext string produced by encrypt().
    - If ciphertext starts with 'enc:', decrypts with Fernet.
    - Otherwise returns the value as-is (plaintext legacy credential).
    Returns empty string on failure.
    """
    if not ciphertext:
        return ciphertext
    if not ciphertext.startswith("enc:"):
        return ciphertext
    try:
        token = ciphertext[4:].encode("ascii")
        return _get_fernet().decrypt(token).decode("utf-8")
    except InvalidToken:
        logger.error("Credential decryption failed - wrong key or corrupted data.")
        return ""
    except Exception as e:
        logger.error(f"Credential decryption error: {e}")
        return ""


def is_encrypted(value: str) -> bool:
    """Return True if the value was produced by encrypt()."""
    return isinstance(value, str) and value.startswith("enc:")