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__)

_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
);
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."""
    conn = sqlite3.connect(db_path)
    try:
        conn.executescript(_SCHEMA)
        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"]),
    )


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 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"]),
        })
    return result