Refactoring reset and session tokens, currently broken.

passkey/fastapi/api_handlers.py

@@ -8,12 +8,12 @@ This module contains all the HTTP API endpoints for:
 - Login/logout functionality
 """
 
-from fastapi import Request, Response
+from fastapi import FastAPI, Request, Response
 
 from .. import aaguid
 from ..db import sql
 from ..util.jwt import refresh_session_token, validate_session_token
-from .session_manager import (
+from .session import (
     clear_session_cookie,
     get_current_user,
     get_session_token_from_bearer,
@@ -22,98 +22,169 @@ from .session_manager import (
 )
 
 
-async def get_user_info(request: Request) -> dict:
-    """Get user information and credentials from session cookie."""
-    try:
-        user = await get_current_user(request)
-        if not user:
-            return {"error": "Not authenticated"}
+def register_api_routes(app: FastAPI):
+    """Register all API routes on the FastAPI app."""
 
-        # Get current session credential ID
-        current_credential_id = None
-        session_token = get_session_token_from_cookie(request)
-        if session_token:
+    @app.post("/auth/user-info")
+    async def api_user_info(request: Request, response: Response):
+        """Get user information and credentials from session cookie."""
+        try:
+            user = await get_current_user(request)
+            if not user:
+                return {"error": "Not authenticated"}
+
+            # Get current session credential ID
+            current_credential_id = None
+            session_token = get_session_token_from_cookie(request)
+            if session_token:
+                token_data = validate_session_token(session_token)
+                if token_data:
+                    current_credential_id = token_data.get("credential_id")
+
+            # Get all credentials for the user
+            credential_ids = await sql.get_user_credentials(user.user_id)
+
+            credentials = []
+            user_aaguids = set()
+
+            for cred_id in credential_ids:
+                stored_cred = await sql.get_credential_by_id(cred_id)
+
+                # Convert AAGUID to string format
+                aaguid_str = str(stored_cred.aaguid)
+                user_aaguids.add(aaguid_str)
+
+                # Check if this is the current session credential
+                is_current_session = current_credential_id == stored_cred.credential_id
+
+                credentials.append(
+                    {
+                        "credential_id": stored_cred.credential_id.hex(),
+                        "aaguid": aaguid_str,
+                        "created_at": stored_cred.created_at.isoformat(),
+                        "last_used": stored_cred.last_used.isoformat()
+                        if stored_cred.last_used
+                        else None,
+                        "last_verified": stored_cred.last_verified.isoformat()
+                        if stored_cred.last_verified
+                        else None,
+                        "sign_count": stored_cred.sign_count,
+                        "is_current_session": is_current_session,
+                    }
+                )
+
+            # Get AAGUID information for only the AAGUIDs that the user has
+            aaguid_info = aaguid.filter(user_aaguids)
+
+            # Sort credentials by creation date (earliest first, most recently created last)
+            credentials.sort(key=lambda cred: cred["created_at"])
+
+            return {
+                "status": "success",
+                "user": {
+                    "user_id": str(user.user_id),
+                    "user_name": user.user_name,
+                    "created_at": user.created_at.isoformat()
+                    if user.created_at
+                    else None,
+                    "last_seen": user.last_seen.isoformat() if user.last_seen else None,
+                    "visits": user.visits,
+                },
+                "credentials": credentials,
+                "aaguid_info": aaguid_info,
+            }
+        except Exception as e:
+            return {"error": f"Failed to get user info: {str(e)}"}
+
+    @app.post("/auth/logout")
+    async def api_logout(response: Response):
+        """Log out the current user by clearing the session cookie."""
+        clear_session_cookie(response)
+        return {"status": "success", "message": "Logged out successfully"}
+
+    @app.post("/auth/set-session")
+    async def api_set_session(request: Request, response: Response):
+        """Set session cookie using JWT token from request body or Authorization header."""
+        try:
+            session_token = await get_session_token_from_bearer(request)
+
+            if not session_token:
+                return {"error": "No session token provided"}
+
+            # Validate the session token
             token_data = validate_session_token(session_token)
-            if token_data:
-                current_credential_id = token_data.get("credential_id")
+            if not token_data:
+                return {"error": "Invalid or expired session token"}
 
-        # Get all credentials for the user
-        credential_ids = await sql.get_user_credentials(user.user_id)
+            # Set the HTTP-only cookie
+            set_session_cookie(response, session_token)
 
-        credentials = []
-        user_aaguids = set()
+            return {
+                "status": "success",
+                "message": "Session cookie set successfully",
+                "user_id": str(token_data["user_id"]),
+            }
 
-        for cred_id in credential_ids:
-            stored_cred = await sql.get_credential_by_id(cred_id)
+        except Exception as e:
+            return {"error": f"Failed to set session: {str(e)}"}
 
-            # Convert AAGUID to string format
-            aaguid_str = str(stored_cred.aaguid)
-            user_aaguids.add(aaguid_str)
+    @app.post("/auth/delete-credential")
+    async def api_delete_credential(request: Request):
+        """Delete a specific credential for the current user."""
+        try:
+            user = await get_current_user(request)
+            if not user:
+                return {"error": "Not authenticated"}
+
+            # Get the credential ID from the request body
+            try:
+                body = await request.json()
+                credential_id = body.get("credential_id")
+                if not credential_id:
+                    return {"error": "credential_id is required"}
+            except Exception:
+                return {"error": "Invalid request body"}
+
+            # Convert credential_id from hex string to bytes
+            try:
+                credential_id_bytes = bytes.fromhex(credential_id)
+            except ValueError:
+                return {"error": "Invalid credential_id format"}
+
+            # First, verify the credential belongs to the current user
+            try:
+                stored_cred = await sql.get_credential_by_id(credential_id_bytes)
+                if stored_cred.user_id != user.user_id:
+                    return {"error": "Credential not found or access denied"}
+            except ValueError:
+                return {"error": "Credential not found"}
 
             # Check if this is the current session credential
-            is_current_session = current_credential_id == stored_cred.credential_id
+            session_token = get_session_token_from_cookie(request)
+            if session_token:
+                token_data = validate_session_token(session_token)
+                if (
+                    token_data
+                    and token_data.get("credential_id") == credential_id_bytes
+                ):
+                    return {"error": "Cannot delete current session credential"}
 
-            credentials.append(
-                {
-                    "credential_id": stored_cred.credential_id.hex(),
-                    "aaguid": aaguid_str,
-                    "created_at": stored_cred.created_at.isoformat(),
-                    "last_used": stored_cred.last_used.isoformat()
-                    if stored_cred.last_used
-                    else None,
-                    "last_verified": stored_cred.last_verified.isoformat()
-                    if stored_cred.last_verified
-                    else None,
-                    "sign_count": stored_cred.sign_count,
-                    "is_current_session": is_current_session,
-                }
-            )
+            # Get user's remaining credentials count
+            remaining_credentials = await sql.get_user_credentials(user.user_id)
+            if len(remaining_credentials) <= 1:
+                return {"error": "Cannot delete last remaining credential"}
 
-        # Get AAGUID information for only the AAGUIDs that the user has
-        aaguid_info = aaguid.filter(user_aaguids)
+            # Delete the credential
+            await sql.delete_user_credential(credential_id_bytes)
 
-        # Sort credentials by creation date (earliest first, most recently created last)
-        credentials.sort(key=lambda cred: cred["created_at"])
+            return {"status": "success", "message": "Credential deleted successfully"}
 
-        return {
-            "status": "success",
-            "user": {
-                "user_id": str(user.user_id),
-                "user_name": user.user_name,
-                "created_at": user.created_at.isoformat() if user.created_at else None,
-                "last_seen": user.last_seen.isoformat() if user.last_seen else None,
-                "visits": user.visits,
-            },
-            "credentials": credentials,
-            "aaguid_info": aaguid_info,
-        }
-    except Exception as e:
-        return {"error": f"Failed to get user info: {str(e)}"}
+        except Exception as e:
+            return {"error": f"Failed to delete credential: {str(e)}"}
 
 
-async def refresh_token(request: Request, response: Response) -> dict:
-    """Refresh the session token."""
-    try:
-        session_token = get_session_token_from_cookie(request)
-        if not session_token:
-            return {"error": "No session token found"}
-
-        # Validate and refresh the token
-        new_token = refresh_session_token(session_token)
-
-        if new_token:
-            set_session_cookie(response, new_token)
-            return {"status": "success", "refreshed": True}
-        else:
-            clear_session_cookie(response)
-            return {"error": "Invalid or expired session token"}
-
-    except Exception as e:
-        return {"error": f"Failed to refresh token: {str(e)}"}
-
-
-async def validate_token(request: Request) -> dict:
-    """Validate a session token and return user info."""
+async def validate_token(request: Request, response: Response) -> dict:
+    """Validate a session token and return user info. Also refreshes the token if valid."""
     try:
         session_token = get_session_token_from_cookie(request)
         if not session_token:
@@ -122,11 +193,18 @@ async def validate_token(request: Request) -> dict:
         # Validate the session token
         token_data = validate_session_token(session_token)
         if not token_data:
+            clear_session_cookie(response)
             return {"error": "Invalid or expired session token"}
 
+        # Refresh the token if valid
+        new_token = refresh_session_token(session_token)
+        if new_token:
+            set_session_cookie(response, new_token)
+
         return {
             "status": "success",
             "valid": True,
+            "refreshed": bool(new_token),
             "user_id": str(token_data["user_id"]),
             "credential_id": token_data["credential_id"].hex(),
             "issued_at": token_data["issued_at"],
@@ -135,86 +213,3 @@ async def validate_token(request: Request) -> dict:
 
     except Exception as e:
         return {"error": f"Failed to validate token: {str(e)}"}
-
-
-async def logout(response: Response) -> dict:
-    """Log out the current user by clearing the session cookie."""
-    clear_session_cookie(response)
-    return {"status": "success", "message": "Logged out successfully"}
-
-
-async def set_session(request: Request, response: Response) -> dict:
-    """Set session cookie using JWT token from request body or Authorization header."""
-    try:
-        session_token = await get_session_token_from_bearer(request)
-
-        if not session_token:
-            return {"error": "No session token provided"}
-
-        # Validate the session token
-        token_data = validate_session_token(session_token)
-        if not token_data:
-            return {"error": "Invalid or expired session token"}
-
-        # Set the HTTP-only cookie
-        set_session_cookie(response, session_token)
-
-        return {
-            "status": "success",
-            "message": "Session cookie set successfully",
-            "user_id": str(token_data["user_id"]),
-        }
-
-    except Exception as e:
-        return {"error": f"Failed to set session: {str(e)}"}
-
-
-async def delete_credential(request: Request) -> dict:
-    """Delete a specific credential for the current user."""
-    try:
-        user = await get_current_user(request)
-        if not user:
-            return {"error": "Not authenticated"}
-
-        # Get the credential ID from the request body
-        try:
-            body = await request.json()
-            credential_id = body.get("credential_id")
-            if not credential_id:
-                return {"error": "credential_id is required"}
-        except Exception:
-            return {"error": "Invalid request body"}
-
-        # Convert credential_id from hex string to bytes
-        try:
-            credential_id_bytes = bytes.fromhex(credential_id)
-        except ValueError:
-            return {"error": "Invalid credential_id format"}
-
-        # First, verify the credential belongs to the current user
-        try:
-            stored_cred = await sql.get_credential_by_id(credential_id_bytes)
-            if stored_cred.user_id != user.user_id:
-                return {"error": "Credential not found or access denied"}
-        except ValueError:
-            return {"error": "Credential not found"}
-
-        # Check if this is the current session credential
-        session_token = get_session_token_from_cookie(request)
-        if session_token:
-            token_data = validate_session_token(session_token)
-            if token_data and token_data.get("credential_id") == credential_id_bytes:
-                return {"error": "Cannot delete current session credential"}
-
-        # Get user's remaining credentials count
-        remaining_credentials = await sql.get_user_credentials(user.user_id)
-        if len(remaining_credentials) <= 1:
-            return {"error": "Cannot delete last remaining credential"}
-
-        # Delete the credential
-        await sql.delete_user_credential(credential_id_bytes)
-
-        return {"status": "success", "message": "Credential deleted successfully"}
-
-    except Exception as e:
-        return {"error": f"Failed to delete credential: {str(e)}"}

passkey/fastapi/main.py

@@ -18,25 +18,18 @@ from fastapi import (
     Request,
     Response,
 )
-from fastapi import (
-    Path as FastAPIPath,
-)
 from fastapi.responses import (
     FileResponse,
-    RedirectResponse,
+    JSONResponse,
 )
 from fastapi.staticfiles import StaticFiles
 
 from ..db import sql
 from .api_handlers import (
-    delete_credential,
-    get_user_info,
-    logout,
-    refresh_token,
-    set_session,
+    register_api_routes,
     validate_token,
 )
-from .reset_handlers import create_device_addition_link, validate_device_addition_token
+from .reset_handlers import register_reset_routes
 from .ws_handlers import ws_app
 
 STATIC_DIR = Path(__file__).parent.parent / "frontend-build"
@@ -48,34 +41,23 @@ async def lifespan(app: FastAPI):
     yield
 
 
-app = FastAPI(title="Passkey Auth", lifespan=lifespan)
+app = FastAPI(lifespan=lifespan)
 
 # Mount the WebSocket subapp
 app.mount("/auth/ws", ws_app)
 
-
-@app.get("/auth/user-info")
-async def api_get_user_info(request: Request):
-    """Get user information and credentials from session cookie."""
-    return await get_user_info(request)
-
-
-@app.post("/auth/refresh-token")
-async def api_refresh_token(request: Request, response: Response):
-    """Refresh the session token."""
-    return await refresh_token(request, response)
-
-
-@app.get("/auth/validate-token")
-async def api_validate_token(request: Request):
-    """Validate a session token and return user info."""
-    return await validate_token(request)
+# Register API routes
+register_api_routes(app)
+register_reset_routes(app)
 
 
 @app.get("/auth/forward-auth")
 async def forward_authentication(request: Request):
     """A verification endpoint to use with Caddy forward_auth or Nginx auth_request."""
-    result = await validate_token(request)
+    # Create a dummy response object for internal validation (we won't use it for cookies)
+    response = Response()
+
+    result = await validate_token(request, response)
     if result.get("status") != "success":
         # Serve the index.html of the authentication app if not authenticated
         return FileResponse(
@@ -91,52 +73,6 @@ async def forward_authentication(request: Request):
     )
 
 
-@app.post("/auth/logout")
-async def api_logout(response: Response):
-    """Log out the current user by clearing the session cookie."""
-    return await logout(response)
-
-
-@app.post("/auth/set-session")
-async def api_set_session(request: Request, response: Response):
-    """Set session cookie using JWT token from request body or Authorization header."""
-    return await set_session(request, response)
-
-
-@app.post("/auth/delete-credential")
-async def api_delete_credential(request: Request):
-    """Delete a specific credential for the current user."""
-    return await delete_credential(request)
-
-
-@app.post("/auth/create-device-link")
-async def api_create_device_link(request: Request):
-    """Create a device addition link for the authenticated user."""
-    return await create_device_addition_link(request)
-
-
-@app.post("/auth/validate-device-token")
-async def api_validate_device_token(request: Request):
-    """Validate a device addition token."""
-    return await validate_device_addition_token(request)
-
-
-@app.get("/auth/{passphrase}")
-async def reset_authentication(
-    passphrase: str = FastAPIPath(pattern=r"^\w+(\.\w+){2,}$"),
-):
-    response = RedirectResponse(url="/", status_code=303)
-    response.set_cookie(
-        key="auth-token",
-        value=passphrase,
-        httponly=False,
-        secure=True,
-        samesite="strict",
-        max_age=2,
-    )
-    return response
-
-
 # Serve static files
 app.mount(
     "/auth/assets", StaticFiles(directory=STATIC_DIR / "assets"), name="static assets"
@@ -154,7 +90,7 @@ async def redirect_to_index():
 async def spa_handler(request: Request, path: str):
     """Serve the Vue SPA for all routes (except API and static)"""
     if "text/html" not in request.headers.get("accept", ""):
-        return Response(content="Not Found", status_code=404)
+        return JSONResponse({"error": "Not Found"}, status_code=404)
     return FileResponse(STATIC_DIR / "index.html")
 
 

passkey/fastapi/reset_handlers.py

@@ -9,73 +9,87 @@ This module provides endpoints for authenticated users to:
 
 from datetime import datetime, timedelta
 
-from fastapi import Request
+from fastapi import FastAPI, Path, Request
+from fastapi.responses import RedirectResponse
 
 from ..db import sql
 from ..util.passphrase import generate
-from .session_manager import get_current_user
+from .session import get_current_user
 
 
-async def create_device_addition_link(request: Request) -> dict:
-    """Create a device addition link for the authenticated user."""
-    try:
-        # Require authentication
-        user = await get_current_user(request)
-        if not user:
-            return {"error": "Authentication required"}
+def register_reset_routes(app: FastAPI):
+    """Register all device addition/reset routes on the FastAPI app."""
 
-        # Generate a human-readable token
-        token = generate(n=4, sep=".")  # e.g., "able-ocean-forest-dawn"
+    @app.post("/auth/create-device-link")
+    async def api_create_device_link(request: Request):
+        """Create a device addition link for the authenticated user."""
+        try:
+            # Require authentication
+            user = await get_current_user(request)
+            if not user:
+                return {"error": "Authentication required"}
 
-        # Create reset token in database
-        await sql.create_reset_token(user.user_id, token)
+            # Generate a human-readable token
+            token = generate(n=4, sep=".")  # e.g., "able-ocean-forest-dawn"
 
-        # Generate the device addition link with pretty URL
-        addition_link = f"{request.headers.get('origin', '')}/auth/{token}"
+            # Create reset token in database
+            await sql.create_reset_token(user.user_id, token)
 
-        return {
-            "status": "success",
-            "message": "Device addition link generated successfully",
-            "addition_link": addition_link,
-            "expires_in_hours": 24,
-        }
+            # Generate the device addition link with pretty URL
+            addition_link = f"{request.headers.get('origin', '')}/auth/{token}"
 
-    except Exception as e:
-        return {"error": f"Failed to create device addition link: {str(e)}"}
+            return {
+                "status": "success",
+                "message": "Device addition link generated successfully",
+                "addition_link": addition_link,
+                "expires_in_hours": 24,
+            }
 
+        except Exception as e:
+            return {"error": f"Failed to create device addition link: {str(e)}"}
 
-async def validate_device_addition_token(request: Request) -> dict:
-    """Validate a device addition token and return user info."""
-    try:
-        body = await request.json()
-        token = body.get("token")
+    @app.get("/auth/device-session-check")
+    async def check_device_session(request: Request):
+        """Check if the current session is for device addition."""
+        from .session import is_device_addition_session
 
-        if not token:
-            return {"error": "Device addition token is required"}
+        is_device_session = await is_device_addition_session(request)
+        return {"device_addition_session": is_device_session}
 
-        # Get reset token
-        reset_token = await sql.get_reset_token(token)
-        if not reset_token:
-            return {"error": "Invalid or expired device addition token"}
+    @app.get("/auth/{passphrase}")
+    async def reset_authentication(
+        passphrase: str = Path(pattern=r"^\w+(\.\w+){2,}$"),
+    ):
+        try:
+            # Get reset token to validate it exists and get user_id
+            reset_token = await sql.get_reset_token(passphrase)
+            if not reset_token:
+                # Token doesn't exist, redirect to home
+                return RedirectResponse(url="/", status_code=303)
 
-        # Check if token is expired (24 hours)
-        expiry_time = reset_token.created_at + timedelta(hours=24)
-        if datetime.now() > expiry_time:
-            return {"error": "Device addition token has expired"}
+            # Check if token is expired (24 hours)
+            expiry_time = reset_token.created_at + timedelta(hours=24)
+            if datetime.now() > expiry_time:
+                # Token expired, clean it up and redirect to home
+                await sql.delete_reset_token(passphrase)
+                return RedirectResponse(url="/", status_code=303)
 
-        # Get user info
-        user = await sql.get_user_by_id(reset_token.user_id)
+            # Create a device addition session token for the user
+            from ..util.jwt import create_device_addition_token
 
-        return {
-            "status": "success",
-            "valid": True,
-            "user_id": str(user.user_id),
-            "user_name": user.user_name,
-            "token": token,
-        }
+            session_token = create_device_addition_token(reset_token.user_id)
 
-    except Exception as e:
-        return {"error": f"Failed to validate device addition token: {str(e)}"}
+            # Create response and set session cookie
+            response = RedirectResponse(url="/auth/", status_code=303)
+            from .session import set_session_cookie
+
+            set_session_cookie(response, session_token)
+
+            return response
+
+        except Exception:
+            # On any error, redirect to home
+            return RedirectResponse(url="/", status_code=303)
 
 
 async def use_device_addition_token(token: str) -> dict:

passkey/fastapi/session.py

@@ -96,3 +96,53 @@ async def get_user_from_cookie_string(cookie_header: str) -> UUID | None:
         return None
 
     return token_data["user_id"]
+
+
+async def is_device_addition_session(request: Request) -> bool:
+    """Check if the current session is for device addition."""
+    session_token = request.cookies.get(COOKIE_NAME)
+    if not session_token:
+        return False
+
+    token_data = validate_session_token(session_token)
+    if not token_data:
+        return False
+
+    return token_data.get("device_addition", False)
+
+
+async def get_device_addition_user_id(request: Request) -> UUID | None:
+    """Get user ID from device addition session."""
+    session_token = request.cookies.get(COOKIE_NAME)
+    if not session_token:
+        return None
+
+    token_data = validate_session_token(session_token)
+    if not token_data or not token_data.get("device_addition"):
+        return None
+
+    return token_data.get("user_id")
+
+
+async def get_device_addition_user_id_from_cookie(cookie_header: str) -> UUID | None:
+    """Parse cookie header and return user ID if valid device addition session exists."""
+    if not cookie_header:
+        return None
+
+    # Parse cookies from header (simple implementation)
+    cookies = {}
+    for cookie in cookie_header.split(";"):
+        cookie = cookie.strip()
+        if "=" in cookie:
+            name, value = cookie.split("=", 1)
+            cookies[name] = value
+
+    session_token = cookies.get(COOKIE_NAME)
+    if not session_token:
+        return None
+
+    token_data = validate_session_token(session_token)
+    if not token_data or not token_data.get("device_addition"):
+        return None
+
+    return token_data["user_id"]

passkey/fastapi/ws_handlers.py

@@ -20,7 +20,7 @@ from ..db import sql
 from ..db.sql import User
 from ..sansio import Passkey
 from ..util.jwt import create_session_token
-from .session_manager import get_user_from_cookie_string
+from .session import get_user_from_cookie_string
 
 # Create a FastAPI subapp for WebSocket endpoints
 ws_app = FastAPI()
@@ -181,6 +181,56 @@ async def websocket_add_device_credential(ws: WebSocket, token: str):
         await ws.send_json({"error": "Internal Server Error"})
 
 
+@ws_app.websocket("/add_device_credential_session")
+async def websocket_add_device_credential_session(ws: WebSocket):
+    """Add a new credential for an existing user via device addition session."""
+    await ws.accept()
+    origin = ws.headers.get("origin")
+    try:
+        # Get device addition user ID from session cookie
+        cookie_header = ws.headers.get("cookie", "")
+        from .session import get_device_addition_user_id_from_cookie
+
+        user_id = await get_device_addition_user_id_from_cookie(cookie_header)
+
+        if not user_id:
+            await ws.send_json({"error": "No valid device addition session found"})
+            return
+
+        # Get user information
+        user = await sql.get_user_by_id(user_id)
+        if not user:
+            await ws.send_json({"error": "User not found"})
+            return
+
+        # WebAuthn registration
+        # Fetch challenge IDs for the user
+        challenge_ids = await sql.get_user_credentials(user_id)
+
+        credential = await register_chat(
+            ws, user_id, user.user_name, challenge_ids, origin=origin
+        )
+
+        # Store the new credential in the database
+        await sql.create_credential_for_user(credential)
+
+        await ws.send_json(
+            {
+                "status": "success",
+                "user_id": str(user_id),
+                "credential_id": credential.credential_id.hex(),
+                "message": "New credential added successfully via device addition session",
+            }
+        )
+    except ValueError as e:
+        await ws.send_json({"error": str(e)})
+    except WebSocketDisconnect:
+        pass
+    except Exception:
+        logging.exception("Internal Server Error")
+        await ws.send_json({"error": "Internal Server Error"})
+
+
 @ws_app.websocket("/authenticate")
 async def websocket_authenticate(ws: WebSocket):
     await ws.accept()