Rhinox13/chatapi

1from __future__ import annotations
2
3import base64
4import binascii
5import hashlib
6import hmac
7import struct
8import time
9from functools import wraps
10from typing import Any, Callable
11
12from flask import jsonify, request, session
13
14
15class AuthContext:
16    def __init__(self, store: Any | None = None, user_store: Any | None = None):
17        self.store = store
18        self.user_store = user_store
19
20    def current_user(self) -> dict[str, str] | None:
21        user_id = str(session.get("user_id", "") or "").strip()
22        username = str(session.get("username", "") or "").strip()
23        role = str(session.get("role", "") or "").strip()
24        if not user_id or not username:
25            return None
26        if self.user_store is None:
27            return {"id": user_id, "username": username, "role": role}
28
29        db_user = self.user_store.get_user(user_id)
30        if db_user is None:
31            session.clear()
32            return None
33
34        if username != db_user.username or role != db_user.role:
35            session["username"] = db_user.username
36            session["role"] = db_user.role
37
38        return {"id": db_user.id, "username": db_user.username, "role": db_user.role}
39
40    def request_api_key(self) -> str:
41        return str(
42            request.headers.get("Authorization", "").removeprefix("Bearer ")
43            or request.headers.get("X-API-Key", "")
44            or ""
45        ).strip()
46
47    def resolve_owner_from_api_key(self) -> str | None:
48        api_key = self.request_api_key()
49        if not api_key or self.user_store is None:
50            return None
51        return self.user_store.resolve_api_key_owner(api_key)
52
53    def request_api_key_name(self) -> str:
54        api_key = self.request_api_key()
55        if not api_key or self.user_store is None:
56            return ""
57        return self.user_store.resolve_api_key_name(api_key) or ""
58
59    def is_request_authorized_by_api_key(self) -> bool:
60        return self.resolve_owner_from_api_key() is not None
61
62    def owner_id(self) -> str:
63        user = self.current_user()
64        if user:
65            return user["id"]
66        owner = self.resolve_owner_from_api_key()
67        if owner:
68            return owner
69        return "anonymous"
70
71    def is_admin(self) -> bool:
72        user = self.current_user()
73        return user is not None and user.get("role") == "admin"
74
75    def require_auth(self, view: Callable[..., Any]):
76        @wraps(view)
77        def wrapped(*args, **kwargs):
78            if self.current_user() is None and not self.is_request_authorized_by_api_key():
79                return jsonify({"error": "unauthorized"}), 401
80            return view(*args, **kwargs)
81        return wrapped
82
83    def require_session_auth(self, view: Callable[..., Any]):
84        @wraps(view)
85        def wrapped(*args, **kwargs):
86            if self.current_user() is None:
87                return jsonify({"error": "unauthorized"}), 401
88            return view(*args, **kwargs)
89
90        return wrapped
91
92    def require_admin(self, view: Callable[..., Any]):
93        @wraps(view)
94        def wrapped(*args, **kwargs):
95            if not self.is_admin():
96                return jsonify({"error": "forbidden"}), 403
97            return view(*args, **kwargs)
98        return wrapped
99
100    def request_headers_snapshot(self) -> dict[str, str]:
101        return {
102            "user_agent": str(request.headers.get("User-Agent", "")).strip(),
103            "content_type": str(request.headers.get("Content-Type", "")).strip(),
104            "origin": str(request.headers.get("Origin", "")).strip(),
105            "referer": str(request.headers.get("Referer", "")).strip(),
106        }
107
108
109def _normalize_totp_secret(secret: str) -> bytes:
110    normalized = "".join(secret.split()).upper()
111    if not normalized:
112        return b""
113    padding = "=" * ((8 - len(normalized) % 8) % 8)
114    try:
115        return base64.b32decode(f"{normalized}{padding}", casefold=True)
116    except (binascii.Error, ValueError):
117        return secret.encode("utf-8")
118
119
120def _totp_code(secret: bytes, counter: int, digits: int = 6) -> str:
121    digest = hmac.new(secret, struct.pack(">Q", counter), hashlib.sha1).digest()
122    offset = digest[-1] & 0x0F
123    binary = struct.unpack(">I", digest[offset : offset + 4])[0] & 0x7FFFFFFF
124    return str(binary % (10**digits)).zfill(digits)
125
126
127def generate_totp_secret() -> str:
128    raw = base64.b32encode(hashlib.sha256(time.time_ns().to_bytes(16, "big")).digest()[:20]).decode()
129    return raw.rstrip("=")
130
131
132def build_totp_uri(secret: str, username: str, issuer: str = "ChatAPI") -> str:
133    import urllib.parse
134    label = urllib.parse.quote(f"{issuer}:{username}")
135    issuer_encoded = urllib.parse.quote(issuer)
136    return f"otpauth://totp/{label}?secret={secret}&issuer={issuer_encoded}&algorithm=SHA1&digits=6&period=30"
137
138
139def verify_totp_code(
140    secret: str,
141    code: str,
142    *,
143    for_time: int | None = None,
144    step: int = 30,
145    window: int = 1,
146    digits: int = 6,
147) -> bool:
148    secret_bytes = _normalize_totp_secret(secret)
149    if not secret_bytes:
150        return False
151
152    candidate = "".join(str(code).split())
153    if not candidate.isdigit() or len(candidate) != digits:
154        return False
155
156    now = int(time.time() if for_time is None else for_time)
157    counter = now // step
158    for drift in range(-window, window + 1):
159        if hmac.compare_digest(_totp_code(secret_bytes, counter + drift, digits=digits), candidate):
160            return True
161    return False
162