gemstone-utils 0.4.0__py3-none-any.whl

gemstone_utils/crypto.py

@@ -0,0 +1,188 @@
+# SPDX-License-Identifier: MPL-2.0
+# Copyright 2026, Clinton Bunch. All rights reserved.
+# gemstone_utils/crypto.py
+
+from __future__ import annotations
+
+import os
+from base64 import urlsafe_b64encode, urlsafe_b64decode
+from typing import Any, Callable, Dict, Final, Mapping, NamedTuple, Optional, Tuple
+
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+
+# ---- KEK derivation ---------------------------------------------------------
+
+# OWASP-style order of magnitude for PBKDF2-HMAC-SHA256 when params omit iterations
+# (used by key_mgmt persisted KDF defaults).
+DEFAULT_PBKDF2_ITERATIONS_STRONG = 600_000
+
+
+def derive_pbkdf2_hmac_sha256(
+    passphrase: str,
+    salt: bytes,
+    *,
+    iterations: int,
+    length: int = 32,
+) -> bytes:
+    """
+    Derive key bytes using cryptography's PBKDF2HMAC (SHA-256).
+    """
+    if not isinstance(passphrase, str):
+        raise TypeError("passphrase must be a str")
+    if not isinstance(salt, (bytes, bytearray)):
+        raise TypeError("salt must be bytes")
+
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=length,
+        salt=salt,
+        iterations=iterations,
+    )
+    return kdf.derive(passphrase.encode("utf-8"))
+
+
+# ---- AES-GCM primitives -----------------------------------------------------
+
+
+def aesgcm_encrypt(dk: bytes, plaintext: bytes, aad: Optional[bytes] = None) -> bytes:
+    aesgcm = AESGCM(dk)
+    nonce = os.urandom(12)
+    ciphertext = aesgcm.encrypt(nonce, plaintext, aad)
+    return nonce + ciphertext
+
+
+def aesgcm_decrypt(dk: bytes, blob: bytes, aad: Optional[bytes] = None) -> bytes:
+    if len(blob) < 12 + 16:
+        raise ValueError("ciphertext blob too short")
+
+    nonce = blob[:12]
+    ciphertext = blob[12:]
+    aesgcm = AESGCM(dk)
+    return aesgcm.decrypt(nonce, ciphertext, aad)
+
+
+# ---- Symmetric algorithm registry -------------------------------------------
+
+
+def _a256_validate_sym_params(params: Dict[str, Any]) -> None:
+    if params:
+        raise ValueError(f"A256GCM does not accept algorithm parameters (got {params!r})")
+
+
+def _a256_encrypt_impl(
+    key: bytes, plaintext: bytes, _params: Dict[str, Any]
+) -> Tuple[bytes, Dict[str, Any]]:
+    return aesgcm_encrypt(key, plaintext), {}
+
+
+def _a256_decrypt_impl(key: bytes, ciphertext: bytes, _params: Dict[str, Any]) -> bytes:
+    return aesgcm_decrypt(key, ciphertext)
+
+
+class SymAlgSpec(NamedTuple):
+    """Registered symmetric algorithm: key size, param validation, and crypto ops."""
+
+    key_length: int
+    validate_sym_params: Callable[[Dict[str, Any]], None]
+    encrypt_impl: Callable[[bytes, bytes, Dict[str, Any]], Tuple[bytes, Dict[str, Any]]]
+    decrypt_impl: Callable[[bytes, bytes, Dict[str, Any]], bytes]
+
+
+SYM_ALG_REGISTRY: Dict[str, SymAlgSpec] = {
+    "A256GCM": SymAlgSpec(
+        key_length=32,
+        validate_sym_params=_a256_validate_sym_params,
+        encrypt_impl=_a256_encrypt_impl,
+        decrypt_impl=_a256_decrypt_impl,
+    ),
+}
+
+RECOMMENDED_DATA_ALG: Final[str] = "A256GCM"
+assert RECOMMENDED_DATA_ALG in SYM_ALG_REGISTRY
+
+SUPPORTED_SYM_ALGS: frozenset[str] = frozenset(SYM_ALG_REGISTRY.keys())
+
+
+def is_supported_sym_alg(alg: str) -> bool:
+    return alg in SYM_ALG_REGISTRY
+
+
+def sym_alg_key_length(alg: str) -> int:
+    spec = SYM_ALG_REGISTRY.get(alg)
+    if spec is None:
+        raise ValueError(f"Unsupported symmetric alg: {alg}")
+    return spec.key_length
+
+
+def recommended_data_alg() -> str:
+    """
+    Symmetric algorithm id recommended for **new** application field encryption.
+
+    Matches the default for :attr:`~gemstone_utils.types.KeyContext.alg` and
+    persisted :attr:`~gemstone_utils.sqlalchemy.key_storage.GemstoneKeyRecord.data_alg`.
+    """
+    return RECOMMENDED_DATA_ALG
+
+
+def generate_key_by_alg(alg: str) -> bytes:
+    """Return ``os.urandom(key_length)`` for the registered algorithm."""
+    return os.urandom(sym_alg_key_length(alg))
+
+
+def encrypt_alg(
+    alg: str,
+    key: bytes,
+    plaintext: bytes,
+    params: Optional[Mapping[str, Any]] = None,
+) -> Tuple[bytes, Dict[str, Any]]:
+    """
+    Encrypt with a registered symmetric algorithm.
+
+    Returns ``(ciphertext, updated_params)``. Callers persist ``updated_params``
+    in the wire JSON segment when nonces or metadata are stored outside the blob.
+    """
+    spec = SYM_ALG_REGISTRY.get(alg)
+    if spec is None:
+        raise ValueError(f"Unsupported symmetric alg: {alg}")
+    p = dict(params) if params is not None else {}
+    spec.validate_sym_params(p)
+    return spec.encrypt_impl(key, plaintext, p)
+
+
+def decrypt_alg(
+    alg: str,
+    key: bytes,
+    ciphertext: bytes,
+    params: Optional[Mapping[str, Any]] = None,
+) -> bytes:
+    spec = SYM_ALG_REGISTRY.get(alg)
+    if spec is None:
+        raise ValueError(f"Unsupported symmetric alg: {alg}")
+    p = dict(params) if params is not None else {}
+    spec.validate_sym_params(p)
+    return spec.decrypt_impl(key, ciphertext, p)
+
+
+def encrypt_with_alg(alg: str, key: bytes, plaintext: bytes) -> bytes:
+    """Backward-compatible: same as ``encrypt_alg`` but returns ciphertext only."""
+    blob, _params = encrypt_alg(alg, key, plaintext, None)
+    return blob
+
+
+def decrypt_with_alg(alg: str, key: bytes, blob: bytes) -> bytes:
+    """Backward-compatible: decrypt with empty symmetric parameters."""
+    return decrypt_alg(alg, key, blob, None)
+
+
+# ---- Base64 helpers ---------------------------------------------------------
+
+
+def b64encode(data: bytes) -> str:
+    return urlsafe_b64encode(data).decode("ascii")
+
+
+def b64decode(data: str) -> bytes:
+    return urlsafe_b64decode(data.encode("ascii"))

