237 lines
9.5 KiB
Python
237 lines
9.5 KiB
Python
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
|