v_2.2
This commit is contained in:
@@ -0,0 +1,587 @@
|
||||
"""
|
||||
PST Indexer — Admin Panel (port 8001).
|
||||
|
||||
Separate FastAPI instance on its own port so operators can place it behind a
|
||||
more restrictive firewall rule or VPN than the main app.
|
||||
|
||||
Exposes:
|
||||
- Admin login (separate session cookie, scoped "admin")
|
||||
- List, create, delete users
|
||||
- Toggle admin role
|
||||
- Reset passwords
|
||||
- Enroll / confirm / disable MFA for any user (admin-managed flow)
|
||||
- Self-service MFA for the currently logged-in admin
|
||||
"""
|
||||
import io
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import qrcode
|
||||
from fastapi import (
|
||||
Depends, FastAPI, HTTPException, Request, Response, status,
|
||||
)
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse, JSONResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from pydantic import BaseModel
|
||||
|
||||
import auth
|
||||
from auth import (
|
||||
ADMIN_SESSION_COOKIE_NAME, SESSION_TTL_SECONDS,
|
||||
User, begin_mfa_enrollment, confirm_mfa_enrollment, count_admins,
|
||||
count_users, create_session_token, create_user, delete_user,
|
||||
disable_mfa, find_user_by_id, find_user_by_username,
|
||||
get_totp_secret, list_users, login_limiter, record_login,
|
||||
require_admin_user, set_user_role, totp_provisioning_uri,
|
||||
update_password, verify_password, verify_totp,
|
||||
)
|
||||
|
||||
ADMIN_FRONTEND_DIR = Path(os.environ.get("ADMIN_FRONTEND_DIR", "/app/admin_frontend"))
|
||||
COOKIE_SECURE = os.environ.get("COOKIE_SECURE", "false").lower() == "true"
|
||||
|
||||
app = FastAPI(title="PST Indexer Admin Panel")
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=[],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def _startup():
|
||||
auth.init_auth_schema()
|
||||
|
||||
|
||||
def _set_admin_cookie(response: Response, token: str):
|
||||
response.set_cookie(
|
||||
key=ADMIN_SESSION_COOKIE_NAME,
|
||||
value=token,
|
||||
max_age=SESSION_TTL_SECONDS,
|
||||
httponly=True,
|
||||
secure=COOKIE_SECURE,
|
||||
samesite="lax",
|
||||
path="/",
|
||||
)
|
||||
|
||||
|
||||
def _clear_admin_cookie(response: Response):
|
||||
response.delete_cookie(key=ADMIN_SESSION_COOKIE_NAME, path="/")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Auth
|
||||
# ---------------------------------------------------------------------------
|
||||
class LoginPayload(BaseModel):
|
||||
username: str
|
||||
password: str
|
||||
totp_code: Optional[str] = None
|
||||
|
||||
|
||||
@app.get("/api/health")
|
||||
def health():
|
||||
return {"status": "ok", "admins": count_admins(), "users": count_users()}
|
||||
|
||||
|
||||
@app.get("/api/setup/status")
|
||||
def setup_status():
|
||||
"""True if no users exist. Admin UI redirects to main app to complete setup."""
|
||||
return {"needs_setup": count_users() == 0}
|
||||
|
||||
|
||||
@app.post("/api/login")
|
||||
def admin_login(payload: LoginPayload, request: Request, response: Response):
|
||||
ip = request.client.host if request.client else "unknown"
|
||||
rate_key = f"admin:{ip}:{payload.username.lower()}"
|
||||
if not login_limiter.check(rate_key):
|
||||
raise HTTPException(429, "Too many login attempts. Please wait a few minutes.")
|
||||
|
||||
found = find_user_by_username(payload.username)
|
||||
if not found:
|
||||
login_limiter.record(rate_key)
|
||||
raise HTTPException(401, "Invalid credentials")
|
||||
user, pw_hash = found
|
||||
|
||||
if not verify_password(pw_hash, payload.password):
|
||||
login_limiter.record(rate_key)
|
||||
raise HTTPException(401, "Invalid credentials")
|
||||
|
||||
if not user.is_admin:
|
||||
raise HTTPException(403, "Admin access required")
|
||||
|
||||
if user.mfa_enabled:
|
||||
if not payload.totp_code:
|
||||
return JSONResponse(
|
||||
status_code=401,
|
||||
content={"detail": "MFA required", "mfa_required": True},
|
||||
)
|
||||
secret = get_totp_secret(user.id)
|
||||
if not secret or not verify_totp(secret, payload.totp_code):
|
||||
login_limiter.record(rate_key)
|
||||
raise HTTPException(401, "Invalid MFA code")
|
||||
|
||||
login_limiter.reset(rate_key)
|
||||
record_login(user.id)
|
||||
token = create_session_token(user, scope="admin")
|
||||
_set_admin_cookie(response, token)
|
||||
return {"user": user.public_dict()}
|
||||
|
||||
|
||||
@app.post("/api/logout")
|
||||
def admin_logout(response: Response):
|
||||
_clear_admin_cookie(response)
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
@app.get("/api/me")
|
||||
def me(user: User = Depends(require_admin_user)):
|
||||
return {"user": user.public_dict()}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# User management
|
||||
# ---------------------------------------------------------------------------
|
||||
class CreateUserPayload(BaseModel):
|
||||
username: str
|
||||
password: str
|
||||
role: str = "user"
|
||||
|
||||
|
||||
class UpdatePasswordPayload(BaseModel):
|
||||
password: str
|
||||
|
||||
|
||||
class UpdateRolePayload(BaseModel):
|
||||
role: str
|
||||
|
||||
|
||||
@app.get("/api/users")
|
||||
def api_list_users(_: User = Depends(require_admin_user)):
|
||||
return [u.public_dict() for u in list_users()]
|
||||
|
||||
|
||||
@app.post("/api/users")
|
||||
def api_create_user(
|
||||
payload: CreateUserPayload,
|
||||
admin: User = Depends(require_admin_user),
|
||||
):
|
||||
try:
|
||||
user = create_user(payload.username, payload.password, role=payload.role)
|
||||
except ValueError as e:
|
||||
raise HTTPException(400, str(e))
|
||||
return user.public_dict()
|
||||
|
||||
|
||||
@app.delete("/api/users/{user_id}")
|
||||
def api_delete_user(user_id: str, admin: User = Depends(require_admin_user)):
|
||||
target = find_user_by_id(user_id)
|
||||
if not target:
|
||||
raise HTTPException(404, "User not found")
|
||||
|
||||
# Prevent deleting yourself
|
||||
if target.id == admin.id:
|
||||
raise HTTPException(400, "You cannot delete your own account")
|
||||
|
||||
# Prevent deleting the last admin
|
||||
if target.is_admin and count_admins() <= 1:
|
||||
raise HTTPException(400, "Cannot delete the last remaining admin")
|
||||
|
||||
delete_user(user_id)
|
||||
return {"deleted": user_id}
|
||||
|
||||
|
||||
@app.put("/api/users/{user_id}/role")
|
||||
def api_set_role(
|
||||
user_id: str,
|
||||
payload: UpdateRolePayload,
|
||||
admin: User = Depends(require_admin_user),
|
||||
):
|
||||
target = find_user_by_id(user_id)
|
||||
if not target:
|
||||
raise HTTPException(404, "User not found")
|
||||
|
||||
if target.id == admin.id and payload.role != "admin":
|
||||
raise HTTPException(400, "You cannot demote your own account")
|
||||
|
||||
if (target.is_admin and payload.role != "admin"
|
||||
and count_admins() <= 1):
|
||||
raise HTTPException(400, "Cannot demote the last remaining admin")
|
||||
|
||||
try:
|
||||
set_user_role(user_id, payload.role)
|
||||
except ValueError as e:
|
||||
raise HTTPException(400, str(e))
|
||||
|
||||
return find_user_by_id(user_id).public_dict()
|
||||
|
||||
|
||||
@app.put("/api/users/{user_id}/password")
|
||||
def api_update_password(
|
||||
user_id: str,
|
||||
payload: UpdatePasswordPayload,
|
||||
admin: User = Depends(require_admin_user),
|
||||
):
|
||||
if not find_user_by_id(user_id):
|
||||
raise HTTPException(404, "User not found")
|
||||
try:
|
||||
update_password(user_id, payload.password)
|
||||
except ValueError as e:
|
||||
raise HTTPException(400, str(e))
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MFA management
|
||||
# ---------------------------------------------------------------------------
|
||||
class ConfirmMfaPayload(BaseModel):
|
||||
totp_code: str
|
||||
|
||||
|
||||
@app.post("/api/users/{user_id}/mfa/begin")
|
||||
def api_begin_mfa(user_id: str, admin: User = Depends(require_admin_user)):
|
||||
"""Provision a new TOTP secret for the target user (admin action).
|
||||
Returns the secret + otpauth URI so the admin can show it to the user."""
|
||||
target = find_user_by_id(user_id)
|
||||
if not target:
|
||||
raise HTTPException(404, "User not found")
|
||||
secret, uri = begin_mfa_enrollment(user_id)
|
||||
return {
|
||||
"secret": secret,
|
||||
"otpauth_uri": uri,
|
||||
"username": target.username,
|
||||
"qr_url": f"/api/users/{user_id}/mfa/qr",
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/users/{user_id}/mfa/qr")
|
||||
def api_mfa_qr(user_id: str, admin: User = Depends(require_admin_user)):
|
||||
"""Render a PNG QR code of the target user's current (unconfirmed) TOTP URI."""
|
||||
target = find_user_by_id(user_id)
|
||||
if not target:
|
||||
raise HTTPException(404, "User not found")
|
||||
secret = get_totp_secret(user_id)
|
||||
if not secret:
|
||||
raise HTTPException(400, "No MFA enrollment in progress")
|
||||
uri = totp_provisioning_uri(secret, target.username)
|
||||
|
||||
qr = qrcode.QRCode(box_size=6, border=2)
|
||||
qr.add_data(uri)
|
||||
qr.make(fit=True)
|
||||
img = qr.make_image(fill_color="black", back_color="white")
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
buf.seek(0)
|
||||
return StreamingResponse(buf, media_type="image/png")
|
||||
|
||||
|
||||
@app.post("/api/users/{user_id}/mfa/confirm")
|
||||
def api_confirm_mfa(
|
||||
user_id: str,
|
||||
payload: ConfirmMfaPayload,
|
||||
admin: User = Depends(require_admin_user),
|
||||
):
|
||||
if not find_user_by_id(user_id):
|
||||
raise HTTPException(404, "User not found")
|
||||
if not confirm_mfa_enrollment(user_id, payload.totp_code):
|
||||
raise HTTPException(400, "Invalid MFA code — enrollment not confirmed")
|
||||
return find_user_by_id(user_id).public_dict()
|
||||
|
||||
|
||||
@app.post("/api/users/{user_id}/mfa/disable")
|
||||
def api_disable_mfa(user_id: str, admin: User = Depends(require_admin_user)):
|
||||
target = find_user_by_id(user_id)
|
||||
if not target:
|
||||
raise HTTPException(404, "User not found")
|
||||
disable_mfa(user_id)
|
||||
return find_user_by_id(user_id).public_dict()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Folder management (admin)
|
||||
# ---------------------------------------------------------------------------
|
||||
import uuid as _uuid
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path as _Path
|
||||
from auth import get_db
|
||||
|
||||
|
||||
UPLOAD_DIR = _Path(os.environ.get("DATA_DIR", "/app/data")) / "uploads"
|
||||
|
||||
|
||||
class AdminFolderPayload(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
def _validate_name(name: str) -> str:
|
||||
name = (name or "").strip()
|
||||
if not name:
|
||||
raise HTTPException(400, "Folder name cannot be empty")
|
||||
if len(name) > 80:
|
||||
raise HTTPException(400, "Folder name too long (max 80 characters)")
|
||||
return name
|
||||
|
||||
|
||||
@app.get("/api/folders")
|
||||
def admin_list_folders(_: User = Depends(require_admin_user)):
|
||||
with get_db() as conn:
|
||||
rows = conn.execute(
|
||||
"""SELECT f.id, f.name, f.created_at, f.created_by,
|
||||
COUNT(p.id) AS file_count
|
||||
FROM pst_folders f
|
||||
LEFT JOIN pst_files p ON p.folder_id = f.id
|
||||
GROUP BY f.id
|
||||
ORDER BY f.name COLLATE NOCASE"""
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
@app.post("/api/folders")
|
||||
def admin_create_folder(payload: AdminFolderPayload,
|
||||
admin: User = Depends(require_admin_user)):
|
||||
name = _validate_name(payload.name)
|
||||
with get_db() as conn:
|
||||
existing = conn.execute(
|
||||
"SELECT id FROM pst_folders WHERE name = ? COLLATE NOCASE", (name,)
|
||||
).fetchone()
|
||||
if existing:
|
||||
raise HTTPException(409, "A folder with that name already exists")
|
||||
folder_id = _uuid.uuid4().hex
|
||||
conn.execute(
|
||||
"""INSERT INTO pst_folders (id, name, created_at, created_by)
|
||||
VALUES (?, ?, ?, ?)""",
|
||||
(folder_id, name, datetime.now(timezone.utc).isoformat(), admin.id),
|
||||
)
|
||||
conn.commit()
|
||||
return {"id": folder_id, "name": name, "file_count": 0}
|
||||
|
||||
|
||||
@app.put("/api/folders/{folder_id}")
|
||||
def admin_rename_folder(folder_id: str, payload: AdminFolderPayload,
|
||||
admin: User = Depends(require_admin_user)):
|
||||
name = _validate_name(payload.name)
|
||||
with get_db() as conn:
|
||||
if not conn.execute("SELECT id FROM pst_folders WHERE id = ?",
|
||||
(folder_id,)).fetchone():
|
||||
raise HTTPException(404, "Folder not found")
|
||||
clash = conn.execute(
|
||||
"SELECT id FROM pst_folders WHERE name = ? COLLATE NOCASE AND id != ?",
|
||||
(name, folder_id),
|
||||
).fetchone()
|
||||
if clash:
|
||||
raise HTTPException(409, "A folder with that name already exists")
|
||||
conn.execute("UPDATE pst_folders SET name = ? WHERE id = ?",
|
||||
(name, folder_id))
|
||||
conn.commit()
|
||||
return {"id": folder_id, "name": name}
|
||||
|
||||
|
||||
@app.delete("/api/folders/{folder_id}")
|
||||
def admin_delete_folder(folder_id: str,
|
||||
admin: User = Depends(require_admin_user)):
|
||||
"""Delete a folder. PSTs inside it become uncategorized (folder_id = NULL)."""
|
||||
with get_db() as conn:
|
||||
if not conn.execute("SELECT id FROM pst_folders WHERE id = ?",
|
||||
(folder_id,)).fetchone():
|
||||
raise HTTPException(404, "Folder not found")
|
||||
conn.execute("UPDATE pst_files SET folder_id = NULL WHERE folder_id = ?",
|
||||
(folder_id,))
|
||||
conn.execute("DELETE FROM pst_folders WHERE id = ?", (folder_id,))
|
||||
conn.commit()
|
||||
return {"deleted": folder_id}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PST file management (admin)
|
||||
# ---------------------------------------------------------------------------
|
||||
@app.get("/api/pst-files")
|
||||
def admin_list_pst_files(_: User = Depends(require_admin_user)):
|
||||
with get_db() as conn:
|
||||
rows = conn.execute(
|
||||
"""SELECT p.id, p.original_name, p.uploaded_at, p.size_bytes,
|
||||
p.status, p.message_count, p.error,
|
||||
p.folder_id, f.name AS folder_name,
|
||||
p.uploaded_by
|
||||
FROM pst_files p
|
||||
LEFT JOIN pst_folders f ON f.id = p.folder_id
|
||||
ORDER BY p.uploaded_at DESC"""
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
@app.delete("/api/pst-files/{pst_id}")
|
||||
def admin_delete_pst(pst_id: str, admin: User = Depends(require_admin_user)):
|
||||
with get_db() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT filename FROM pst_files WHERE id = ?", (pst_id,)
|
||||
).fetchone()
|
||||
if not row:
|
||||
raise HTTPException(404, "PST not found")
|
||||
conn.execute("DELETE FROM messages WHERE pst_id = ?", (pst_id,))
|
||||
conn.execute("DELETE FROM pst_files WHERE id = ?", (pst_id,))
|
||||
conn.commit()
|
||||
try:
|
||||
(UPLOAD_DIR / row["filename"]).unlink(missing_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
return {"deleted": pst_id}
|
||||
|
||||
|
||||
class AdminAssignPayload(BaseModel):
|
||||
folder_id: Optional[str] = None
|
||||
|
||||
|
||||
@app.put("/api/pst-files/{pst_id}/folder")
|
||||
def admin_assign_pst_to_folder(pst_id: str, payload: AdminAssignPayload,
|
||||
admin: User = Depends(require_admin_user)):
|
||||
with get_db() as conn:
|
||||
if not conn.execute("SELECT id FROM pst_files WHERE id = ?",
|
||||
(pst_id,)).fetchone():
|
||||
raise HTTPException(404, "PST not found")
|
||||
if payload.folder_id:
|
||||
if not conn.execute("SELECT id FROM pst_folders WHERE id = ?",
|
||||
(payload.folder_id,)).fetchone():
|
||||
raise HTTPException(404, "Folder not found")
|
||||
conn.execute("UPDATE pst_files SET folder_id = ? WHERE id = ?",
|
||||
(payload.folder_id, pst_id))
|
||||
conn.commit()
|
||||
return {"pst_id": pst_id, "folder_id": payload.folder_id}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Folder permissions (admin-only)
|
||||
# Controls which user can access which folder.
|
||||
# Admins ALWAYS have access to every folder regardless of this table —
|
||||
# these grants are for non-admin users.
|
||||
# ---------------------------------------------------------------------------
|
||||
class PermissionToggle(BaseModel):
|
||||
granted: bool
|
||||
|
||||
|
||||
@app.get("/api/permissions")
|
||||
def admin_list_permissions(_: User = Depends(require_admin_user)):
|
||||
"""Return the full matrix: users, folders, and all grants."""
|
||||
with get_db() as conn:
|
||||
users = conn.execute("SELECT * FROM users ORDER BY created_at ASC").fetchall()
|
||||
folders = conn.execute(
|
||||
"SELECT id, name FROM pst_folders ORDER BY name COLLATE NOCASE"
|
||||
).fetchall()
|
||||
grants = conn.execute(
|
||||
"SELECT user_id, folder_id, granted_at FROM folder_permissions"
|
||||
).fetchall()
|
||||
|
||||
from auth import _row_to_user
|
||||
user_list = []
|
||||
for u in users:
|
||||
pub = _row_to_user(u).public_dict()
|
||||
user_list.append(pub)
|
||||
|
||||
return {
|
||||
"users": user_list,
|
||||
"folders": [dict(f) for f in folders],
|
||||
"grants": [dict(g) for g in grants],
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/users/{user_id}/folders")
|
||||
def admin_user_folders(user_id: str, _: User = Depends(require_admin_user)):
|
||||
"""Folders a given user has been granted."""
|
||||
target = find_user_by_id(user_id)
|
||||
if not target:
|
||||
raise HTTPException(404, "User not found")
|
||||
with get_db() as conn:
|
||||
rows = conn.execute(
|
||||
"""SELECT f.id, f.name, p.granted_at
|
||||
FROM folder_permissions p
|
||||
JOIN pst_folders f ON f.id = p.folder_id
|
||||
WHERE p.user_id = ?
|
||||
ORDER BY f.name COLLATE NOCASE""",
|
||||
(user_id,),
|
||||
).fetchall()
|
||||
return {
|
||||
"user": target.public_dict(),
|
||||
"folders": [dict(r) for r in rows],
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/folders/{folder_id}/users")
|
||||
def admin_folder_users(folder_id: str, _: User = Depends(require_admin_user)):
|
||||
"""Users granted on a given folder."""
|
||||
with get_db() as conn:
|
||||
folder = conn.execute(
|
||||
"SELECT id, name FROM pst_folders WHERE id = ?", (folder_id,)
|
||||
).fetchone()
|
||||
if not folder:
|
||||
raise HTTPException(404, "Folder not found")
|
||||
|
||||
grant_rows = conn.execute(
|
||||
"""SELECT user_id, granted_at
|
||||
FROM folder_permissions WHERE folder_id = ?""",
|
||||
(folder_id,),
|
||||
).fetchall()
|
||||
|
||||
# Hydrate user details (usernames are encrypted — need decrypt via find_user_by_id)
|
||||
users_out = []
|
||||
for r in grant_rows:
|
||||
u = find_user_by_id(r["user_id"])
|
||||
if u:
|
||||
users_out.append({
|
||||
**u.public_dict(),
|
||||
"granted_at": r["granted_at"],
|
||||
})
|
||||
users_out.sort(key=lambda x: x["username"].lower())
|
||||
return {
|
||||
"folder": dict(folder),
|
||||
"users": users_out,
|
||||
}
|
||||
|
||||
|
||||
@app.put("/api/permissions/{user_id}/{folder_id}")
|
||||
def admin_toggle_permission(
|
||||
user_id: str,
|
||||
folder_id: str,
|
||||
payload: PermissionToggle,
|
||||
admin: User = Depends(require_admin_user),
|
||||
):
|
||||
target = find_user_by_id(user_id)
|
||||
if not target:
|
||||
raise HTTPException(404, "User not found")
|
||||
|
||||
if target.is_admin:
|
||||
# Admins implicitly have access to everything; refuse explicit grants
|
||||
# so the UI doesn't confuse admins with checked-but-meaningless rows.
|
||||
raise HTTPException(
|
||||
400,
|
||||
"Admins have access to all folders automatically; permissions "
|
||||
"only apply to non-admin users.",
|
||||
)
|
||||
|
||||
with get_db() as conn:
|
||||
if not conn.execute("SELECT id FROM pst_folders WHERE id = ?",
|
||||
(folder_id,)).fetchone():
|
||||
raise HTTPException(404, "Folder not found")
|
||||
|
||||
if payload.granted:
|
||||
conn.execute(
|
||||
"""INSERT OR IGNORE INTO folder_permissions
|
||||
(user_id, folder_id, granted_at, granted_by)
|
||||
VALUES (?, ?, ?, ?)""",
|
||||
(user_id, folder_id,
|
||||
datetime.now(timezone.utc).isoformat(), admin.id),
|
||||
)
|
||||
else:
|
||||
conn.execute(
|
||||
"DELETE FROM folder_permissions WHERE user_id = ? AND folder_id = ?",
|
||||
(user_id, folder_id),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
return {"user_id": user_id, "folder_id": folder_id, "granted": payload.granted}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Static admin frontend
|
||||
# ---------------------------------------------------------------------------
|
||||
if ADMIN_FRONTEND_DIR.exists():
|
||||
app.mount("/", StaticFiles(directory=str(ADMIN_FRONTEND_DIR), html=True), name="admin")
|
||||
@@ -0,0 +1,835 @@
|
||||
"""
|
||||
PST Indexer — Main application (port 8000).
|
||||
|
||||
Serves:
|
||||
- The web UI (login, setup wizard, archive browser)
|
||||
- REST API for PST upload, indexing, search
|
||||
- All data-access endpoints require an authenticated session
|
||||
|
||||
First-run behavior:
|
||||
- If no users exist, the UI routes users to a setup wizard.
|
||||
- /api/setup is only callable while the users table is empty; after the first
|
||||
admin is created, the endpoint returns 409.
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import uuid
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import aiofiles
|
||||
import pypff
|
||||
from fastapi import (
|
||||
BackgroundTasks, Depends, FastAPI, File, HTTPException,
|
||||
Request, Response, UploadFile, status,
|
||||
)
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import JSONResponse, FileResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from pydantic import BaseModel
|
||||
|
||||
import auth
|
||||
from auth import (
|
||||
SESSION_COOKIE_NAME, SESSION_TTL_SECONDS,
|
||||
User, count_users, count_admins, create_user, find_user_by_username,
|
||||
find_user_by_id, get_totp_secret, verify_password, verify_totp,
|
||||
record_login, create_session_token, login_limiter, require_app_user,
|
||||
get_db,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
DATA_DIR = Path(os.environ.get("DATA_DIR", "/app/data"))
|
||||
UPLOAD_DIR = DATA_DIR / "uploads"
|
||||
FRONTEND_DIR = Path(os.environ.get("FRONTEND_DIR", "/app/frontend"))
|
||||
|
||||
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
|
||||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Cookie security. Flip to True when serving over HTTPS.
|
||||
COOKIE_SECURE = os.environ.get("COOKIE_SECURE", "false").lower() == "true"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# App DB schema (messages + pst_files)
|
||||
# ---------------------------------------------------------------------------
|
||||
def init_app_schema():
|
||||
with get_db() as conn:
|
||||
conn.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS pst_folders (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
created_by TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS folder_permissions (
|
||||
user_id TEXT NOT NULL,
|
||||
folder_id TEXT NOT NULL,
|
||||
granted_at TEXT NOT NULL,
|
||||
granted_by TEXT,
|
||||
PRIMARY KEY (user_id, folder_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_folder_perms_user ON folder_permissions(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_folder_perms_folder ON folder_permissions(folder_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS pst_files (
|
||||
id TEXT PRIMARY KEY,
|
||||
filename TEXT NOT NULL,
|
||||
original_name TEXT NOT NULL,
|
||||
uploaded_at TEXT NOT NULL,
|
||||
size_bytes INTEGER NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
message_count INTEGER DEFAULT 0,
|
||||
error TEXT,
|
||||
uploaded_by TEXT,
|
||||
folder_id TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
pst_id TEXT NOT NULL,
|
||||
folder TEXT,
|
||||
subject TEXT,
|
||||
sender TEXT,
|
||||
recipients TEXT,
|
||||
sent_date TEXT,
|
||||
body TEXT,
|
||||
attachments TEXT,
|
||||
FOREIGN KEY (pst_id) REFERENCES pst_files(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_pst ON messages(pst_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_sent_date ON messages(sent_date);
|
||||
CREATE INDEX IF NOT EXISTS idx_pst_files_folder ON pst_files(folder_id);
|
||||
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
|
||||
subject, sender, recipients, body, folder,
|
||||
content='messages',
|
||||
content_rowid='id',
|
||||
tokenize='porter unicode61'
|
||||
);
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
|
||||
INSERT INTO messages_fts(rowid, subject, sender, recipients, body, folder)
|
||||
VALUES (new.id, new.subject, new.sender, new.recipients, new.body, new.folder);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
|
||||
INSERT INTO messages_fts(messages_fts, rowid, subject, sender, recipients, body, folder)
|
||||
VALUES ('delete', old.id, old.subject, old.sender, old.recipients, old.body, old.folder);
|
||||
END;
|
||||
""")
|
||||
# Backfill: add folder_id column to DBs created before this feature existed.
|
||||
cols = [r["name"] for r in conn.execute("PRAGMA table_info(pst_files)").fetchall()]
|
||||
if "folder_id" not in cols:
|
||||
conn.execute("ALTER TABLE pst_files ADD COLUMN folder_id TEXT")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_pst_files_folder ON pst_files(folder_id)")
|
||||
conn.commit()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Access control helpers — which folders / PSTs is a given user allowed to see?
|
||||
#
|
||||
# Rules:
|
||||
# - Admins see everything.
|
||||
# - Regular users see only folders they have been explicitly granted,
|
||||
# plus the implicit "Uncategorized" bucket (folder_id = NULL) — this is
|
||||
# deliberately shared because otherwise newly-uploaded PSTs would be
|
||||
# invisible until someone moved them into a permitted folder.
|
||||
# - A PST is visible iff its folder is visible (or the PST is uncategorized).
|
||||
# ---------------------------------------------------------------------------
|
||||
def allowed_folder_ids_for(user) -> set[str]:
|
||||
"""Return the set of folder_ids a user is explicitly permitted on.
|
||||
For admins this is not used — they see everything unconditionally."""
|
||||
with get_db() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT folder_id FROM folder_permissions WHERE user_id = ?",
|
||||
(user.id,),
|
||||
).fetchall()
|
||||
return {r["folder_id"] for r in rows}
|
||||
|
||||
|
||||
def user_can_see_folder(user, folder_id: Optional[str]) -> bool:
|
||||
"""Can this user see content inside this folder? folder_id may be None
|
||||
(Uncategorized), which is visible to everyone."""
|
||||
if user.is_admin:
|
||||
return True
|
||||
if folder_id is None:
|
||||
return True # Uncategorized is shared
|
||||
return folder_id in allowed_folder_ids_for(user)
|
||||
|
||||
|
||||
def user_can_see_pst(user, pst_id: str) -> bool:
|
||||
if user.is_admin:
|
||||
return True
|
||||
with get_db() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT folder_id FROM pst_files WHERE id = ?", (pst_id,)
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
return user_can_see_folder(user, row["folder_id"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PST indexing helpers (unchanged from pre-auth version)
|
||||
# ---------------------------------------------------------------------------
|
||||
def _safe_str(value) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, bytes):
|
||||
try:
|
||||
return value.decode("utf-8", errors="replace")
|
||||
except Exception:
|
||||
return value.decode("latin-1", errors="replace")
|
||||
return str(value)
|
||||
|
||||
|
||||
def _extract_recipients(message) -> str:
|
||||
parts = []
|
||||
try:
|
||||
headers = _safe_str(getattr(message, "transport_headers", ""))
|
||||
if headers:
|
||||
for line in headers.split("\n"):
|
||||
low = line.lower()
|
||||
if low.startswith("to:") or low.startswith("cc:") or low.startswith("bcc:"):
|
||||
parts.append(line.strip())
|
||||
except Exception:
|
||||
pass
|
||||
return " | ".join(parts) if parts else ""
|
||||
|
||||
|
||||
def _extract_body(message) -> str:
|
||||
for attr in ("plain_text_body", "html_body", "rtf_body"):
|
||||
try:
|
||||
body = getattr(message, attr, None)
|
||||
if body:
|
||||
text = _safe_str(body)
|
||||
if attr == "html_body":
|
||||
text = re.sub(r"<[^>]+>", " ", text)
|
||||
text = re.sub(r"\s+", " ", text).strip()
|
||||
if text:
|
||||
return text
|
||||
except Exception:
|
||||
continue
|
||||
return ""
|
||||
|
||||
|
||||
def _extract_sender(message) -> str:
|
||||
for attr in ("sender_name", "sender_email_address"):
|
||||
try:
|
||||
v = _safe_str(getattr(message, attr, ""))
|
||||
if v:
|
||||
return v
|
||||
except Exception:
|
||||
continue
|
||||
try:
|
||||
headers = _safe_str(getattr(message, "transport_headers", ""))
|
||||
for line in headers.split("\n"):
|
||||
if line.lower().startswith("from:"):
|
||||
return line.split(":", 1)[1].strip()
|
||||
except Exception:
|
||||
pass
|
||||
return ""
|
||||
|
||||
|
||||
def _extract_date(message) -> Optional[str]:
|
||||
try:
|
||||
d = getattr(message, "delivery_time", None) or getattr(message, "client_submit_time", None)
|
||||
if d:
|
||||
if isinstance(d, datetime):
|
||||
return d.isoformat()
|
||||
return _safe_str(d)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _extract_attachments(message) -> str:
|
||||
names = []
|
||||
try:
|
||||
count = getattr(message, "number_of_attachments", 0)
|
||||
for i in range(count):
|
||||
try:
|
||||
att = message.get_attachment(i)
|
||||
name = _safe_str(getattr(att, "name", "") or getattr(att, "long_filename", ""))
|
||||
if name:
|
||||
names.append(name)
|
||||
except Exception:
|
||||
continue
|
||||
except Exception:
|
||||
pass
|
||||
return " | ".join(names)
|
||||
|
||||
|
||||
def _walk_folders(folder, path=""):
|
||||
try:
|
||||
name = _safe_str(getattr(folder, "name", "") or "")
|
||||
except Exception:
|
||||
name = ""
|
||||
current_path = f"{path}/{name}" if name else (path or "/")
|
||||
try:
|
||||
for i in range(folder.number_of_sub_messages):
|
||||
try:
|
||||
msg = folder.get_sub_message(i)
|
||||
yield current_path, msg
|
||||
except Exception:
|
||||
continue
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
for i in range(folder.number_of_sub_folders):
|
||||
try:
|
||||
sub = folder.get_sub_folder(i)
|
||||
yield from _walk_folders(sub, current_path)
|
||||
except Exception:
|
||||
continue
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def index_pst_file(pst_id: str, file_path: Path):
|
||||
pff = None
|
||||
try:
|
||||
pff = pypff.file()
|
||||
pff.open(str(file_path))
|
||||
root = pff.get_root_folder()
|
||||
count = 0
|
||||
batch = []
|
||||
BATCH_SIZE = 500
|
||||
|
||||
with get_db() as conn:
|
||||
for folder_path, msg in _walk_folders(root):
|
||||
try:
|
||||
batch.append((
|
||||
pst_id, folder_path,
|
||||
_safe_str(getattr(msg, "subject", "")),
|
||||
_extract_sender(msg),
|
||||
_extract_recipients(msg),
|
||||
_extract_date(msg),
|
||||
_extract_body(msg),
|
||||
_extract_attachments(msg),
|
||||
))
|
||||
count += 1
|
||||
if len(batch) >= BATCH_SIZE:
|
||||
conn.executemany(
|
||||
"""INSERT INTO messages
|
||||
(pst_id, folder, subject, sender, recipients, sent_date, body, attachments)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
batch,
|
||||
)
|
||||
conn.commit()
|
||||
batch.clear()
|
||||
except Exception as e:
|
||||
print(f"[index] skip: {e}")
|
||||
continue
|
||||
|
||||
if batch:
|
||||
conn.executemany(
|
||||
"""INSERT INTO messages
|
||||
(pst_id, folder, subject, sender, recipients, sent_date, body, attachments)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
batch,
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
conn.execute(
|
||||
"UPDATE pst_files SET status = ?, message_count = ? WHERE id = ?",
|
||||
("ready", count, pst_id),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
print(f"[index] {pst_id}: {count} messages")
|
||||
except Exception as e:
|
||||
print(f"[index] failed: {e}")
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"UPDATE pst_files SET status = ?, error = ? WHERE id = ?",
|
||||
("failed", str(e), pst_id),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
if pff is not None:
|
||||
try:
|
||||
pff.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FastAPI
|
||||
# ---------------------------------------------------------------------------
|
||||
app = FastAPI(title="PST Indexer")
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=[],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def _startup():
|
||||
auth.init_auth_schema()
|
||||
init_app_schema()
|
||||
|
||||
|
||||
def _set_session_cookie(response: Response, token: str):
|
||||
response.set_cookie(
|
||||
key=SESSION_COOKIE_NAME,
|
||||
value=token,
|
||||
max_age=SESSION_TTL_SECONDS,
|
||||
httponly=True,
|
||||
secure=COOKIE_SECURE,
|
||||
samesite="lax",
|
||||
path="/",
|
||||
)
|
||||
|
||||
|
||||
def _clear_session_cookie(response: Response):
|
||||
response.delete_cookie(key=SESSION_COOKIE_NAME, path="/")
|
||||
|
||||
|
||||
# ----- Setup / auth endpoints -----
|
||||
|
||||
class SetupRequest(BaseModel):
|
||||
username: str
|
||||
password: str
|
||||
|
||||
|
||||
class LoginRequest(BaseModel):
|
||||
username: str
|
||||
password: str
|
||||
totp_code: Optional[str] = None
|
||||
|
||||
|
||||
@app.get("/api/setup/status")
|
||||
def setup_status():
|
||||
"""Lets the frontend decide whether to show the setup wizard."""
|
||||
return {"needs_setup": count_users() == 0}
|
||||
|
||||
|
||||
@app.post("/api/setup")
|
||||
def setup_admin(payload: SetupRequest, response: Response):
|
||||
if count_users() > 0:
|
||||
raise HTTPException(status.HTTP_409_CONFLICT, "Setup has already been completed")
|
||||
|
||||
try:
|
||||
user = create_user(payload.username, payload.password, role="admin")
|
||||
except ValueError as e:
|
||||
raise HTTPException(400, str(e))
|
||||
|
||||
record_login(user.id)
|
||||
token = create_session_token(user, scope="app")
|
||||
_set_session_cookie(response, token)
|
||||
return {"user": user.public_dict()}
|
||||
|
||||
|
||||
@app.post("/api/login")
|
||||
def login(payload: LoginRequest, request: Request, response: Response):
|
||||
ip = request.client.host if request.client else "unknown"
|
||||
rate_key = f"{ip}:{payload.username.lower()}"
|
||||
if not login_limiter.check(rate_key):
|
||||
raise HTTPException(429, "Too many login attempts. Please wait a few minutes.")
|
||||
|
||||
found = find_user_by_username(payload.username)
|
||||
if not found:
|
||||
login_limiter.record(rate_key)
|
||||
raise HTTPException(401, "Invalid credentials")
|
||||
|
||||
user, stored_pw_hash = found
|
||||
if not verify_password(stored_pw_hash, payload.password):
|
||||
login_limiter.record(rate_key)
|
||||
raise HTTPException(401, "Invalid credentials")
|
||||
|
||||
if user.mfa_enabled:
|
||||
if not payload.totp_code:
|
||||
# Tell the client MFA is required
|
||||
return JSONResponse(
|
||||
status_code=401,
|
||||
content={"detail": "MFA required", "mfa_required": True},
|
||||
)
|
||||
secret = get_totp_secret(user.id)
|
||||
if not secret or not verify_totp(secret, payload.totp_code):
|
||||
login_limiter.record(rate_key)
|
||||
raise HTTPException(401, "Invalid MFA code")
|
||||
|
||||
login_limiter.reset(rate_key)
|
||||
record_login(user.id)
|
||||
token = create_session_token(user, scope="app")
|
||||
_set_session_cookie(response, token)
|
||||
return {"user": user.public_dict()}
|
||||
|
||||
|
||||
@app.post("/api/logout")
|
||||
def logout(response: Response):
|
||||
_clear_session_cookie(response)
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
@app.get("/api/me")
|
||||
def me(user: User = Depends(require_app_user)):
|
||||
return {"user": user.public_dict()}
|
||||
|
||||
|
||||
# ----- PST endpoints (all require auth) -----
|
||||
|
||||
@app.get("/api/health")
|
||||
def health():
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
@app.get("/api/pst-files")
|
||||
def list_pst_files(user: User = Depends(require_app_user)):
|
||||
with get_db() as conn:
|
||||
if user.is_admin:
|
||||
rows = conn.execute(
|
||||
"""SELECT id, original_name, uploaded_at, size_bytes, status,
|
||||
message_count, error, folder_id
|
||||
FROM pst_files ORDER BY uploaded_at DESC"""
|
||||
).fetchall()
|
||||
else:
|
||||
allowed = allowed_folder_ids_for(user)
|
||||
if allowed:
|
||||
placeholders = ",".join("?" for _ in allowed)
|
||||
rows = conn.execute(
|
||||
f"""SELECT id, original_name, uploaded_at, size_bytes, status,
|
||||
message_count, error, folder_id
|
||||
FROM pst_files
|
||||
WHERE folder_id IS NULL OR folder_id IN ({placeholders})
|
||||
ORDER BY uploaded_at DESC""",
|
||||
list(allowed),
|
||||
).fetchall()
|
||||
else:
|
||||
rows = conn.execute(
|
||||
"""SELECT id, original_name, uploaded_at, size_bytes, status,
|
||||
message_count, error, folder_id
|
||||
FROM pst_files WHERE folder_id IS NULL
|
||||
ORDER BY uploaded_at DESC"""
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
# ----- Folder endpoints (shared with admin panel) -----
|
||||
|
||||
class FolderCreate(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
class FolderRename(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
class FolderAssign(BaseModel):
|
||||
folder_id: Optional[str] = None # null = move to uncategorized
|
||||
|
||||
|
||||
def _validate_folder_name(name: str) -> str:
|
||||
name = (name or "").strip()
|
||||
if not name:
|
||||
raise HTTPException(400, "Folder name cannot be empty")
|
||||
if len(name) > 80:
|
||||
raise HTTPException(400, "Folder name too long (max 80 characters)")
|
||||
return name
|
||||
|
||||
|
||||
@app.get("/api/folders")
|
||||
def list_folders(user: User = Depends(require_app_user)):
|
||||
with get_db() as conn:
|
||||
if user.is_admin:
|
||||
rows = conn.execute(
|
||||
"""SELECT f.id, f.name, f.created_at, f.created_by,
|
||||
COUNT(p.id) AS file_count
|
||||
FROM pst_folders f
|
||||
LEFT JOIN pst_files p ON p.folder_id = f.id
|
||||
GROUP BY f.id
|
||||
ORDER BY f.name COLLATE NOCASE"""
|
||||
).fetchall()
|
||||
else:
|
||||
allowed = allowed_folder_ids_for(user)
|
||||
if not allowed:
|
||||
return []
|
||||
placeholders = ",".join("?" for _ in allowed)
|
||||
rows = conn.execute(
|
||||
f"""SELECT f.id, f.name, f.created_at, f.created_by,
|
||||
COUNT(p.id) AS file_count
|
||||
FROM pst_folders f
|
||||
LEFT JOIN pst_files p ON p.folder_id = f.id
|
||||
WHERE f.id IN ({placeholders})
|
||||
GROUP BY f.id
|
||||
ORDER BY f.name COLLATE NOCASE""",
|
||||
list(allowed),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
@app.post("/api/folders")
|
||||
def create_folder(payload: FolderCreate, user: User = Depends(require_app_user)):
|
||||
# Only admins can create folders — the admin panel is the authoritative
|
||||
# management surface. (We still allow creation from the main app sidebar,
|
||||
# but only for admin users.)
|
||||
if not user.is_admin:
|
||||
raise HTTPException(403, "Only administrators can create folders")
|
||||
name = _validate_folder_name(payload.name)
|
||||
with get_db() as conn:
|
||||
existing = conn.execute(
|
||||
"SELECT id FROM pst_folders WHERE name = ? COLLATE NOCASE", (name,)
|
||||
).fetchone()
|
||||
if existing:
|
||||
raise HTTPException(409, "A folder with that name already exists")
|
||||
folder_id = uuid.uuid4().hex
|
||||
conn.execute(
|
||||
"""INSERT INTO pst_folders (id, name, created_at, created_by)
|
||||
VALUES (?, ?, ?, ?)""",
|
||||
(folder_id, name, datetime.now(timezone.utc).isoformat(), user.id),
|
||||
)
|
||||
conn.commit()
|
||||
return {"id": folder_id, "name": name, "file_count": 0}
|
||||
|
||||
|
||||
@app.put("/api/folders/{folder_id}")
|
||||
def rename_folder(folder_id: str, payload: FolderRename,
|
||||
user: User = Depends(require_app_user)):
|
||||
if not user.is_admin:
|
||||
raise HTTPException(403, "Only administrators can rename folders")
|
||||
name = _validate_folder_name(payload.name)
|
||||
with get_db() as conn:
|
||||
if not conn.execute("SELECT id FROM pst_folders WHERE id = ?",
|
||||
(folder_id,)).fetchone():
|
||||
raise HTTPException(404, "Folder not found")
|
||||
clash = conn.execute(
|
||||
"SELECT id FROM pst_folders WHERE name = ? COLLATE NOCASE AND id != ?",
|
||||
(name, folder_id),
|
||||
).fetchone()
|
||||
if clash:
|
||||
raise HTTPException(409, "A folder with that name already exists")
|
||||
conn.execute("UPDATE pst_folders SET name = ? WHERE id = ?",
|
||||
(name, folder_id))
|
||||
conn.commit()
|
||||
return {"id": folder_id, "name": name}
|
||||
|
||||
|
||||
@app.delete("/api/folders/{folder_id}")
|
||||
def delete_folder(folder_id: str, user: User = Depends(require_app_user)):
|
||||
"""Delete a folder. Any PSTs inside it become uncategorized (folder_id = NULL)."""
|
||||
if not user.is_admin:
|
||||
raise HTTPException(403, "Only administrators can delete folders")
|
||||
with get_db() as conn:
|
||||
if not conn.execute("SELECT id FROM pst_folders WHERE id = ?",
|
||||
(folder_id,)).fetchone():
|
||||
raise HTTPException(404, "Folder not found")
|
||||
conn.execute("UPDATE pst_files SET folder_id = NULL WHERE folder_id = ?",
|
||||
(folder_id,))
|
||||
conn.execute("DELETE FROM folder_permissions WHERE folder_id = ?",
|
||||
(folder_id,))
|
||||
conn.execute("DELETE FROM pst_folders WHERE id = ?", (folder_id,))
|
||||
conn.commit()
|
||||
return {"deleted": folder_id}
|
||||
|
||||
|
||||
@app.put("/api/pst-files/{pst_id}/folder")
|
||||
def assign_pst_to_folder(pst_id: str, payload: FolderAssign,
|
||||
user: User = Depends(require_app_user)):
|
||||
# Check the user can see the PST's current location AND the destination
|
||||
if not user_can_see_pst(user, pst_id):
|
||||
raise HTTPException(404, "PST not found")
|
||||
if payload.folder_id and not user_can_see_folder(user, payload.folder_id):
|
||||
raise HTTPException(403, "You do not have access to the target folder")
|
||||
with get_db() as conn:
|
||||
if not conn.execute("SELECT id FROM pst_files WHERE id = ?",
|
||||
(pst_id,)).fetchone():
|
||||
raise HTTPException(404, "PST not found")
|
||||
if payload.folder_id:
|
||||
if not conn.execute("SELECT id FROM pst_folders WHERE id = ?",
|
||||
(payload.folder_id,)).fetchone():
|
||||
raise HTTPException(404, "Folder not found")
|
||||
conn.execute("UPDATE pst_files SET folder_id = ? WHERE id = ?",
|
||||
(payload.folder_id, pst_id))
|
||||
conn.commit()
|
||||
return {"pst_id": pst_id, "folder_id": payload.folder_id}
|
||||
|
||||
|
||||
@app.post("/api/pst-files")
|
||||
async def upload_pst(
|
||||
background_tasks: BackgroundTasks,
|
||||
file: UploadFile = File(...),
|
||||
folder_id: Optional[str] = None,
|
||||
user: User = Depends(require_app_user),
|
||||
):
|
||||
if not file.filename:
|
||||
raise HTTPException(400, "No filename provided")
|
||||
if not file.filename.lower().endswith((".pst", ".ost")):
|
||||
raise HTTPException(400, "Only .pst and .ost files are supported")
|
||||
|
||||
# Validate folder if provided (and verify user has access to it)
|
||||
if folder_id:
|
||||
with get_db() as conn:
|
||||
exists = conn.execute("SELECT id FROM pst_folders WHERE id = ?",
|
||||
(folder_id,)).fetchone()
|
||||
if not exists:
|
||||
raise HTTPException(400, "Target folder does not exist")
|
||||
if not user_can_see_folder(user, folder_id):
|
||||
raise HTTPException(403, "You do not have access to the target folder")
|
||||
|
||||
pst_id = uuid.uuid4().hex
|
||||
stored_name = f"{pst_id}_{file.filename}"
|
||||
stored_path = UPLOAD_DIR / stored_name
|
||||
|
||||
size = 0
|
||||
async with aiofiles.open(stored_path, "wb") as out:
|
||||
while chunk := await file.read(1024 * 1024):
|
||||
await out.write(chunk)
|
||||
size += len(chunk)
|
||||
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"""INSERT INTO pst_files
|
||||
(id, filename, original_name, uploaded_at, size_bytes, status,
|
||||
uploaded_by, folder_id)
|
||||
VALUES (?, ?, ?, ?, ?, 'indexing', ?, ?)""",
|
||||
(pst_id, stored_name, file.filename,
|
||||
datetime.now(timezone.utc).isoformat(), size, user.id, folder_id),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
background_tasks.add_task(asyncio.to_thread, index_pst_file, pst_id, stored_path)
|
||||
return {"id": pst_id, "original_name": file.filename,
|
||||
"status": "indexing", "size_bytes": size, "folder_id": folder_id}
|
||||
|
||||
|
||||
@app.delete("/api/pst-files/{pst_id}")
|
||||
def delete_pst(pst_id: str, user: User = Depends(require_app_user)):
|
||||
if not user_can_see_pst(user, pst_id):
|
||||
raise HTTPException(404, "PST not found")
|
||||
with get_db() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT filename FROM pst_files WHERE id = ?", (pst_id,)
|
||||
).fetchone()
|
||||
if not row:
|
||||
raise HTTPException(404, "PST not found")
|
||||
conn.execute("DELETE FROM messages WHERE pst_id = ?", (pst_id,))
|
||||
conn.execute("DELETE FROM pst_files WHERE id = ?", (pst_id,))
|
||||
conn.commit()
|
||||
try:
|
||||
(UPLOAD_DIR / row["filename"]).unlink(missing_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
return {"deleted": pst_id}
|
||||
|
||||
|
||||
@app.get("/api/pst-files/{pst_id}/folders")
|
||||
def get_folders(pst_id: str, user: User = Depends(require_app_user)):
|
||||
if not user_can_see_pst(user, pst_id):
|
||||
raise HTTPException(404, "PST not found")
|
||||
with get_db() as conn:
|
||||
rows = conn.execute(
|
||||
"""SELECT folder, COUNT(*) as count FROM messages
|
||||
WHERE pst_id = ? GROUP BY folder ORDER BY folder""",
|
||||
(pst_id,),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
def _escape_fts(query: str) -> str:
|
||||
tokens = []
|
||||
for raw in query.split():
|
||||
if re.fullmatch(r"[A-Za-z0-9]+", raw):
|
||||
tokens.append(f'"{raw}"*')
|
||||
else:
|
||||
sanitized = raw.replace('"', '""')
|
||||
tokens.append(f'"{sanitized}"')
|
||||
return " ".join(tokens)
|
||||
|
||||
|
||||
@app.get("/api/pst-files/{pst_id}/messages")
|
||||
def search_messages(
|
||||
pst_id: str,
|
||||
q: Optional[str] = None,
|
||||
folder: Optional[str] = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
user: User = Depends(require_app_user),
|
||||
):
|
||||
if not user_can_see_pst(user, pst_id):
|
||||
raise HTTPException(404, "PST not found")
|
||||
limit = max(1, min(limit, 200))
|
||||
offset = max(0, offset)
|
||||
|
||||
with get_db() as conn:
|
||||
p = conn.execute("SELECT id FROM pst_files WHERE id = ?", (pst_id,)).fetchone()
|
||||
if not p:
|
||||
raise HTTPException(404, "PST not found")
|
||||
|
||||
params = [pst_id]
|
||||
|
||||
if q and q.strip():
|
||||
fts_q = _escape_fts(q.strip())
|
||||
base = """
|
||||
FROM messages m
|
||||
JOIN messages_fts f ON f.rowid = m.id
|
||||
WHERE m.pst_id = ? AND messages_fts MATCH ?
|
||||
"""
|
||||
params.append(fts_q)
|
||||
if folder:
|
||||
base += " AND m.folder = ?"
|
||||
params.append(folder)
|
||||
|
||||
total = conn.execute("SELECT COUNT(*) as c " + base, params).fetchone()["c"]
|
||||
|
||||
data_sql = (
|
||||
"SELECT m.id, m.folder, m.subject, m.sender, m.recipients, m.sent_date, "
|
||||
" m.attachments, snippet(messages_fts, 3, '<mark>', '</mark>', '…', 24) AS snippet "
|
||||
+ base + " ORDER BY rank LIMIT ? OFFSET ?"
|
||||
)
|
||||
rows = conn.execute(data_sql, params + [limit, offset]).fetchall()
|
||||
else:
|
||||
where = ["m.pst_id = ?"]
|
||||
if folder:
|
||||
where.append("m.folder = ?")
|
||||
params.append(folder)
|
||||
where_sql = " AND ".join(where)
|
||||
total = conn.execute(
|
||||
f"SELECT COUNT(*) as c FROM messages m WHERE {where_sql}", params
|
||||
).fetchone()["c"]
|
||||
rows = conn.execute(
|
||||
f"""SELECT m.id, m.folder, m.subject, m.sender, m.recipients,
|
||||
m.sent_date, m.attachments, substr(m.body, 1, 200) AS snippet
|
||||
FROM messages m WHERE {where_sql}
|
||||
ORDER BY m.sent_date DESC LIMIT ? OFFSET ?""",
|
||||
params + [limit, offset],
|
||||
).fetchall()
|
||||
|
||||
return {"total": total, "limit": limit, "offset": offset,
|
||||
"items": [dict(r) for r in rows]}
|
||||
|
||||
|
||||
@app.get("/api/messages/{message_id}")
|
||||
def get_message(message_id: int, user: User = Depends(require_app_user)):
|
||||
with get_db() as conn:
|
||||
row = conn.execute(
|
||||
"""SELECT id, pst_id, folder, subject, sender, recipients,
|
||||
sent_date, body, attachments
|
||||
FROM messages WHERE id = ?""",
|
||||
(message_id,),
|
||||
).fetchone()
|
||||
if not row:
|
||||
raise HTTPException(404, "Message not found")
|
||||
if not user_can_see_pst(user, row["pst_id"]):
|
||||
raise HTTPException(404, "Message not found")
|
||||
return dict(row)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Static frontend (anonymous access — but index.html enforces auth client-side
|
||||
# AND the APIs enforce it server-side, which is what actually matters)
|
||||
# ---------------------------------------------------------------------------
|
||||
if FRONTEND_DIR.exists():
|
||||
app.mount("/", StaticFiles(directory=str(FRONTEND_DIR), html=True), name="frontend")
|
||||
@@ -0,0 +1,513 @@
|
||||
"""
|
||||
auth.py — Shared authentication & encryption module.
|
||||
|
||||
Responsibilities:
|
||||
- Derive or generate a master secret key (persisted to disk with chmod 600)
|
||||
- Provide Fernet-based encryption for user PII at rest
|
||||
- Argon2id password hashing (constant-time verification)
|
||||
- TOTP secret generation, QR provisioning URI, 6-digit code verification
|
||||
- JWT session tokens (signed with master key)
|
||||
- User CRUD on the shared SQLite database
|
||||
- FastAPI dependencies to enforce authentication and admin role
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import os
|
||||
import secrets
|
||||
import sqlite3
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import jwt
|
||||
import pyotp
|
||||
from argon2 import PasswordHasher
|
||||
from argon2.exceptions import VerifyMismatchError, InvalidHashError
|
||||
from cryptography.fernet import Fernet, InvalidToken
|
||||
from cryptography.hazmat.primitives import hashes
|
||||
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
||||
from fastapi import Depends, HTTPException, Request, status
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
DATA_DIR = Path(os.environ.get("DATA_DIR", "/app/data"))
|
||||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
SECRET_KEY_PATH = DATA_DIR / "secret.key"
|
||||
DB_PATH = DATA_DIR / "pst_index.db"
|
||||
|
||||
SESSION_TTL_SECONDS = int(os.environ.get("SESSION_TTL_SECONDS", 60 * 60 * 8)) # 8 hours
|
||||
SESSION_COOKIE_NAME = "pst_session"
|
||||
ADMIN_SESSION_COOKIE_NAME = "pst_admin_session"
|
||||
|
||||
JWT_ALGORITHM = "HS256"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Master secret
|
||||
# ---------------------------------------------------------------------------
|
||||
def _load_or_create_master_secret() -> bytes:
|
||||
"""Return a 32-byte master secret. Created on first run, persisted to disk."""
|
||||
# Allow overriding via env var (useful for orchestration); still persist so
|
||||
# restarts without the env var keep working.
|
||||
env_secret = os.environ.get("MASTER_SECRET")
|
||||
if env_secret:
|
||||
return hashlib.sha256(env_secret.encode("utf-8")).digest()
|
||||
|
||||
if SECRET_KEY_PATH.exists():
|
||||
data = SECRET_KEY_PATH.read_bytes()
|
||||
if len(data) == 32:
|
||||
return data
|
||||
# Legacy or truncated: rewrap
|
||||
return hashlib.sha256(data).digest()
|
||||
|
||||
secret = secrets.token_bytes(32)
|
||||
SECRET_KEY_PATH.write_bytes(secret)
|
||||
try:
|
||||
os.chmod(SECRET_KEY_PATH, 0o600)
|
||||
except Exception:
|
||||
pass
|
||||
return secret
|
||||
|
||||
|
||||
MASTER_SECRET: bytes = _load_or_create_master_secret()
|
||||
|
||||
|
||||
def _derive_subkey(info: bytes, length: int = 32) -> bytes:
|
||||
"""Derive a sub-key from the master secret using HKDF-SHA256."""
|
||||
return HKDF(
|
||||
algorithm=hashes.SHA256(),
|
||||
length=length,
|
||||
salt=b"pst-indexer-salt-v1",
|
||||
info=info,
|
||||
).derive(MASTER_SECRET)
|
||||
|
||||
|
||||
# Key used for encrypting user PII at rest
|
||||
_FERNET_KEY = base64.urlsafe_b64encode(_derive_subkey(b"fernet-user-pii"))
|
||||
_FERNET = Fernet(_FERNET_KEY)
|
||||
|
||||
# Key used for JWT signatures
|
||||
_JWT_KEY = _derive_subkey(b"jwt-sessions")
|
||||
|
||||
# Key used for deterministic username lookup hashes (so we can index without
|
||||
# exposing plaintext usernames in the DB).
|
||||
_USERNAME_HMAC_KEY = _derive_subkey(b"username-lookup")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Encryption helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
def encrypt(plaintext: str) -> str:
|
||||
"""Fernet-encrypt a string; returns base64 ciphertext."""
|
||||
if plaintext is None:
|
||||
return ""
|
||||
return _FERNET.encrypt(plaintext.encode("utf-8")).decode("utf-8")
|
||||
|
||||
|
||||
def decrypt(ciphertext: str) -> str:
|
||||
"""Fernet-decrypt. Returns empty string for empty input."""
|
||||
if not ciphertext:
|
||||
return ""
|
||||
try:
|
||||
return _FERNET.decrypt(ciphertext.encode("utf-8")).decode("utf-8")
|
||||
except InvalidToken:
|
||||
raise ValueError("Failed to decrypt — master key may have changed")
|
||||
|
||||
|
||||
def username_lookup_hash(username: str) -> str:
|
||||
"""Deterministic HMAC-SHA256 over the normalized username, hex-encoded."""
|
||||
normalized = username.strip().lower().encode("utf-8")
|
||||
return hmac.new(_USERNAME_HMAC_KEY, normalized, hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Password hashing
|
||||
# ---------------------------------------------------------------------------
|
||||
_hasher = PasswordHasher(
|
||||
time_cost=3,
|
||||
memory_cost=64 * 1024, # 64 MiB
|
||||
parallelism=2,
|
||||
)
|
||||
|
||||
|
||||
def hash_password(password: str) -> str:
|
||||
return _hasher.hash(password)
|
||||
|
||||
|
||||
def verify_password(stored_hash: str, password: str) -> bool:
|
||||
try:
|
||||
_hasher.verify(stored_hash, password)
|
||||
return True
|
||||
except (VerifyMismatchError, InvalidHashError):
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TOTP (MFA)
|
||||
# ---------------------------------------------------------------------------
|
||||
def generate_totp_secret() -> str:
|
||||
return pyotp.random_base32()
|
||||
|
||||
|
||||
def totp_provisioning_uri(secret: str, username: str, issuer: str = "PST Archive") -> str:
|
||||
return pyotp.TOTP(secret).provisioning_uri(name=username, issuer_name=issuer)
|
||||
|
||||
|
||||
def verify_totp(secret: str, code: str) -> bool:
|
||||
if not secret or not code:
|
||||
return False
|
||||
try:
|
||||
return pyotp.TOTP(secret).verify(code.strip(), valid_window=1)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Database
|
||||
# ---------------------------------------------------------------------------
|
||||
@contextmanager
|
||||
def get_db():
|
||||
conn = sqlite3.connect(str(DB_PATH))
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def init_auth_schema():
|
||||
with get_db() as conn:
|
||||
conn.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
id TEXT PRIMARY KEY,
|
||||
username_hash TEXT NOT NULL UNIQUE,
|
||||
username_enc TEXT NOT NULL,
|
||||
password_hash_enc TEXT NOT NULL,
|
||||
totp_secret_enc TEXT,
|
||||
mfa_enabled INTEGER NOT NULL DEFAULT 0,
|
||||
role TEXT NOT NULL DEFAULT 'user',
|
||||
created_at TEXT NOT NULL,
|
||||
last_login TEXT
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_users_username_hash ON users(username_hash);
|
||||
CREATE INDEX IF NOT EXISTS idx_users_role ON users(role);
|
||||
""")
|
||||
conn.commit()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# User model
|
||||
# ---------------------------------------------------------------------------
|
||||
@dataclass
|
||||
class User:
|
||||
id: str
|
||||
username: str
|
||||
role: str
|
||||
mfa_enabled: bool
|
||||
created_at: str
|
||||
last_login: Optional[str]
|
||||
|
||||
@property
|
||||
def is_admin(self) -> bool:
|
||||
return self.role == "admin"
|
||||
|
||||
def public_dict(self) -> dict:
|
||||
return {
|
||||
"id": self.id,
|
||||
"username": self.username,
|
||||
"role": self.role,
|
||||
"mfa_enabled": self.mfa_enabled,
|
||||
"created_at": self.created_at,
|
||||
"last_login": self.last_login,
|
||||
}
|
||||
|
||||
|
||||
def _row_to_user(row: sqlite3.Row) -> User:
|
||||
return User(
|
||||
id=row["id"],
|
||||
username=decrypt(row["username_enc"]),
|
||||
role=row["role"],
|
||||
mfa_enabled=bool(row["mfa_enabled"]),
|
||||
created_at=row["created_at"],
|
||||
last_login=row["last_login"],
|
||||
)
|
||||
|
||||
|
||||
def count_users() -> int:
|
||||
with get_db() as conn:
|
||||
return conn.execute("SELECT COUNT(*) AS c FROM users").fetchone()["c"]
|
||||
|
||||
|
||||
def count_admins() -> int:
|
||||
with get_db() as conn:
|
||||
return conn.execute("SELECT COUNT(*) AS c FROM users WHERE role='admin'").fetchone()["c"]
|
||||
|
||||
|
||||
def find_user_by_username(username: str) -> Optional[tuple[User, str]]:
|
||||
"""Return (User, decrypted_password_hash) if found, else None."""
|
||||
uh = username_lookup_hash(username)
|
||||
with get_db() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM users WHERE username_hash = ?", (uh,)
|
||||
).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
user = _row_to_user(row)
|
||||
try:
|
||||
pw_hash = decrypt(row["password_hash_enc"])
|
||||
except ValueError:
|
||||
return None
|
||||
return user, pw_hash
|
||||
|
||||
|
||||
def find_user_by_id(user_id: str) -> Optional[User]:
|
||||
with get_db() as conn:
|
||||
row = conn.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
return _row_to_user(row)
|
||||
|
||||
|
||||
def get_totp_secret(user_id: str) -> Optional[str]:
|
||||
with get_db() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT totp_secret_enc FROM users WHERE id = ?", (user_id,)
|
||||
).fetchone()
|
||||
if not row or not row["totp_secret_enc"]:
|
||||
return None
|
||||
try:
|
||||
return decrypt(row["totp_secret_enc"])
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def list_users() -> list[User]:
|
||||
with get_db() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM users ORDER BY created_at ASC"
|
||||
).fetchall()
|
||||
return [_row_to_user(r) for r in rows]
|
||||
|
||||
|
||||
def create_user(username: str, password: str, role: str = "user") -> User:
|
||||
username = username.strip()
|
||||
if not username or len(username) < 3:
|
||||
raise ValueError("Username must be at least 3 characters")
|
||||
if len(username) > 64:
|
||||
raise ValueError("Username too long")
|
||||
if role not in ("user", "admin"):
|
||||
raise ValueError("Invalid role")
|
||||
if len(password) < 8:
|
||||
raise ValueError("Password must be at least 8 characters")
|
||||
|
||||
uh = username_lookup_hash(username)
|
||||
with get_db() as conn:
|
||||
existing = conn.execute(
|
||||
"SELECT id FROM users WHERE username_hash = ?", (uh,)
|
||||
).fetchone()
|
||||
if existing:
|
||||
raise ValueError("Username already taken")
|
||||
|
||||
user_id = secrets.token_hex(16)
|
||||
now = datetime.now(timezone.utc).isoformat()
|
||||
pw_hash = hash_password(password)
|
||||
|
||||
conn.execute(
|
||||
"""INSERT INTO users
|
||||
(id, username_hash, username_enc, password_hash_enc,
|
||||
totp_secret_enc, mfa_enabled, role, created_at, last_login)
|
||||
VALUES (?, ?, ?, ?, NULL, 0, ?, ?, NULL)""",
|
||||
(user_id, uh, encrypt(username), encrypt(pw_hash), role, now),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
return User(
|
||||
id=user_id, username=username, role=role, mfa_enabled=False,
|
||||
created_at=now, last_login=None,
|
||||
)
|
||||
|
||||
|
||||
def delete_user(user_id: str) -> None:
|
||||
with get_db() as conn:
|
||||
# Also clean up any folder permission grants this user had
|
||||
conn.execute("DELETE FROM folder_permissions WHERE user_id = ?", (user_id,))
|
||||
conn.execute("DELETE FROM users WHERE id = ?", (user_id,))
|
||||
conn.commit()
|
||||
|
||||
|
||||
def set_user_role(user_id: str, role: str) -> None:
|
||||
if role not in ("user", "admin"):
|
||||
raise ValueError("Invalid role")
|
||||
with get_db() as conn:
|
||||
conn.execute("UPDATE users SET role = ? WHERE id = ?", (role, user_id))
|
||||
conn.commit()
|
||||
|
||||
|
||||
def update_password(user_id: str, new_password: str) -> None:
|
||||
if len(new_password) < 8:
|
||||
raise ValueError("Password must be at least 8 characters")
|
||||
pw_hash = hash_password(new_password)
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"UPDATE users SET password_hash_enc = ? WHERE id = ?",
|
||||
(encrypt(pw_hash), user_id),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def begin_mfa_enrollment(user_id: str) -> tuple[str, str]:
|
||||
"""Generate a new TOTP secret (not yet enabled) and return (secret, provisioning_uri)."""
|
||||
user = find_user_by_id(user_id)
|
||||
if not user:
|
||||
raise ValueError("User not found")
|
||||
secret = generate_totp_secret()
|
||||
# Stash it but leave mfa_enabled = 0 until confirmed
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"UPDATE users SET totp_secret_enc = ?, mfa_enabled = 0 WHERE id = ?",
|
||||
(encrypt(secret), user_id),
|
||||
)
|
||||
conn.commit()
|
||||
uri = totp_provisioning_uri(secret, user.username)
|
||||
return secret, uri
|
||||
|
||||
|
||||
def confirm_mfa_enrollment(user_id: str, code: str) -> bool:
|
||||
secret = get_totp_secret(user_id)
|
||||
if not secret or not verify_totp(secret, code):
|
||||
return False
|
||||
with get_db() as conn:
|
||||
conn.execute("UPDATE users SET mfa_enabled = 1 WHERE id = ?", (user_id,))
|
||||
conn.commit()
|
||||
return True
|
||||
|
||||
|
||||
def disable_mfa(user_id: str) -> None:
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"UPDATE users SET mfa_enabled = 0, totp_secret_enc = NULL WHERE id = ?",
|
||||
(user_id,),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def record_login(user_id: str) -> None:
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"UPDATE users SET last_login = ? WHERE id = ?",
|
||||
(datetime.now(timezone.utc).isoformat(), user_id),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JWT sessions
|
||||
# ---------------------------------------------------------------------------
|
||||
def create_session_token(user: User, scope: str = "app") -> str:
|
||||
"""scope='app' for main app, scope='admin' for admin panel."""
|
||||
now = datetime.now(timezone.utc)
|
||||
payload = {
|
||||
"sub": user.id,
|
||||
"role": user.role,
|
||||
"scope": scope,
|
||||
"iat": int(now.timestamp()),
|
||||
"exp": int((now + timedelta(seconds=SESSION_TTL_SECONDS)).timestamp()),
|
||||
}
|
||||
return jwt.encode(payload, _JWT_KEY, algorithm=JWT_ALGORITHM)
|
||||
|
||||
|
||||
def verify_session_token(token: str, expected_scope: str) -> Optional[dict]:
|
||||
try:
|
||||
payload = jwt.decode(token, _JWT_KEY, algorithms=[JWT_ALGORITHM])
|
||||
if payload.get("scope") != expected_scope:
|
||||
return None
|
||||
return payload
|
||||
except jwt.PyJWTError:
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FastAPI dependencies
|
||||
# ---------------------------------------------------------------------------
|
||||
def _extract_token(request: Request, cookie_name: str) -> Optional[str]:
|
||||
# Prefer cookie, fall back to Authorization header
|
||||
token = request.cookies.get(cookie_name)
|
||||
if token:
|
||||
return token
|
||||
auth = request.headers.get("authorization", "")
|
||||
if auth.lower().startswith("bearer "):
|
||||
return auth[7:].strip()
|
||||
return None
|
||||
|
||||
|
||||
def require_user(scope: str, cookie_name: str):
|
||||
def _dep(request: Request) -> User:
|
||||
token = _extract_token(request, cookie_name)
|
||||
if not token:
|
||||
raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Not authenticated")
|
||||
payload = verify_session_token(token, expected_scope=scope)
|
||||
if not payload:
|
||||
raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Invalid or expired session")
|
||||
user = find_user_by_id(payload["sub"])
|
||||
if not user:
|
||||
raise HTTPException(status.HTTP_401_UNAUTHORIZED, "User no longer exists")
|
||||
return user
|
||||
return _dep
|
||||
|
||||
|
||||
def require_admin(scope: str, cookie_name: str):
|
||||
base = require_user(scope, cookie_name)
|
||||
def _dep(request: Request) -> User:
|
||||
user = base(request)
|
||||
if not user.is_admin:
|
||||
raise HTTPException(status.HTTP_403_FORBIDDEN, "Admin role required")
|
||||
return user
|
||||
return _dep
|
||||
|
||||
|
||||
# Pre-wired dependencies for the two app scopes
|
||||
require_app_user = require_user("app", SESSION_COOKIE_NAME)
|
||||
require_admin_user = require_admin("admin", ADMIN_SESSION_COOKIE_NAME)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Simple rate limiter (per-IP, per-username) to slow brute-force attempts
|
||||
# ---------------------------------------------------------------------------
|
||||
class RateLimiter:
|
||||
def __init__(self, window_seconds: int = 300, max_attempts: int = 10):
|
||||
self.window = window_seconds
|
||||
self.max = max_attempts
|
||||
self._attempts: dict[str, list[float]] = {}
|
||||
|
||||
def check(self, key: str) -> bool:
|
||||
now = time.time()
|
||||
cutoff = now - self.window
|
||||
arr = [t for t in self._attempts.get(key, []) if t > cutoff]
|
||||
if len(arr) >= self.max:
|
||||
self._attempts[key] = arr
|
||||
return False
|
||||
return True
|
||||
|
||||
def record(self, key: str) -> None:
|
||||
now = time.time()
|
||||
cutoff = now - self.window
|
||||
arr = [t for t in self._attempts.get(key, []) if t > cutoff]
|
||||
arr.append(now)
|
||||
self._attempts[key] = arr
|
||||
|
||||
def reset(self, key: str) -> None:
|
||||
self._attempts.pop(key, None)
|
||||
|
||||
|
||||
login_limiter = RateLimiter(window_seconds=300, max_attempts=10)
|
||||
@@ -0,0 +1,10 @@
|
||||
fastapi==0.115.0
|
||||
uvicorn[standard]==0.32.0
|
||||
python-multipart==0.0.12
|
||||
libpff-python==20231205
|
||||
aiofiles==24.1.0
|
||||
argon2-cffi==23.1.0
|
||||
pyotp==2.9.0
|
||||
cryptography==43.0.1
|
||||
PyJWT==2.9.0
|
||||
qrcode[pil]==7.4.2
|
||||
Reference in New Issue
Block a user