from __future__ import annotations import sqlite3 from datetime import datetime, timezone from typing import Any class AccountRepositoryMixin: """Original SQLite account persistence methods, moved without query changes.""" def count_users(self) -> int: with self.connect() as connection: row = connection.execute("SELECT COUNT(*) AS total FROM users").fetchone() return int(row["total"] if row else 0) def first_user_id(self) -> int: with self.connect() as connection: row = connection.execute("SELECT MIN(id) AS id FROM users").fetchone() return int(row["id"] or 0) if row else 0 def create_user( self, username: str, password_salt: str, password_hash: str, ) -> dict[str, Any]: now = datetime.now(timezone.utc).isoformat(timespec="seconds") try: with self.connect() as connection: role = "admin" if int(connection.execute("SELECT COUNT(*) FROM users").fetchone()[0]) == 0 else "user" cursor = connection.execute( """ INSERT INTO users (username, password_salt, password_hash, role, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?) """, (username, password_salt, password_hash, role, now, now), ) user_id = int(cursor.lastrowid) except sqlite3.IntegrityError as exc: raise ValueError("该账号名已被使用。") from exc return {"id": user_id, "username": username, "role": role, "created_at": now} def user_by_username(self, username: str) -> dict[str, Any] | None: with self.connect() as connection: row = connection.execute( """ SELECT id, username, password_salt, password_hash, role, llm_mode, membership_status, membership_plan, membership_starts_at, membership_expires_at, created_at FROM users WHERE username = ? COLLATE NOCASE """, (username,), ).fetchone() return dict(row) if row else None def user_password(self, user_id: int) -> dict[str, str] | None: with self.connect() as connection: row = connection.execute( "SELECT password_salt, password_hash FROM users WHERE id = ?", (user_id,), ).fetchone() return dict(row) if row else None def update_user_password(self, user_id: int, password_salt: str, password_hash: str) -> bool: now = datetime.now(timezone.utc).isoformat(timespec="seconds") with self.connect() as connection: cursor = connection.execute( "UPDATE users SET password_salt = ?, password_hash = ?, updated_at = ? WHERE id = ?", (password_salt, password_hash, now, user_id), ) return cursor.rowcount > 0 def delete_user(self, user_id: int) -> bool: with self.connect() as connection: cursor = connection.execute("DELETE FROM users WHERE id = ?", (user_id,)) return cursor.rowcount > 0 def create_session( self, session_hash: str, user_id: int, csrf_token: str, expires_at: str, ) -> None: now = datetime.now(timezone.utc).isoformat(timespec="seconds") with self.connect() as connection: connection.execute("DELETE FROM user_sessions WHERE expires_at <= ?", (now,)) connection.execute( """ INSERT INTO user_sessions (token_hash, user_id, csrf_token, expires_at, created_at, last_seen_at) VALUES (?, ?, ?, ?, ?, ?) """, (session_hash, user_id, csrf_token, expires_at, now, now), ) def session_user(self, session_hash: str) -> dict[str, Any] | None: now = datetime.now(timezone.utc).isoformat(timespec="seconds") with self.connect() as connection: row = connection.execute( """ SELECT u.id, u.username, u.role, u.llm_mode, u.membership_status, u.membership_plan, u.membership_starts_at, u.membership_expires_at, u.created_at, s.csrf_token, s.expires_at FROM user_sessions AS s JOIN users AS u ON u.id = s.user_id WHERE s.token_hash = ? AND s.expires_at > ? """, (session_hash, now), ).fetchone() if row: connection.execute( "UPDATE user_sessions SET last_seen_at = ? WHERE token_hash = ?", (now, session_hash), ) return dict(row) if row else None def delete_session(self, session_hash: str) -> bool: with self.connect() as connection: cursor = connection.execute( "DELETE FROM user_sessions WHERE token_hash = ?", (session_hash,), ) return cursor.rowcount > 0 def get_user_credentials(self, user_id: int) -> str: with self.connect() as connection: row = connection.execute( "SELECT encrypted_payload FROM user_credentials WHERE user_id = ?", (user_id,), ).fetchone() return str(row["encrypted_payload"]) if row else "" def save_user_credentials(self, user_id: int, encrypted_payload: str) -> None: now = datetime.now(timezone.utc).isoformat(timespec="seconds") with self.connect() as connection: connection.execute( """ INSERT INTO user_credentials (user_id, encrypted_payload, updated_at) VALUES (?, ?, ?) ON CONFLICT(user_id) DO UPDATE SET encrypted_payload = excluded.encrypted_payload, updated_at = excluded.updated_at """, (user_id, encrypted_payload, now), ) def list_user_credentials(self) -> list[dict[str, Any]]: with self.connect() as connection: rows = connection.execute( "SELECT user_id, encrypted_payload FROM user_credentials ORDER BY user_id" ).fetchall() return [dict(row) for row in rows] def user_access(self, user_id: int) -> dict[str, Any] | None: with self.connect() as connection: row = connection.execute( """ SELECT id, username, role, llm_mode, membership_status, membership_plan, membership_starts_at, membership_expires_at, created_at FROM users WHERE id = ? """, (user_id,), ).fetchone() return dict(row) if row else None def list_users(self) -> list[dict[str, Any]]: with self.connect() as connection: rows = connection.execute( """ SELECT id, username, role, llm_mode, membership_status, membership_plan, membership_starts_at, membership_expires_at, created_at FROM users ORDER BY id """ ).fetchall() return [dict(row) for row in rows] def update_user_llm_mode(self, user_id: int, mode: str) -> None: now = datetime.now(timezone.utc).isoformat(timespec="seconds") with self.connect() as connection: connection.execute( "UPDATE users SET llm_mode = ?, updated_at = ? WHERE id = ?", (mode, now, user_id), ) def update_membership( self, user_id: int, status: str, plan: str, starts_at: str | None, expires_at: str | None, ) -> bool: now = datetime.now(timezone.utc).isoformat(timespec="seconds") with self.connect() as connection: cursor = connection.execute( """ UPDATE users SET membership_status = ?, membership_plan = ?, membership_starts_at = ?, membership_expires_at = ?, updated_at = ? WHERE id = ? """, (status, plan, starts_at, expires_at, now, user_id), ) return cursor.rowcount > 0 def get_user_birth_profile(self, user_id: int) -> str: with self.connect() as connection: row = connection.execute( "SELECT encrypted_payload FROM user_birth_profiles WHERE user_id = ?", (user_id,), ).fetchone() return str(row["encrypted_payload"]) if row else "" def save_user_birth_profile(self, user_id: int, encrypted_payload: str) -> None: now = datetime.now(timezone.utc).isoformat(timespec="seconds") with self.connect() as connection: connection.execute( """ INSERT INTO user_birth_profiles (user_id, encrypted_payload, updated_at) VALUES (?, ?, ?) ON CONFLICT(user_id) DO UPDATE SET encrypted_payload = excluded.encrypted_payload, updated_at = excluded.updated_at """, (user_id, encrypted_payload, now), ) def delete_user_birth_profile(self, user_id: int) -> bool: with self.connect() as connection: cursor = connection.execute( "DELETE FROM user_birth_profiles WHERE user_id = ?", (user_id,), ) return cursor.rowcount > 0