From daa20821c110d580042d8a0cbc688f2b43e285f2 Mon Sep 17 00:00:00 2001 From: sshah229 Date: Fri, 6 Feb 2026 02:56:54 -0700 Subject: [PATCH] refactored the code for username/password and added pydantic models and routes for the same --- studio/backend/auth/__init__.py | 24 +++++++ studio/backend/auth/auth.db | Bin 0 -> 12288 bytes studio/backend/auth/hashing.py | 40 ++++++++++++ studio/backend/auth/jwt.py | 34 +++++++--- studio/backend/auth/storage.py | 101 ++++++++++++++++++++++++++++++ studio/backend/main.py | 3 +- studio/backend/models/__init__.py | 9 +++ studio/backend/models/auth.py | 22 +++++++ studio/backend/routes/__init__.py | 3 +- studio/backend/routes/auth.py | 88 ++++++++++++++++++++++++++ 10 files changed, 312 insertions(+), 12 deletions(-) create mode 100644 studio/backend/auth/auth.db create mode 100644 studio/backend/auth/hashing.py create mode 100644 studio/backend/auth/storage.py create mode 100644 studio/backend/models/auth.py create mode 100644 studio/backend/routes/auth.py diff --git a/studio/backend/auth/__init__.py b/studio/backend/auth/__init__.py index e69de29bb2..4ea6ea0a8c 100644 --- a/studio/backend/auth/__init__.py +++ b/studio/backend/auth/__init__.py @@ -0,0 +1,24 @@ +""" +Authentication module for JWT-based auth with SQLite storage. +""" +from .jwt import create_access_token, get_current_subject, reload_secret +from .storage import ( + is_initialized, + create_initial_user, + get_user_and_secret, + load_jwt_secret, +) +from .hashing import hash_password, verify_password + +__all__ = [ + "create_access_token", + "get_current_subject", + "reload_secret", + "is_initialized", + "create_initial_user", + "get_user_and_secret", + "load_jwt_secret", + "hash_password", + "verify_password", +] + diff --git a/studio/backend/auth/auth.db b/studio/backend/auth/auth.db new file mode 100644 index 0000000000000000000000000000000000000000..b525e4b09621984252c0620d4efe40cc4fbfc4f2 GIT binary patch literal 12288 zcmeI$L66cv6ae5UYh2uf#T!TFsEL?P+bOLl3+Wo5pj(UV9t@c>9qCGGX}jQZLy!KX zJ^P#VXxYsYcTMDO-pjo7P5V0Oe7EWQQzy&=EaEg`IabkABuVHs#t5NOkq1TIg;W%e z?k2@`JKq;2^m6&VQvHQW$Jgliy82_^3-BBSKmY_l00ck)1V8`;KmY_l00bTgd_9y7 zYxTPH?PJdTS->vy>G(1W(p^z%ySC-o*t1%Wjd!>4vm^XRgdBGUp8dvl@jJKEx7>5w zv(I1L>)9?lXHkGX`^>|`LFaU6&siuN|93a6)FY< z2!H?xfB*=900@8p2!H?xfB*=9z+VEJwp2#7{-&}q8IQubZ^#qcRAfdNCqfvq!E_QB zx=9pDIBUv+>Y6G03ggt2Wt|W~H9=W3pemy7mB~^c-VnBcZ2!H?xfB*=900@8p2!H?xfB*>mUjpTW K+LN2V7QX>ySfjfD literal 0 HcmV?d00001 diff --git a/studio/backend/auth/hashing.py b/studio/backend/auth/hashing.py new file mode 100644 index 0000000000..c5d629a2a2 --- /dev/null +++ b/studio/backend/auth/hashing.py @@ -0,0 +1,40 @@ +""" +Password hashing utilities using PBKDF2. +""" +import hashlib +import hmac +import secrets +from typing import Tuple + + +def hash_password(password: str, salt: str | None = None) -> Tuple[str, str]: + """ + Hash a password using PBKDF2-HMAC-SHA256. + + Returns (salt, hex_hash) tuple. + """ + if salt is None: + salt = secrets.token_hex(16) + dk = hashlib.pbkdf2_hmac( + "sha256", + password.encode("utf-8"), + salt.encode("utf-8"), + 100_000, # 100k iterations + ) + return salt, dk.hex() + + +def verify_password(password: str, salt: str, hashed: str) -> bool: + """ + Verify a password against a stored salt and hash. + + Uses constant-time comparison to prevent timing attacks. + """ + dk = hashlib.pbkdf2_hmac( + "sha256", + password.encode("utf-8"), + salt.encode("utf-8"), + 100_000, + ) + return hmac.compare_digest(dk.hex(), hashed) + diff --git a/studio/backend/auth/jwt.py b/studio/backend/auth/jwt.py index 2637efa821..33fc68e1e9 100644 --- a/studio/backend/auth/jwt.py +++ b/studio/backend/auth/jwt.py @@ -6,15 +6,20 @@ from fastapi import Depends, HTTPException, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from jose import JWTError, jwt +from .storage import load_jwt_secret -# Ephemeral in-memory secret: -# - Generated fresh on each backend process start -# - Never written to disk -# - Not configurable by the user -SECRET_KEY = secrets.token_urlsafe(64) ALGORITHM = "HS256" ACCESS_TOKEN_EXPIRE_MINUTES = 60 +# Load stable secret from SQLite (set during first-time setup) +# This will raise RuntimeError if auth hasn't been initialized yet +try: + SECRET_KEY = load_jwt_secret() +except RuntimeError: + # Fallback: use a temporary secret until setup is complete + # This allows the app to start, but protected routes will fail until setup + SECRET_KEY = secrets.token_urlsafe(64) + security = HTTPBearer() # Reads Authorization: Bearer @@ -23,10 +28,9 @@ def create_access_token( expires_delta: Optional[timedelta] = None, ) -> str: """ - Create a signed JWT for the given subject (e.g. "local-user"). + Create a signed JWT for the given subject (e.g. username). - Tokens are valid only for the lifetime of this process, because the - SECRET_KEY is regenerated each time the backend restarts. + Tokens are valid across restarts because SECRET_KEY is stored in SQLite. """ to_encode = {"sub": subject} expire = datetime.now(UTC) + ( @@ -36,6 +40,16 @@ def create_access_token( return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM) +def reload_secret() -> None: + """ + Reload the JWT secret from SQLite. + + Call this after setup to ensure new tokens use the persistent secret. + """ + global SECRET_KEY + SECRET_KEY = load_jwt_secret() + + async def get_current_subject( credentials: HTTPAuthorizationCredentials = Depends(security), ) -> str: @@ -63,7 +77,7 @@ async def get_current_subject( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid or expired token", ) -token = create_access_token("local-user") -print(token) +# token = create_access_token("local-user") +# print(token) diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py new file mode 100644 index 0000000000..0e2c9d7388 --- /dev/null +++ b/studio/backend/auth/storage.py @@ -0,0 +1,101 @@ +""" +SQLite storage for authentication data (user credentials + JWT secret). +""" +import sqlite3 +from pathlib import Path +from typing import Optional, Tuple + +DB_PATH = Path(__file__).parent / "auth.db" + + +def get_connection() -> sqlite3.Connection: + """Get a connection to the auth database, creating tables if needed.""" + conn = sqlite3.connect(DB_PATH) + conn.row_factory = sqlite3.Row + conn.execute( + """ + CREATE TABLE IF NOT EXISTS auth_user ( + id INTEGER PRIMARY KEY, + username TEXT UNIQUE NOT NULL, + password_salt TEXT NOT NULL, + password_hash TEXT NOT NULL, + jwt_secret TEXT NOT NULL + ); + """ + ) + conn.commit() + return conn + + +def is_initialized() -> bool: + """Check if auth has been set up (user exists in DB).""" + conn = get_connection() + cur = conn.execute("SELECT COUNT(*) AS c FROM auth_user") + row = cur.fetchone() + conn.close() + return bool(row["c"]) + + +def create_initial_user(username: str, password: str, jwt_secret: str) -> None: + """ + Create the initial admin user in the database. + + Raises sqlite3.IntegrityError if username already exists. + """ + from .hashing import hash_password + + salt, pwd_hash = hash_password(password) + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO auth_user (username, password_salt, password_hash, jwt_secret) + VALUES (?, ?, ?, ?) + """, + (username, salt, pwd_hash, jwt_secret), + ) + conn.commit() + finally: + conn.close() + + +def get_user_and_secret(username: str) -> Optional[Tuple[str, str, str]]: + """ + Get user's password salt, hash, and JWT secret. + + Returns (password_salt, password_hash, jwt_secret) or None if user not found. + """ + conn = get_connection() + try: + cur = conn.execute( + """ + SELECT password_salt, password_hash, jwt_secret + FROM auth_user + WHERE username = ? + """, + (username,), + ) + row = cur.fetchone() + if not row: + return None + return row["password_salt"], row["password_hash"], row["jwt_secret"] + finally: + conn.close() + + +def load_jwt_secret() -> str: + """ + Load the JWT secret from the database. + + Raises RuntimeError if auth is not initialized. + """ + conn = get_connection() + try: + cur = conn.execute("SELECT jwt_secret FROM auth_user LIMIT 1") + row = cur.fetchone() + if not row: + raise RuntimeError("Auth is not initialized. Please set up a password first.") + return row["jwt_secret"] + finally: + conn.close() + diff --git a/studio/backend/main.py b/studio/backend/main.py index f335839f8d..f93f6c2820 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -9,7 +9,7 @@ from pathlib import Path from datetime import datetime # Import routers -from routes import training_router, models_router, inference_router, datasets_router +from routes import training_router, models_router, inference_router, datasets_router, auth_router # Create FastAPI app app = FastAPI( @@ -30,6 +30,7 @@ app.add_middleware( # ============ Register API Routes ============ # Register routers +app.include_router(auth_router, prefix="/api/auth", tags=["auth"]) app.include_router(training_router, prefix="/api/train", tags=["training"]) app.include_router(models_router, prefix="/api/models", tags=["models"]) app.include_router(inference_router, prefix="/api/inference", tags=["inference"]) diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index a47da0242c..55e827bd5a 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -12,6 +12,11 @@ from .models import ( LoRAInfo, LoRAScanResponse, ) +from .auth import ( + AuthSetupRequest, + AuthLoginRequest, + AuthStatusResponse, +) __all__ = [ # Training schemas @@ -23,5 +28,9 @@ __all__ = [ "ModelDetails", "LoRAInfo", "LoRAScanResponse", + # Auth schemas + "AuthSetupRequest", + "AuthLoginRequest", + "AuthStatusResponse", ] diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py new file mode 100644 index 0000000000..f71956db58 --- /dev/null +++ b/studio/backend/models/auth.py @@ -0,0 +1,22 @@ +""" +Pydantic schemas for Authentication API +""" +from pydantic import BaseModel, Field + + +class AuthSetupRequest(BaseModel): + """First-time setup: create the initial admin user + password.""" + username: str = Field(..., description="Admin username") + password: str = Field(..., min_length=8, description="Admin password (minimum 8 characters)") + + +class AuthLoginRequest(BaseModel): + """Login payload: username/password to obtain a JWT.""" + username: str = Field(..., description="Username") + password: str = Field(..., description="Password") + + +class AuthStatusResponse(BaseModel): + """Indicate whether auth has been initialized.""" + initialized: bool = Field(..., description="True if auth setup has been completed") + diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index a865794d54..5a16125a64 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -6,5 +6,6 @@ from routes.training import router as training_router from routes.models import router as models_router from routes.inference import router as inference_router from routes.datasets import router as datasets_router +from routes.auth import router as auth_router -__all__ = ["training_router", "models_router", "inference_router", "datasets_router"] \ No newline at end of file +__all__ = ["training_router", "models_router", "inference_router", "datasets_router", "auth_router"] \ No newline at end of file diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py new file mode 100644 index 0000000000..3f21eceeba --- /dev/null +++ b/studio/backend/routes/auth.py @@ -0,0 +1,88 @@ +""" +Authentication API routes +""" +from fastapi import APIRouter, HTTPException, status +import secrets + +from models.auth import ( + AuthSetupRequest, + AuthLoginRequest, + AuthStatusResponse, +) +from models.users import Token +from auth import storage, hashing +from auth.jwt import create_access_token, reload_secret + +router = APIRouter() + + +@router.get("/status", response_model=AuthStatusResponse) +async def auth_status() -> AuthStatusResponse: + """ + Check whether auth has already been initialized. + + - initialized = False -> frontend should show "Set admin password" screen. + - initialized = True -> frontend should show normal login. + """ + return AuthStatusResponse(initialized=storage.is_initialized()) + + +@router.post("/setup", response_model=Token, status_code=status.HTTP_201_CREATED) +async def setup_auth(payload: AuthSetupRequest) -> Token: + """ + First-time setup: create the admin user and a JWT secret. + + Can only be called once. Subsequent calls will return 400. + """ + if storage.is_initialized(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Auth is already initialized.", + ) + + # Generate a strong random JWT secret for this installation + jwt_secret = secrets.token_urlsafe(64) + + # Save username/password hash and secret in SQLite + try: + storage.create_initial_user( + username=payload.username, + password=payload.password, + jwt_secret=jwt_secret, + ) + except Exception as e: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Failed to create user: {str(e)}", + ) + + # Reload JWT secret from DB (so jwt.py picks it up) + reload_secret() + + # Issue a token for the new user + access_token = create_access_token(subject=payload.username) + return Token(access_token=access_token, token_type="bearer") + + +@router.post("/login", response_model=Token) +async def login(payload: AuthLoginRequest) -> Token: + """ + Login with username/password and receive a JWT. + """ + record = storage.get_user_and_secret(payload.username) + if record is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect username or password", + ) + + salt, pwd_hash, _jwt_secret = record + if not hashing.verify_password(payload.password, salt, pwd_hash): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect username or password", + ) + + access_token = create_access_token(subject=payload.username) + return Token(access_token=access_token, token_type="bearer") +