This commit is contained in:
jpmvaz
2026-09-13 20:07:03 +01:00
commit 5482731e78
37 changed files with 15639 additions and 0 deletions
@@ -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