Sso one time (#658)

* Created functions to create one time codes

* Created exchange endpoint

* Renamed method calls and created unit tests

* Removed some comments

* Created check for empty code

* Enabled SSO

* Added logs

Author: IanWearsHat

apps/api/docker-compose.yml

@@ -28,7 +28,7 @@ services:
         environment:
             MONGODB_URI: mongodb://root:example@mongo:27017
             DEPLOYMENT: LOCAL
-            UCI_SSO_ENABLED: "false"
+            UCI_SSO_ENABLED: "true"
 
 volumes:
     mongodb_data_volume: {}

apps/api/src/routers/saml.py

@@ -1,5 +1,7 @@
 import json
 import os
+import secrets
+import time
 from urllib.parse import urlparse
 from functools import lru_cache
 from logging import getLogger
@@ -28,6 +30,9 @@
 ALLOWED_RELAY_HOSTS = {"zothacks.com", "www.zothacks.com", "localhost"}
 
 
+ONE_TIME_CODE_TTL = 5 * 60  # in seconds
+
+
 def _is_valid_relay_state(relay_state: str) -> bool:
     # allow same-origin relative paths
     if relay_state.startswith("/"):
@@ -113,6 +118,80 @@ async def _update_last_login(user: NativeUser) -> None:
     )
 
 
+async def _generate_one_time_code(user: NativeUser) -> str:
+    """Generate a secure one-time code and store it in MongoDB."""
+    code = secrets.token_urlsafe(32)
+    expires_at = time.time() + ONE_TIME_CODE_TTL
+
+    code_data = {
+        "code": code,
+        "user": {
+            "ucinetid": user.ucinetid,
+            "display_name": user.display_name,
+            "email": user.email,
+            "affiliations": user.affiliations,
+        },
+        "expires_at": expires_at,
+        "created_at": time.time(),
+    }
+
+    try:
+        await mongodb_handler.insert(Collection.CODES, code_data)
+        log.info("Generated one-time code")
+        return code
+    except Exception as e:
+        log.error(f"Failed to store one-time code in MongoDB: {e}")
+        raise HTTPException(
+            status.HTTP_500_INTERNAL_SERVER_ERROR, "Failed to generate one-time code"
+        )
+
+
+async def _validate_one_time_code(code: str) -> NativeUser:
+    """Validate the one-time code and return the associated user."""
+    try:
+        if not code:
+            raise HTTPException(status.HTTP_400_BAD_REQUEST, "Invalid code")
+
+        code_data = await mongodb_handler.retrieve_one(Collection.CODES, {"code": code})
+
+        if not code_data:
+            log.info("Code not found")
+            raise HTTPException(status.HTTP_400_BAD_REQUEST, "Invalid code")
+
+        current_time = time.time()
+        if current_time > code_data["expires_at"]:
+            # Remove if expired
+            await mongodb_handler.raw_update_one(
+                Collection.CODES, {"code": code}, {"$unset": {"code": ""}}
+            )
+            raise HTTPException(status.HTTP_400_BAD_REQUEST, "Code expired")
+
+        user_data = code_data["user"]
+        user = NativeUser(
+            ucinetid=user_data["ucinetid"],
+            display_name=user_data["display_name"],
+            email=user_data["email"],
+            affiliations=user_data["affiliations"],
+        )
+
+        # Remove the code after successful validation (single-use)
+        await mongodb_handler.raw_update_one(
+            Collection.CODES, {"code": code}, {"$unset": {"code": ""}}
+        )
+
+        log.info(f"Validated code: {code}")
+
+        return user
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        log.error(f"Failed to validate one-time code: {e}")
+        raise HTTPException(
+            status.HTTP_500_INTERNAL_SERVER_ERROR, "Failed to validate one-time code"
+        )
+
+
 @router.get("/login")
 async def login(req: Request, return_to: str = "/") -> RedirectResponse:
     """Initiate login to SSO identity provider."""
