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 0000000000..b525e4b096 Binary files /dev/null and b/studio/backend/auth/auth.db differ 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") +