gemstone_utils/db.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: MPL-2.0
+# Copyright 2026,
+# gemstone_utils/db.py
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from sqlalchemy import create_engine, event
+from sqlalchemy.engine import Engine, make_url
+from sqlalchemy.engine.url import URL
+from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
+
+
+class GemstoneDB(DeclarativeBase):
+    """Shared declarative base for gemstone_utils ORM models."""
+
+
+_engine: Optional[Engine] = None
+_session_factory: Optional[sessionmaker[Session]] = None
+
+
+def _is_sqlite(drivername: str) -> bool:
+    return drivername == "sqlite" or drivername.startswith("sqlite+")
+
+
+def _is_mysql_family(drivername: str) -> bool:
+    return drivername.startswith("mysql") or drivername.startswith("mariadb")
+
+
+def _is_postgresql(drivername: str) -> bool:
+    return drivername.startswith("postgresql")
+
+
+def _apply_dialect_engine_kwargs(url: URL, engine_kw: dict[str, Any]) -> URL:
+    """
+    Apply backend-specific defaults. Caller ``**engine_kw`` values win over
+    defaults (via setdefault / merged connect_args).
+    """
+    kw = engine_kw
+    driver = url.drivername
+
+    if _is_mysql_family(driver):
+        kw.setdefault("pool_pre_ping", True)
+        kw.setdefault("pool_recycle", 3600)
+        if "charset" not in url.query:
+            url = url.update_query_dict({"charset": "utf8mb4"})
+
+    elif _is_postgresql(driver):
+        kw.setdefault("pool_pre_ping", True)
+        kw.setdefault("pool_recycle", 3600)
+        merged_connect: dict[str, Any] = {}
+        merged_connect.setdefault("options", "-c timezone=UTC")
+        user_connect = kw.get("connect_args")
+        if user_connect:
+            merged_connect.update(user_connect)
+        kw["connect_args"] = merged_connect
+
+    return url
+
+
+def _register_sqlite_pragmas(engine: Engine) -> None:
+    @event.listens_for(engine, "connect")
+    def _on_sqlite_connect(dbapi_conn, connection_record) -> None:
+        cursor = dbapi_conn.cursor()
+        try:
+            cursor.execute("PRAGMA journal_mode=WAL")
+            cursor.execute("PRAGMA foreign_keys=ON")
+            cursor.execute("PRAGMA busy_timeout=5000")
+        finally:
+            cursor.close()
+
+
+def init_db(db_url: str, *, echo: bool = False, **engine_kw: Any) -> Engine:
+    """
+    Configure the process-global SQLAlchemy engine and session factory, then
+    create any missing tables for all models registered on
+    :attr:`GemstoneDB.metadata` (call after every plugin/module that defines
+    ``GemstoneDB`` subclasses has been imported).
+
+    Applies light dialect-specific defaults (SQLite WAL and pragmas; MySQL /
+    MariaDB utf8mb4 + pool tuning; PostgreSQL UTC session timezone + pool
+    tuning). Pass ``**engine_kw`` to override or extend :func:`create_engine`
+    arguments.
+
+    Returns the new :class:`~sqlalchemy.engine.Engine`.
+    """
+    global _engine, _session_factory
+
+    url = make_url(db_url)
+    kw = dict(engine_kw)
+    url = _apply_dialect_engine_kwargs(url, kw)
+
+    _engine = create_engine(url, echo=echo, **kw)
+
+    if _is_sqlite(url.drivername):
+        _register_sqlite_pragmas(_engine)
+
+    _session_factory = sessionmaker(
+        bind=_engine,
+        autoflush=True,
+        autocommit=False,
+        expire_on_commit=False,
+        class_=Session,
+    )
+    GemstoneDB.metadata.create_all(bind=_engine)
+    return _engine
+
+
+def get_session() -> Session:
+    """
+    Return a new :class:`~sqlalchemy.orm.Session` bound to the engine from
+    :func:`init_db`. The caller should close the session when done (or use it
+    as a context manager: ``with get_session() as session:``).
+    """
+    if _session_factory is None:
+        raise RuntimeError("init_db(...) must be called before get_session()")
+
+    return _session_factory()