@@ -189,9 +268,18 @@ async def acs(
 
     await _update_last_login(user)
 
-    res = RedirectResponse(relay_state, status_code=status.HTTP_303_SEE_OTHER)
-    issue_user_identity(user, res)
-    return res
+    # Generate one-time code if returning to external site
+    if relay_state.startswith("https://zothacks.com"):
+        log.info("Relay starts with zothacks, generating one-time code")
+        code = await _generate_one_time_code(user)
+        redirect_url = f"{relay_state}?code={code}"
+        return RedirectResponse(redirect_url, status_code=status.HTTP_303_SEE_OTHER)
+    else:
+        # Same-domain redirect: set cookie directly
+        log.info("Relay from irvinehacks, issuing identity")
+        res = RedirectResponse(relay_state, status_code=status.HTTP_303_SEE_OTHER)
+        issue_user_identity(user, res)
+        return res
 
 
 @router.get("/sls")
@@ -214,3 +302,14 @@ async def get_saml_metadata() -> Response:
         raise HTTPException(500, "Could not prepare SP metadata")
 
     return Response(metadata, media_type="application/xml")
+
+
+@router.get("/exchange")
+async def exchange_code(code: str) -> RedirectResponse:
+    """Exchange one-time code for JWT."""
+    log.info(f"Attempting exchange with code: {code}")
+    user = await _validate_one_time_code(code)
+    res = RedirectResponse("/", status_code=status.HTTP_303_SEE_OTHER)
+    log.info("Issuing user identity")
+    issue_user_identity(user, res)
+    return res

apps/api/src/services/mongodb_handler.py

@@ -47,6 +47,7 @@ class Collection(str, Enum):
     SETTINGS = "settings"
     EVENTS = "events"
     EMAILS = "emails"
+    CODES = "codes"
 
 
 def get_database() -> AgnosticDatabase[Any]:

apps/api/tests/test_saml.py

@@ -1,11 +1,15 @@
+import time
 from unittest.mock import ANY, AsyncMock, Mock, patch
 
+import pytest
+from fastapi import HTTPException
 from fastapi.testclient import TestClient
 from onelogin.saml2.auth import OneLogin_Saml2_Auth
 from onelogin.saml2.settings import OneLogin_Saml2_Settings
 
 from auth.user_identity import NativeUser, issue_user_identity
 from routers import saml
+from services.mongodb_handler import Collection
 
 SSO_URL = "https://shib.service.uci.edu/idp/profile/SAML2/Redirect/SSO"
 SAMPLE_SETTINGS = OneLogin_Saml2_Settings(
@@ -103,3 +107,155 @@ def test_saml_acs_succeeds(
     )
     # response sets appropriate JWT cookie for user identity
     assert res.headers["Set-Cookie"].startswith("irvinehacks_auth=ey")
+
+
+# One-time code tests
+@patch("services.mongodb_handler.insert", autospec=True)
+async def test_generate_one_time_code_creates_and_stores_code(
+    mock_mongodb_insert: AsyncMock,
+) -> None:
+    """Test that _generate_one_time_code creates a code and stores it in MongoDB."""
+    user = NativeUser(
+        ucinetid="testuser",
+        display_name="Test User",
+        email="test@uci.edu",
+        affiliations=["student"],
+    )
+
+    mock_mongodb_insert.return_value = "test_code_id"
+
+    code = await saml._generate_one_time_code(user)
+
+    # Verify code is generated (should be a URL-safe string)
+    assert isinstance(code, str)
+    assert len(code) > 0
+
+    # Verify MongoDB insert was called with correct data
+    mock_mongodb_insert.assert_awaited_once()
+    call_args = mock_mongodb_insert.await_args
+    assert call_args is not None
+    assert call_args[0][0] == Collection.CODES  # Collection.CODES
+    code_data = call_args[0][1]  # The data dict
+
+    assert code_data["code"] == code
+    assert code_data["user"]["ucinetid"] == "testuser"
+    assert code_data["user"]["display_name"] == "Test User"
+    assert code_data["user"]["email"] == "test@uci.edu"
+    assert code_data["user"]["affiliations"] == ["student"]
+    assert "expires_at" in code_data
+    assert "created_at" in code_data
+
+
+@patch("services.mongodb_handler.retrieve_one", autospec=True)
+@patch("services.mongodb_handler.raw_update_one", autospec=True)
+async def test_validate_one_time_code_with_valid_code(
+    mock_raw_update_one: AsyncMock,
+    mock_retrieve_one: AsyncMock,
+) -> None:
+    """Test that _validate_one_time_code works with a valid code."""
+    current_time = time.time()
+    code_data = {
+        "code": "valid_code_123",
+        "user": {
+            "ucinetid": "testuser",
+            "display_name": "Test User",
+            "email": "test@uci.edu",
+            "affiliations": ["student"],
+        },
+        "expires_at": current_time + 300,  # 5 minutes from now
+        "created_at": current_time,
+    }
+
+    mock_retrieve_one.return_value = code_data
+    mock_raw_update_one.return_value = True
+
+    user = await saml._validate_one_time_code("valid_code_123")
+
+    # Verify user is reconstructed correctly
+    assert user.ucinetid == "testuser"
+    assert user.display_name == "Test User"
+    assert user.email == "test@uci.edu"
+    assert user.affiliations == ["student"]
+
+    # Verify code was removed after validation
+    mock_raw_update_one.assert_awaited_once_with(
+        Collection.CODES, {"code": "valid_code_123"}, {"$unset": {"code": ""}}
+    )
+
+
+@patch("services.mongodb_handler.retrieve_one", autospec=True)
+async def test_validate_one_time_code_with_invalid_code(
+    mock_retrieve_one: AsyncMock,
+) -> None:
+    """Test that _validate_one_time_code fails with an invalid code."""
+    mock_retrieve_one.return_value = None  # Code not found
+
+    with pytest.raises(HTTPException) as exc_info:
+        await saml._validate_one_time_code("invalid_code")
+
+    assert exc_info.value.status_code == 400
+    assert "Invalid code" in str(exc_info.value.detail)
+
+
+@patch("services.mongodb_handler.retrieve_one", autospec=True)
+@patch("services.mongodb_handler.raw_update_one", autospec=True)
+async def test_validate_one_time_code_with_expired_code(
+    mock_raw_update_one: AsyncMock,
+    mock_retrieve_one: AsyncMock,
+) -> None:
+    """Test that _validate_one_time_code fails with an expired code."""
+    current_time = time.time()
+    code_data = {
+        "code": "expired_code_123",
+        "user": {
+            "ucinetid": "testuser",
+            "display_name": "Test User",
+            "email": "test@uci.edu",
+            "affiliations": ["student"],
+        },
+        "expires_at": current_time - 100,  # Expired 100 seconds ago
+        "created_at": current_time - 400,
+    }
+
+    mock_retrieve_one.return_value = code_data
+    mock_raw_update_one.return_value = True
+
+    with pytest.raises(HTTPException) as exc_info:
+        await saml._validate_one_time_code("expired_code_123")
+
+    assert exc_info.value.status_code == 400
+    assert "Code expired" in str(exc_info.value.detail)
+
+    # Verify expired code was cleaned up
+    mock_raw_update_one.assert_awaited_once_with(
+        Collection.CODES, {"code": "expired_code_123"}, {"$unset": {"code": ""}}
+    )
+
+
+@patch("routers.saml._validate_one_time_code", autospec=True)
+@patch("routers.saml.issue_user_identity", autospec=True)
+def test_exchange_code_with_valid_code(
+    mock_issue_user_identity: Mock,
+    mock_validate_one_time_code: AsyncMock,
+) -> None:
+    """Test that /exchange endpoint works with a valid code."""
+    user = NativeUser(
+        ucinetid="testuser",
+        display_name="Test User",
+        email="test@uci.edu",
+        affiliations=["student"],
+    )
+
+    mock_validate_one_time_code.return_value = user
+    mock_issue_user_identity.side_effect = issue_user_identity
+
+    res = client.get("/exchange?code=valid_code_123")
+
+    assert res.status_code == 303
+    assert res.headers["location"] == "/"
+
+    # Verify user identity was issued
+    mock_issue_user_identity.assert_called_once_with(user, ANY)
+
+    # Verify JWT cookie was set
+    assert res.headers["Set-Cookie"].startswith("irvinehacks_auth=ey")