feat: open source multi-user AI notebook with observability
This commit is contained in:
wuyang
+204
@@ -0,0 +1,204 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import secrets
|
||||
import time
|
||||
from collections import defaultdict, deque
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from argon2 import PasswordHasher
|
||||
from argon2.exceptions import InvalidHashError, VerifyMismatchError
|
||||
from fastapi import HTTPException, Request, Response, status
|
||||
|
||||
from .config import Settings
|
||||
from .db import Database, utc_now
|
||||
|
||||
COOKIE_NAME = "xuxiang_session"
|
||||
SESSION_DAYS = 30
|
||||
_password_hasher = PasswordHasher()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AuthUser:
|
||||
id: str
|
||||
label: str
|
||||
role: str
|
||||
debug_sharing: bool
|
||||
|
||||
@property
|
||||
def is_admin(self) -> bool:
|
||||
return self.role == "admin"
|
||||
|
||||
def public(self) -> dict[str, object]:
|
||||
return {
|
||||
"id": self.id,
|
||||
"label": self.label,
|
||||
"role": self.role,
|
||||
"debug_sharing": self.debug_sharing,
|
||||
}
|
||||
|
||||
|
||||
class LoginLimiter:
|
||||
def __init__(self) -> None:
|
||||
self.attempts: dict[str, deque[float]] = defaultdict(deque)
|
||||
|
||||
def check(self, key: str) -> None:
|
||||
now = time.monotonic()
|
||||
attempts = self.attempts[key]
|
||||
while attempts and now - attempts[0] > 900:
|
||||
attempts.popleft()
|
||||
if len(attempts) >= 8:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
detail="请稍后再试",
|
||||
)
|
||||
|
||||
def fail(self, key: str) -> None:
|
||||
self.attempts[key].append(time.monotonic())
|
||||
|
||||
def clear(self, key: str) -> None:
|
||||
self.attempts.pop(key, None)
|
||||
|
||||
|
||||
limiter = LoginLimiter()
|
||||
|
||||
|
||||
def find_user_for_key(access_key: str, db: Database) -> AuthUser | None:
|
||||
for candidate in db.list_auth_users():
|
||||
try:
|
||||
matches = _password_hasher.verify(
|
||||
candidate["access_key_hash"], access_key
|
||||
)
|
||||
except (VerifyMismatchError, InvalidHashError):
|
||||
matches = False
|
||||
if matches:
|
||||
return AuthUser(
|
||||
id=candidate["id"],
|
||||
label=candidate["label"],
|
||||
role=candidate["role"],
|
||||
debug_sharing=bool(candidate["debug_sharing"]),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _token_hash(token: str, settings: Settings) -> str:
|
||||
secret = (settings.session_secret or "").encode()
|
||||
return hmac.new(secret, token.encode(), hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
def create_session(
|
||||
response: Response,
|
||||
db: Database,
|
||||
settings: Settings,
|
||||
user_id: str,
|
||||
) -> None:
|
||||
token = secrets.token_urlsafe(40)
|
||||
expires = datetime.now(UTC) + timedelta(days=SESSION_DAYS)
|
||||
with db.connect() as connection:
|
||||
connection.execute(
|
||||
"DELETE FROM sessions WHERE expires_at < ?", (utc_now(),)
|
||||
)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO sessions (
|
||||
token_hash, user_id, created_at, expires_at
|
||||
) VALUES (?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
_token_hash(token, settings),
|
||||
user_id,
|
||||
utc_now(),
|
||||
expires.isoformat(),
|
||||
),
|
||||
)
|
||||
response.set_cookie(
|
||||
COOKIE_NAME,
|
||||
token,
|
||||
max_age=SESSION_DAYS * 24 * 60 * 60,
|
||||
httponly=True,
|
||||
secure=settings.cookie_secure,
|
||||
samesite="strict",
|
||||
path="/",
|
||||
)
|
||||
|
||||
|
||||
def delete_session(
|
||||
request: Request, response: Response, db: Database, settings: Settings
|
||||
) -> None:
|
||||
token = request.cookies.get(COOKIE_NAME)
|
||||
if token:
|
||||
with db.connect() as connection:
|
||||
connection.execute(
|
||||
"DELETE FROM sessions WHERE token_hash = ?",
|
||||
(_token_hash(token, settings),),
|
||||
)
|
||||
response.delete_cookie(
|
||||
COOKIE_NAME,
|
||||
path="/",
|
||||
secure=settings.cookie_secure,
|
||||
httponly=True,
|
||||
samesite="strict",
|
||||
)
|
||||
|
||||
|
||||
def require_auth(
|
||||
request: Request, db: Database, settings: Settings
|
||||
) -> AuthUser:
|
||||
if settings.auth_disabled:
|
||||
users = db.list_auth_users()
|
||||
if users:
|
||||
user = users[0]
|
||||
return AuthUser(
|
||||
id=user["id"],
|
||||
label=user["label"],
|
||||
role="admin",
|
||||
debug_sharing=True,
|
||||
)
|
||||
return AuthUser(
|
||||
id="development",
|
||||
label="本地开发",
|
||||
role="admin",
|
||||
debug_sharing=True,
|
||||
)
|
||||
if not settings.users or not settings.session_secret:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="应用尚未配置访问身份",
|
||||
)
|
||||
token = request.cookies.get(COOKIE_NAME)
|
||||
if not token:
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
|
||||
with db.connect() as connection:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT u.id, u.label, u.role, u.debug_sharing, s.expires_at
|
||||
FROM sessions s
|
||||
JOIN users u ON u.id = s.user_id
|
||||
WHERE s.token_hash = ?
|
||||
""",
|
||||
(_token_hash(token, settings),),
|
||||
).fetchone()
|
||||
if not row or row["expires_at"] <= utc_now():
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
|
||||
return AuthUser(
|
||||
id=row["id"],
|
||||
label=row["label"],
|
||||
role=row["role"],
|
||||
debug_sharing=bool(row["debug_sharing"]),
|
||||
)
|
||||
|
||||
|
||||
def require_admin(user: AuthUser) -> None:
|
||||
if not user.is_admin:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
|
||||
|
||||
def require_same_origin_intent(request: Request) -> None:
|
||||
if request.method in {"GET", "HEAD", "OPTIONS"}:
|
||||
return
|
||||
if request.url.path == "/api/login":
|
||||
return
|
||||
if request.headers.get("x-note-client") != "xuxiang-web":
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
Reference in New Issue
Block a user