ria-ran--sim-drop/server/database.py

import sqlite3
import logging
from datetime import datetime
from typing import Optional

from server.models import Subscriber

logger = logging.getLogger(__name__)

# status: 'claimed' (on-demand / already handed out) or 'available' (pre-registered
# batch, waiting to be claimed by a device).
_SCHEMA = """
CREATE TABLE IF NOT EXISTS subscribers (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    imsi TEXT UNIQUE NOT NULL,
    iccid TEXT UNIQUE NOT NULL,
    ki TEXT NOT NULL,
    opc TEXT NOT NULL,
    mac_address TEXT,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    synced_to_core BOOLEAN DEFAULT 0,
    status TEXT NOT NULL DEFAULT 'claimed'
);
CREATE INDEX IF NOT EXISTS idx_imsi ON subscribers(imsi);
CREATE INDEX IF NOT EXISTS idx_mac_address ON subscribers(mac_address);
"""


def init_db(db_path: str) -> None:
    """Create the database schema if it does not already exist (and migrate)."""
    conn = sqlite3.connect(db_path)
    try:
        conn.executescript(_SCHEMA)
        # Migrate pre-existing databases that lack the status column. This must
        # happen before creating the status index (a legacy table has neither).
        cols = {row[1] for row in conn.execute("PRAGMA table_info(subscribers)")}
        if "status" not in cols:
            conn.execute(
                "ALTER TABLE subscribers ADD COLUMN status TEXT NOT NULL DEFAULT 'claimed'"
            )
        conn.execute("CREATE INDEX IF NOT EXISTS idx_status ON subscribers(status)")
        conn.commit()
        logger.info(f"Database initialized at {db_path}")
    finally:
        conn.close()


def get_connection(db_path: str) -> sqlite3.Connection:
    """Return an open connection with Row factory enabled.

    Sets a busy timeout so that under a threaded server concurrent writers
    wait for the lock instead of immediately raising "database is locked".
    """
    conn = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES, timeout=5.0)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA busy_timeout = 5000")
    return conn


def _row_to_subscriber(row: sqlite3.Row) -> Subscriber:
    created = row["created_at"]
    if isinstance(created, str):
        created = datetime.fromisoformat(created)
    return Subscriber(
        id=row["id"],
        imsi=row["imsi"],
        iccid=row["iccid"],
        ki=row["ki"],
        opc=row["opc"],
        mac_address=row["mac_address"],
        created_at=created,
        synced_to_core=bool(row["synced_to_core"]),
        status=row["status"],
    )


def add_subscriber(
    conn: sqlite3.Connection,
    imsi: str,
    iccid: str,
    ki: str,
    opc: str,
    mac_address: Optional[str] = None,
) -> int:
    """Insert a new subscriber. Returns the new row id.

    Raises ValueError if imsi or iccid already exists.
    """
    try:
        cur = conn.execute(
            """
            INSERT INTO subscribers (imsi, iccid, ki, opc, mac_address)
            VALUES (?, ?, ?, ?, ?)
            """,
            (imsi, iccid, ki, opc, mac_address),
        )
        conn.commit()
        return cur.lastrowid
    except sqlite3.IntegrityError as exc:
        raise ValueError(f"Duplicate subscriber (imsi or iccid already exists): {exc}") from exc


def add_available_subscriber(
    conn: sqlite3.Connection, imsi: str, iccid: str, ki: str, opc: str,
) -> int:
    """Insert a pre-registered ('available') subscriber with no MAC assigned yet.

    Raises ValueError if imsi or iccid already exists.
    """
    try:
        cur = conn.execute(
            """
            INSERT INTO subscribers (imsi, iccid, ki, opc, mac_address, status)
            VALUES (?, ?, ?, ?, NULL, 'available')
            """,
            (imsi, iccid, ki, opc),
        )
        conn.commit()
        return cur.lastrowid
    except sqlite3.IntegrityError as exc:
        raise ValueError(f"Duplicate subscriber (imsi or iccid already exists): {exc}") from exc


def count_available(conn: sqlite3.Connection) -> int:
    """Return the number of unclaimed pre-registered subscribers."""
    return conn.execute(
        "SELECT COUNT(*) FROM subscribers WHERE status = 'available'"
    ).fetchone()[0]


def claim_available(conn: sqlite3.Connection, mac_address: str) -> Optional[Subscriber]:
    """Atomically claim the oldest available subscriber for ``mac_address``.

    Returns the claimed Subscriber, or None if none are available. Uses a write
    transaction so concurrent claims cannot hand out the same profile twice.
    """
    try:
        conn.execute("BEGIN IMMEDIATE")
        row = conn.execute(
            "SELECT id FROM subscribers WHERE status = 'available' ORDER BY id LIMIT 1"
        ).fetchone()
        if row is None:
            conn.commit()
            return None
        sub_id = row["id"]
        conn.execute(
            "UPDATE subscribers SET status = 'claimed', mac_address = ? WHERE id = ?",
            (mac_address, sub_id),
        )
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    return get_subscriber_by_id(conn, sub_id)


def get_subscriber_by_imsi(conn: sqlite3.Connection, imsi: str) -> Optional[Subscriber]:
    row = conn.execute("SELECT * FROM subscribers WHERE imsi = ?", (imsi,)).fetchone()
    return _row_to_subscriber(row) if row else None


def get_subscriber_by_id(conn: sqlite3.Connection, sub_id: int) -> Optional[Subscriber]:
    row = conn.execute("SELECT * FROM subscribers WHERE id = ?", (sub_id,)).fetchone()
    return _row_to_subscriber(row) if row else None


def list_unsynced(conn: sqlite3.Connection) -> list:
    rows = conn.execute(
        "SELECT * FROM subscribers WHERE synced_to_core = 0"
    ).fetchall()
    return [_row_to_subscriber(r) for r in rows]


def mark_synced(conn: sqlite3.Connection, imsi: str) -> None:
    conn.execute(
        "UPDATE subscribers SET synced_to_core = 1 WHERE imsi = ?", (imsi,)
    )
    conn.commit()


def check_rate_limit(
    conn: sqlite3.Connection, mac_address: str, window_hours: int
) -> bool:
    """Return True if the MAC address is allowed to provision (no record in the window).

    Uses SQLite's ``datetime()`` on both sides so the comparison is a true
    chronological one regardless of whether ``created_at`` was written by
    ``CURRENT_TIMESTAMP`` (``"YYYY-MM-DD HH:MM:SS"``) or as a Python
    ``isoformat()`` string (``"YYYY-MM-DDTHH:MM:SS.ffffff"``).
    """
    cur = conn.execute(
        "SELECT COUNT(*) FROM subscribers "
        "WHERE mac_address = ? AND datetime(created_at) > datetime('now', ?)",
        (mac_address, f"-{window_hours} hours"),
    )
    count = cur.fetchone()[0]
    return count == 0


def export_all(conn: sqlite3.Connection) -> list:
    """Return all subscriber rows as a list of plain dicts."""
    rows = conn.execute("SELECT * FROM subscribers").fetchall()
    result = []
    for row in rows:
        result.append({
            "id": row["id"],
            "imsi": row["imsi"],
            "iccid": row["iccid"],
            "ki": row["ki"],
            "opc": row["opc"],
            "mac_address": row["mac_address"],
            "created_at": str(row["created_at"]),
            "synced_to_core": bool(row["synced_to_core"]),
            "status": row["status"],
        })
    return result