245 lines
9.0 KiB
Python
245 lines
9.0 KiB
Python
import hashlib
|
|
import secrets
|
|
import time
|
|
|
|
from fastapi import Depends, HTTPException, Request
|
|
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
|
from jose import JWTError, jwt
|
|
from passlib.context import CryptContext
|
|
|
|
from app import database
|
|
from app.config import get_settings
|
|
from app.models import ApiToken, Setting, User
|
|
|
|
pwd = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
|
_bearer = HTTPBearer(auto_error=False)
|
|
_secret: str | None = None
|
|
TOKEN_DAYS = 30
|
|
|
|
|
|
def secret_key() -> str:
|
|
global _secret
|
|
if _secret:
|
|
return _secret
|
|
if get_settings().secret_key:
|
|
_secret = get_settings().secret_key
|
|
return _secret
|
|
with database.session() as db:
|
|
row = db.get(Setting, "secret_key")
|
|
if row is None:
|
|
row = Setting(key="secret_key", value=secrets.token_urlsafe(48))
|
|
db.add(row)
|
|
db.commit()
|
|
_secret = row.value
|
|
return _secret
|
|
|
|
|
|
def ensure_admin():
|
|
"""设置了 ADMIN_PASSWORD 且还没有账号时，创建 admin。"""
|
|
pw = get_settings().admin_password
|
|
with database.session() as db:
|
|
if pw and db.query(User).count() == 0:
|
|
db.add(User(username="admin", password_hash=pwd.hash(pw), role="admin", created=int(time.time())))
|
|
db.commit()
|
|
print("[auth] 已按 ADMIN_PASSWORD 创建管理员 admin")
|
|
|
|
|
|
def make_token(username: str, version: int | None = None) -> str:
|
|
return jwt.encode({"sub": username, "ver": version or 0,
|
|
"exp": int(time.time()) + TOKEN_DAYS * 86400}, secret_key(), "HS256")
|
|
|
|
|
|
# 普通用户（role=user）只读：除了改自己的密码，所有写操作都要管理员；这些读接口也只给管理员
|
|
USER_WRITABLE = {"/api/auth/password"}
|
|
ADMIN_ONLY_GET = ("/api/settings", "/api/audit/notify", "/api/audit/oplog", "/api/auth/users", "/api/auth/tokens", "/api/compliance")
|
|
API_PREFIX = "rdp_" # 只读 API 令牌的前缀；登录令牌是 JWT（eyJ 开头），两者不会混淆
|
|
_api_cache: dict[str, tuple[str | None, float]] = {} # 令牌哈希 -> (名称，不存在为 None, 缓存时间)
|
|
_api_used: dict[int, float] = {} # 令牌 id -> 上次写 last_used 的时间（5 分钟写一次）
|
|
|
|
|
|
def new_api_token() -> tuple[str, str]:
|
|
"""(令牌原文, SHA-256)"""
|
|
t = API_PREFIX + secrets.token_urlsafe(32)
|
|
return t, hashlib.sha256(t.encode()).hexdigest()
|
|
|
|
|
|
def forget_api_tokens():
|
|
_api_cache.clear()
|
|
|
|
|
|
def _api_token_name(token: str) -> str | None:
|
|
h = hashlib.sha256(token.encode()).hexdigest()
|
|
hit = _api_cache.get(h)
|
|
if hit and time.time() - hit[1] < 30:
|
|
return hit[0]
|
|
with database.session() as db:
|
|
row = db.query(ApiToken).filter(ApiToken.token_hash == h).first()
|
|
name = row.name if row else None
|
|
if row and time.time() - _api_used.get(row.id, 0) > 300:
|
|
_api_used[row.id] = time.time()
|
|
row.last_used = int(time.time())
|
|
db.commit()
|
|
if len(_api_cache) > 1000:
|
|
_api_cache.clear()
|
|
_api_cache[h] = (name, time.time())
|
|
return name
|
|
|
|
|
|
def _read_only(request: Request):
|
|
path = request.url.path
|
|
if request.method not in ("GET", "HEAD") and path not in USER_WRITABLE:
|
|
raise HTTPException(403, "只有管理员可以修改")
|
|
if request.method in ("GET", "HEAD") and path.startswith(ADMIN_ONLY_GET):
|
|
raise HTTPException(403, "只有管理员可以查看")
|
|
_roles: dict[str, tuple[str | None, int, float]] = {} # username -> (role，账号不存在时为 None, token_version, 缓存时间)
|
|
|
|
|
|
def user_info(username: str) -> tuple[str, int] | None:
|
|
"""账号的 (角色, token 版本)；账号不存在返回 None。缓存 30 秒，改账号后调用 forget_roles()"""
|
|
hit = _roles.get(username)
|
|
if hit and time.time() - hit[2] < 30:
|
|
return (hit[0], hit[1]) if hit[0] is not None else None
|
|
with database.session() as db:
|
|
u = db.query(User).filter(User.username == username).first()
|
|
if u is None:
|
|
_roles[username] = (None, 0, time.time())
|
|
return None
|
|
_roles[username] = ((u.role or "admin"), u.token_version or 0, time.time())
|
|
return _roles[username][0], _roles[username][1]
|
|
|
|
|
|
def role_of(username: str) -> str | None:
|
|
info = user_info(username)
|
|
return info[0] if info else None
|
|
|
|
|
|
def forget_roles():
|
|
_roles.clear()
|
|
|
|
|
|
def current_user(request: Request, cred: HTTPAuthorizationCredentials | None = Depends(_bearer)) -> str:
|
|
if cred is None:
|
|
raise HTTPException(401, "未登录")
|
|
if cred.credentials.startswith(API_PREFIX):
|
|
# 只读 API 令牌：等同只读用户，连改自己密码都不行
|
|
name = _api_token_name(cred.credentials)
|
|
if name is None:
|
|
raise HTTPException(401, "API 令牌无效或已被删除")
|
|
request.state.role = "user"
|
|
if request.method not in ("GET", "HEAD"):
|
|
raise HTTPException(403, "API 令牌只能读取数据")
|
|
_read_only(request)
|
|
return "api:" + name
|
|
try:
|
|
data = jwt.decode(cred.credentials, secret_key(), algorithms=["HS256"])
|
|
except JWTError:
|
|
raise HTTPException(401, "登录已过期，请重新登录")
|
|
user = data.get("sub")
|
|
if not user:
|
|
raise HTTPException(401, "登录已过期，请重新登录")
|
|
info = user_info(user)
|
|
if info is None:
|
|
raise HTTPException(401, "账号不存在或已被删除，请重新登录")
|
|
role, ver = info
|
|
if data.get("ver", 0) != ver:
|
|
# 改密码 / 管理员改权限后版本号会 +1，旧 token 立即失效
|
|
raise HTTPException(401, "登录状态已失效（密码或权限已变更），请重新登录")
|
|
request.state.role = role
|
|
if role != "admin":
|
|
_read_only(request)
|
|
return user
|
|
|
|
|
|
def require_admin(request: Request, user: str = Depends(current_user)) -> str:
|
|
if getattr(request.state, "role", None) != "admin":
|
|
raise HTTPException(403, "只有管理员可以操作")
|
|
return user
|
|
|
|
|
|
# ---------------- 登录限速：同一 IP + 用户名连续失败 5 次后锁定 15 分钟 ----------------
|
|
MAX_LOGIN_FAILS = 5
|
|
LOGIN_LOCK_SECONDS = 900
|
|
_login_fails: dict[tuple[str, str], list] = {} # (ip, username) -> [连续失败次数, 锁定到]
|
|
# 不看来源地址、只按用户名的总失败数：攻击者换 IP 池爆破同一个账号时兜底
|
|
MAX_USER_FAILS = 20
|
|
_user_fails: dict[str, list] = {} # username -> [失败次数, 锁定到]
|
|
|
|
|
|
def _fail_key(request: Request, username: str) -> tuple[str, str]:
|
|
from app.oplog import client_ip
|
|
return client_ip(request), username
|
|
|
|
|
|
def _trim(d: dict):
|
|
"""满了淘汰最早的一批（以前整表清空，攻击者用 1 万个假用户名就能把目标账号的锁定全部重置）"""
|
|
if len(d) > 10_000:
|
|
for k in list(d)[:1000]:
|
|
d.pop(k, None)
|
|
|
|
|
|
def check_login_lock(request: Request, username: str):
|
|
now = time.time()
|
|
u = _user_fails.get(username)
|
|
if u and u[0] >= MAX_USER_FAILS:
|
|
if now < u[1]:
|
|
raise HTTPException(429, f"这个账号失败次数过多，请 {max(int((u[1] - now) / 60), 1)} 分钟后再试")
|
|
_user_fails.pop(username, None)
|
|
key = _fail_key(request, username)
|
|
e = _login_fails.get(key)
|
|
if not e:
|
|
return
|
|
if e[0] >= MAX_LOGIN_FAILS:
|
|
if now < e[1]:
|
|
raise HTTPException(429, f"失败次数过多，请 {max(int((e[1] - now) / 60), 1)} 分钟后再试")
|
|
_login_fails.pop(key, None) # 锁定期已过，重新计数
|
|
|
|
|
|
def note_login_fail(request: Request, username: str):
|
|
_trim(_login_fails)
|
|
_trim(_user_fails)
|
|
u = _user_fails.get(username) or [0, 0.0]
|
|
u[0] += 1
|
|
if u[0] >= MAX_USER_FAILS:
|
|
u[1] = time.time() + LOGIN_LOCK_SECONDS
|
|
_user_fails[username] = u
|
|
key = _fail_key(request, username)
|
|
e = _login_fails.get(key) or [0, 0.0]
|
|
if e[0] >= MAX_LOGIN_FAILS and time.time() >= e[1]:
|
|
e = [0, 0.0]
|
|
e[0] += 1
|
|
if e[0] >= MAX_LOGIN_FAILS:
|
|
e[1] = time.time() + LOGIN_LOCK_SECONDS
|
|
_login_fails[key] = e
|
|
|
|
|
|
def note_login_ok(request: Request, username: str):
|
|
_login_fails.pop(_fail_key(request, username), None)
|
|
_user_fails.pop(username, None)
|
|
|
|
|
|
# ---------------- 首次设置的安装码 ----------------
|
|
# 安装向导和创建管理员的接口在设置完成前不需要登录：谁先打开网页谁就能把实例接管。内网访问照旧不受影响；
|
|
# 从公网地址访问（例如已经通过反向代理暴露出去）时，要填启动日志里打印的一次性安装码
|
|
_setup_code = secrets.token_hex(4)
|
|
|
|
|
|
def setup_code() -> str:
|
|
return _setup_code
|
|
|
|
|
|
def check_setup_code(request: Request):
|
|
import hmac
|
|
import ipaddress
|
|
|
|
from app.oplog import client_ip
|
|
ip = client_ip(request)
|
|
try:
|
|
a = ipaddress.ip_address(ip)
|
|
if a.is_private or a.is_loopback or a.is_link_local:
|
|
return
|
|
except ValueError:
|
|
return # 测试客户端等没有真实地址的情况
|
|
code = (request.headers.get("x-setup-code") or "").strip().lower()
|
|
if not hmac.compare_digest(code, _setup_code):
|
|
raise HTTPException(403, "setup_code_required")
|