refactored the code for username/password and added pydantic models and routes for the same

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",
+]
+

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)
+

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 <token>
 
 
@@ -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)
 
 

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()
+

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"])

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",
 ]
 

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")
+

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"]
+__all__ = ["training_router", "models_router", "inference_router", "datasets_router", "auth_router"]

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")
+