gemstone_utils/election.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: MPL-2.0
+# Copyright 2026,
+# gemstone_utils/election.py
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Optional, Iterator
+from uuid import UUID
+
+from sqlalchemy import DateTime, String, delete, select
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm import Mapped, Session, mapped_column
+
+from .db import GemstoneDB, get_session
+
+
+def _utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def _ns(ns: Optional[str]) -> str:
+    return ns or "default"
+
+
+_expire_seconds: int = 60
+
+
+def set_expire(sec: int) -> None:
+    """
+    Set the candidate and leader lease expiration window (seconds).
+
+    The application should call this once at startup.
+    """
+    global _expire_seconds
+    if not isinstance(sec, int) or sec <= 0:
+        raise ValueError("expire seconds must be a positive int")
+    _expire_seconds = sec
+
+
+class ElectionCandidate(GemstoneDB):
+    __tablename__ = "gemstone_election_candidate"
+
+    ns: Mapped[str] = mapped_column(String(255), primary_key=True)
+    candidate_id: Mapped[str] = mapped_column(String(36), primary_key=True)
+
+    last_heartbeat_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
+    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
+
+
+class ElectionLeader(GemstoneDB):
+    __tablename__ = "gemstone_election_leader"
+
+    ns: Mapped[str] = mapped_column(String(255), primary_key=True)
+
+    leader_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)
+    lease_expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)
+    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
+
+
+@dataclass(frozen=True)
+class _Lease:
+    now: datetime
+    expires_at: datetime
+
+
+def _lease() -> _Lease:
+    now = _utcnow()
+    return _Lease(now=now, expires_at=now + timedelta(seconds=_expire_seconds))
+
+
+@contextmanager
+def _session_scope(session: Optional[Session]) -> Iterator[Session]:
+    if session is not None:
+        yield session
+        return
+    s = get_session()
+    try:
+        yield s
+    finally:
+        s.close()
+
+
+def register_candidate(candidate_id: UUID, ns: Optional[str] = None, *, session: Optional[Session] = None) -> None:
+    """
+    Register or refresh a candidate in the election namespace.
+    """
+    n = _ns(ns)
+    cid = str(candidate_id)
+    lease = _lease()
+
+    with _session_scope(session) as s:
+        with s.begin():
+            row = s.get(ElectionCandidate, {"ns": n, "candidate_id": cid})
+            if row is None:
+                s.add(
+                    ElectionCandidate(
+                        ns=n,
+                        candidate_id=cid,
+                        last_heartbeat_at=lease.now,
+                        expires_at=lease.expires_at,
+                    )
+                )
+            else:
+                row.last_heartbeat_at = lease.now
+                row.expires_at = lease.expires_at
+
+
+def heartbeat(candidate_id: UUID, ns: Optional[str] = None, *, session: Optional[Session] = None) -> None:
+    """
+    Refresh a candidate heartbeat and extend its expiry window.
+
+    If the candidate is missing, this behaves like register_candidate().
+    """
+    register_candidate(candidate_id, ns, session=session)
+
+
+def unregister_candidate(candidate_id: UUID, ns: Optional[str] = None, *, session: Optional[Session] = None) -> None:
+    """
+    Remove a candidate from the registry. If it is currently leader, clear the
+    leader row for faster failover (best-effort).
+    """
+    n = _ns(ns)
+    cid = str(candidate_id)
+    now = _utcnow()
+
+    with _session_scope(session) as s:
+        with s.begin():
+            s.execute(delete(ElectionCandidate).where(ElectionCandidate.ns == n, ElectionCandidate.candidate_id == cid))
+
+            leader = s.get(ElectionLeader, n)
+            if leader and leader.leader_id == cid:
+                leader.leader_id = None
+                leader.lease_expires_at = None
+                leader.updated_at = now
+
+
+def list_candidates(ns: Optional[str] = None, *, session: Optional[Session] = None) -> list[UUID]:
+    """
+    Return the currently-active candidates in the namespace (expires_at > now).
+    """
+    n = _ns(ns)
+    now = _utcnow()
+
+    with _session_scope(session) as s:
+        rows = s.execute(
+            select(ElectionCandidate.candidate_id)
+            .where(ElectionCandidate.ns == n, ElectionCandidate.expires_at > now)
+            .order_by(ElectionCandidate.candidate_id)
+        ).scalars()
+        return [UUID(x) for x in rows]
+
+
+def is_leader(candidate_id: UUID, ns: Optional[str] = None, *, session: Optional[Session] = None) -> bool:
+    """
+    Return True if candidate_id holds an unexpired lease for the namespace.
+    """
+    n = _ns(ns)
+    cid = str(candidate_id)
+    now = _utcnow()
+
+    with _session_scope(session) as s:
+        leader = s.get(ElectionLeader, n)
+        if leader is None or leader.leader_id is None or leader.lease_expires_at is None:
+            return False
+        return leader.leader_id == cid and leader.lease_expires_at > now
+
+
+def elect(candidate_id: UUID, ns: Optional[str] = None, *, session: Optional[Session] = None) -> UUID:
+    """
+    Attempt to acquire (or renew) leadership for candidate_id.
+
+    Returns the UUID of the current leader for the namespace after this call.
+    """
+    n = _ns(ns)
+    cid = str(candidate_id)
+    with _session_scope(session) as s:
+        # Contention-safe: if two candidates attempt to create the leader row at
+        # the same time, one may hit IntegrityError on insert/flush. Retry once
+        # with a fresh transaction.
+        for _ in range(2):
+            lease = _lease()
+            try:
+                with s.begin():
+                    # Ensure the candidate is present/active.
+                    row = s.get(ElectionCandidate, {"ns": n, "candidate_id": cid})
+                    if row is None:
+                        s.add(
+                            ElectionCandidate(
+                                ns=n,
+                                candidate_id=cid,
+                                last_heartbeat_at=lease.now,
+                                expires_at=lease.expires_at,
+                            )
+                        )
+                    else:
+                        row.last_heartbeat_at = lease.now
+                        row.expires_at = lease.expires_at
+
+                    # Acquire a row lock on the leader entry if supported.
+                    try:
+                        leader = s.execute(
+                            select(ElectionLeader).where(ElectionLeader.ns == n).with_for_update()
+                        ).scalar_one_or_none()
+                    except Exception:
+                        # SQLite and some drivers may not support FOR UPDATE cleanly.
+                        leader = s.get(ElectionLeader, n)
+
+                    if leader is None:
+                        leader = ElectionLeader(
+                            ns=n,
+                            leader_id=cid,
+                            lease_expires_at=lease.expires_at,
+                            updated_at=lease.now,
+                        )
+                        s.add(leader)
+                        s.flush()  # may raise IntegrityError under contention
+
+                    # Renew if we are leader, or take over if expired/cleared.
+                    if leader.leader_id is None or leader.lease_expires_at is None:
+                        leader.leader_id = cid
+                        leader.lease_expires_at = lease.expires_at
+                        leader.updated_at = lease.now
+                    elif leader.leader_id == cid:
+                        leader.lease_expires_at = lease.expires_at
+                        leader.updated_at = lease.now
+                    elif leader.lease_expires_at <= lease.now:
+                        leader.leader_id = cid
+                        leader.lease_expires_at = lease.expires_at
+                        leader.updated_at = lease.now
+
+                    # Ensure a usable leader_id is returned.
+                    if leader.leader_id is None:
+                        leader.leader_id = cid
+                        leader.lease_expires_at = lease.expires_at
+                        leader.updated_at = lease.now
+
+                    return UUID(leader.leader_id)
+            except IntegrityError:
+                # Auto-rollback occurs; retry once.
+                continue
+
+        # If we still can't elect due to persistent contention, fall back to
+        # reading current leader without attempting insert again.
+        leader = s.get(ElectionLeader, n)
+        if leader and leader.leader_id and leader.lease_expires_at and leader.lease_expires_at > _utcnow():
+            return UUID(leader.leader_id)
+        return UUID(cid)
+

gemstone_utils/encrypted_fields.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: MPL-2.0
+# Copyright 2026, Clinton Bunch. All rights reserved.
+# gemstone_utils/encrypted_fields.py
+
+from __future__ import annotations
+
+import json
+import warnings
+from typing import Any, Dict, Optional, Tuple
+
+from .crypto import (
+    SUPPORTED_SYM_ALGS,
+    SYM_ALG_REGISTRY,
+    b64decode,
+    b64encode,
+    decrypt_alg,
+    encrypt_alg,
+    is_supported_sym_alg,
+)
+from .key_id import normalize_key_id
+from .types import KeyContext
+
+
+def is_encrypted_prefix(value: str) -> bool:
+    if not isinstance(value, str) or not value.startswith("$"):
+        return False
+    parts = value.split("$")
+    if len(parts) < 4 or parts[0] != "":
+        return False
+    return parts[1] in SUPPORTED_SYM_ALGS
+
+
+def _params_json_bytes(params: Dict[str, Any]) -> bytes:
+    return json.dumps(params, separators=(",", ":"), sort_keys=True).encode("utf-8")
+
+
+def _encode_params_segment(params: Dict[str, Any]) -> str:
+    return b64encode(_params_json_bytes(params))
+
+
+def _decode_params_segment(segment: str) -> Dict[str, Any]:
+    raw = b64decode(segment)
+    try:
+        data = json.loads(raw.decode("utf-8"))
+    except json.JSONDecodeError as e:
+        raise ValueError("invalid algorithm parameters (not valid JSON)") from e
+    if not isinstance(data, dict):
+        raise ValueError("algorithm parameters must be a JSON object")
+    return data
+
+
+def _parse_key_id_segment(seg: str) -> str:
+    if seg.isdigit():
+        raise ValueError(
+            "legacy integer key id in encrypted field; migrate ciphertext to UUID "
+            "key ids (see gemstone_utils migration docs) before using this version"
+        )
+    return normalize_key_id(seg)
+
+
+def format_encrypted_field(
+    alg: str,
+    keyid: str,
+    blob: bytes,
+    params: Optional[Dict[str, Any]] = None,
+) -> str:
+    if not is_supported_sym_alg(alg):
+        raise ValueError(f"Unsupported symmetric alg: {alg}")
+    p = {} if params is None else params
+    kid = normalize_key_id(keyid)
+    return f"${alg}${kid}${_encode_params_segment(p)}${b64encode(blob)}"
+
+
+def parse_encrypted_field(value: str) -> Tuple[str, str, Dict[str, Any], bytes]:
+    parts = value.split("$")
+    if parts[0] != "":
+        raise ValueError("invalid encrypted field format")
+
+    if len(parts) == 4:
+        warnings.warn(
+            "Four-part encrypted fields (no algorithm-parameters segment) are deprecated "
+            "and will be removed in gemstone_utils 0.9.0; re-encrypt or run a key rotation "
+            "with gemstone_utils >= 0.3.0 to migrate.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        alg_id = parts[1]
+        keyid = _parse_key_id_segment(parts[2])
+        blob = b64decode(parts[3])
+        return alg_id, keyid, {}, blob
+
+    if len(parts) == 5:
+        alg_id = parts[1]
+        keyid = _parse_key_id_segment(parts[2])
+        params = _decode_params_segment(parts[3])
+        blob = b64decode(parts[4])
+        return alg_id, keyid, params, blob
+
+    raise ValueError("invalid encrypted field format")
+
+
+def _validate_alg_params(alg_id: str, params: Dict[str, Any]) -> None:
+    spec = SYM_ALG_REGISTRY.get(alg_id)
+    if spec is None:
+        raise ValueError(f"unsupported algorithm: {alg_id}")
+    spec.validate_sym_params(params)
+
+
+def encrypt_string(plaintext: Optional[str], keyctx: KeyContext) -> Optional[str]:
+    if plaintext is None:
+        return None
+    blob, out_params = encrypt_alg(
+        keyctx.alg, keyctx.key, plaintext.encode("utf-8"), None
+    )
+    return format_encrypted_field(keyctx.alg, keyctx.keyid, blob, out_params)
+
+
+def decrypt_string(value: Optional[str], keyctx: KeyContext) -> Optional[str]:
+    if value is None:
+        return None
+
+    alg_id, keyid, params, blob = parse_encrypted_field(value)
+
+    if alg_id != keyctx.alg:
+        raise ValueError(f"unsupported algorithm: {alg_id}")
+    if keyid != keyctx.keyid:
+        raise ValueError(f"unexpected keyid {keyid}")
+
+    _validate_alg_params(alg_id, params)
+
+    return decrypt_alg(keyctx.alg, keyctx.key, blob, params).decode("utf-8")