Merge pull request #113 from nebulabroadcast/sso-support

SSO

Author: martastain

backend/api/auth/__init__.py

@@ -1,5 +1,14 @@
-__all__ = ["LoginRequest", "LogoutRequest", "SetPasswordRequest"]
+__all__ = [
+    "LoginRequest",
+    "LogoutRequest",
+    "SetPasswordRequest",
+    "SSOLoginRequest",
+    "SSOLoginCallback",
+    "TokenExchangeRequest",
+]
 
 from .login_request import LoginRequest
 from .logout_request import LogoutRequest
 from .set_password_request import SetPasswordRequest
+from .sso import SSOLoginCallback, SSOLoginRequest
+from .token_exchange import TokenExchangeRequest

backend/api/auth/login_request.py

@@ -1,39 +1,14 @@
 import time
 
 from fastapi import Request
-from pydantic import Field
 
 import nebula
 from server.clientinfo import get_real_ip
-from server.models import RequestModel, ResponseModel
+from server.models.login import LoginRequestModel, LoginResponseModel
 from server.request import APIRequest
 from server.session import Session
 
 
-class LoginRequestModel(RequestModel):
-    username: str = Field(
-        ...,
-        title="Username",
-        examples=["admin"],
-        pattern=r"^[a-zA-Z0-9_\-\.]{2,}$",
-    )
-    password: str = Field(
-        ...,
-        title="Password",
-        description="Password in plain text",
-        examples=["Password.123"],
-    )
-
-
-class LoginResponseModel(ResponseModel):
-    access_token: str = Field(
-        ...,
-        title="Access token",
-        description="Access token to be used in Authorization header"
-        "for the subsequent requests",
-    )
-
-
 async def check_failed_login(ip_address: str) -> None:
     banned_until = await nebula.redis.get("banned-ip-until", ip_address)
     if banned_until is None:

backend/api/auth/sso.py

@@ -0,0 +1,59 @@
+from urllib.parse import urlparse
+
+from authlib.integrations.starlette_client import OAuthError
+from fastapi import Request
+from fastapi.responses import RedirectResponse
+
+import nebula
+from server.request import APIRequest
+from server.session import Session
+from server.sso import NebulaSSO
+
+
+class SSOLoginRequest(APIRequest):
+    name = "sso_login"
+    path = "/api/sso/login/{provider}"
+    methods = ["GET"]
+
+    async def handle(self, request: Request, provider: str) -> RedirectResponse:
+        client = NebulaSSO.client(provider)
+
+        referer = request.headers.get("referer")
+        if referer:
+            parsed_url = urlparse(referer)
+            base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
+        else:
+            base_url = "http://localhost:4455"
+
+        # We cannot use request.url_for here because it screws the frontend
+        # dev server proxy
+
+        redirect_uri = f"{base_url}/api/sso/callback/{provider}"
+        nebula.log.debug(f"Redirect URI: {redirect_uri}")
+        return await client.authorize_redirect(request, redirect_uri)
+
+
+class SSOLoginCallback(APIRequest):
+    name = "sso_callback"
+    path = "/api/sso/callback/{provider}"
+    methods = ["GET"]
+
+    async def handle(self, request: Request, provider: str) -> RedirectResponse:
+        client = NebulaSSO.client(provider)
+
+        try:
+            token = await client.authorize_access_token(request)
+        except OAuthError as error:
+            return RedirectResponse(f"/?error={error}")
+        user = token.get("userinfo", {})
+        email = user.get("email")
+
+        if not email:
+            return RedirectResponse("/?error=User email not found")
+
+        try:
+            user = await nebula.User.by_email(email)
+        except nebula.NotFoundException:
+            return RedirectResponse("/?error=User not found")
+        session = await Session.create(user, request, transient=True)
+        return RedirectResponse(f"/?authorize={session.token}")

backend/api/auth/token_exchange.py

@@ -0,0 +1,32 @@
+from fastapi import Request
+
+import nebula
+from server.models.login import LoginResponseModel, TokenExchangeRequestModel
+from server.request import APIRequest
+from server.session import Session
+
+
+class TokenExchangeRequest(APIRequest):
+    """Exachange a transient access token for a normal one
+
+    This request will exchange an access token for a new one.
+    The original access token will be invalidated.
+    """
+
+    name: str = "token-exchange"
+    response_model = LoginResponseModel
+
+    async def handle(
+        self,
+        request: Request,
+        payload: TokenExchangeRequestModel,
+    ) -> LoginResponseModel:
+        session = await Session.check(payload.access_token, request, transient=True)
+        if not session:
+            raise nebula.UnauthorizedException("Invalid token")
+        user_id = session.user["id"]
+        user = await nebula.User.load(user_id)
+        session = await Session.create(user, request)
+        nebula.log.debug(f"{user} token exchanged")
+        await Session.delete(payload.access_token)
+        return LoginResponseModel(access_token=session.token)

backend/api/init/init_request.py

@@ -1,4 +1,4 @@
-from typing import Annotated, Any, get_args
+from typing import Annotated, get_args
 
 import fastapi
 from pydantic import Field
@@ -11,6 +11,7 @@
 from server.dependencies import CurrentUserOptional
 from server.models import ResponseModel, UserModel
 from server.request import APIRequest
+from server.sso import NebulaSSO, SSOOption
 
 from .client_settings import ClientSettingsModel, get_client_settings
 
@@ -63,10 +64,10 @@ class InitResponseModel(ResponseModel):
         ),
     ] = None
 
-    oauth2_options: Annotated[
-        list[dict[str, Any]] | None,
+    sso_options: Annotated[
+        list[SSOOption] | None,
         Field(
-            title="OAuth2 options",
+            title="SSO options",
         ),
     ] = None
 
@@ -104,10 +105,11 @@ async def handle(
                 return InitResponseModel(installed=False)
 
         # Not logged in. Only return motd and oauth2 options.
-        # TODO: return oauth2 options
         if user is None:
+            sso_options = await NebulaSSO.options() or None
             return InitResponseModel(
                 motd=motd,
+                sso_options=sso_options,
                 experimental=nebula.config.enable_experimental or None,
             )
 

backend/nebula/config.py

@@ -41,6 +41,11 @@ class NebulaConfig(BaseModel):
         description="Password hashing method",
     )
 
+    session_secret: str = Field(
+        default_factory=lambda: os.urandom(32).hex(),
+        description="Session secret. MUST be set when the server is scaled",
+    )
+
     max_failed_login_attempts: int = Field(
         10,
         description="Maximum number of failed login attempts before the IP is banned",

backend/nebula/objects/user.py

@@ -88,6 +88,19 @@ async def by_api_key(cls, api_key: str) -> "User":
             raise NotFoundException(f"User with API key {api_key} not found")
         return cls.from_row(row[0])
 
+    @classmethod
+    async def by_email(cls, email: str) -> "User":
+        """Return the user with the given email."""
+        row = await db.fetch(
+            """
+            SELECT meta FROM users WHERE LOWER(meta->>'email') = $1
+            """,
+            email.lower(),
+        )
+        if not row:
+            raise NotFoundException(f"User with email {email} not found")
+        return cls.from_row(row[0])
+
     @classmethod
     async def login(cls, username: str, password: str) -> "User":
         """Return a User instance based on username and password."""

backend/nebula/settings/models.py

@@ -93,6 +93,14 @@ class BaseSystemSettings(SettingsModel):
     )
 
 
+class SSOProvider(SettingsModel):
+    name: str
+    title: str | None = None
+    entrypoint: str
+    client_id: str
+    client_secret: str
+
+
 class SystemSettings(BaseSystemSettings):
     """System settings.
 
@@ -107,6 +115,7 @@ class SystemSettings(BaseSystemSettings):
     upload_storage: int | None = Field(default=None)
     upload_dir: str | None = Field(default=None)
     upload_base_name: str = Field(default="{id}")
+    sso_providers: list[SSOProvider] = Field(default_factory=list)
 
     smtp_host: str | None = Field(
         default=None,

backend/pyproject.toml

@@ -7,11 +7,13 @@ requires-python = ">=3.11,<3.13"
 dependencies = [
     "aiofiles >=24.1.0",
     "asyncpg >=0.29.0",
+    "authlib>=1.5.1",
     "email-validator >=2.1.1",
     "fastapi >=0.115.0",
     "geoip2 >=4.8.0",
     "gunicorn >=22.0.0",
     "httpx >=0.27.2",
+    "itsdangerous>=2.2.0",
     "mistune >=3.0.1",
     "nxtools >=1.6",
     "pydantic >=2.9.2",
@@ -33,6 +35,7 @@ dev = [
     "pytest-asyncio >=0.20.3",
     "ruff >=0.6.8",
     "types-aiofiles >=23.2.0.20240311",
+    "types-authlib>=1.4.0.20241230",
     "types-requests >=2.31.0.20240311",
 ]
 

backend/server/middleware/session.py

@@ -0,0 +1,120 @@
+import json
+import typing
+from base64 import b64decode, b64encode
+
+import itsdangerous
+from itsdangerous.exc import BadSignature
+from starlette.datastructures import MutableHeaders
+from starlette.requests import HTTPConnection
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+import nebula
+from nebula.common import create_hash
+
+
+async def get_session_key() -> str:
+    key_candidate = create_hash()
+
+    query = """
+        INSERT INTO settings (key, value)
+        VALUES ('.session_key', $1)
+        ON CONFLICT (key) DO UPDATE
+        SET value = settings.value
+        RETURNING value
+    """
+
+    res = await nebula.db.fetchrow(query, key_candidate)
+    assert res, "Failed to retrieve session key. This shouldn't happen"
+
+    if res["value"] == key_candidate:
+        nebula.log.info("Created new session key")
+
+    return res["value"]
+
+
+class SessionMiddleware:
+    """Custom session middleware for Nebula.
+
+    The main difference with the default Starlette session middleware is that
+    it loads the session key from the database so it can be shared between
+    multiple replicas of the application.
+    """
+
+    _signer: itsdangerous.Signer | None = None
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        session_cookie: str = "session",
+        max_age: int | None = 14 * 24 * 60 * 60,  # 14 days, in seconds
+        path: str = "/",
+        same_site: typing.Literal["lax", "strict", "none"] = "lax",
+        https_only: bool = False,
+        domain: str | None = None,
+    ) -> None:
+        self.app = app
+        self.session_cookie = session_cookie
+        self.max_age = max_age
+        self.path = path
+        self.security_flags = "httponly; samesite=" + same_site
+        if https_only:  # Secure flag can be used with HTTPS only
+            self.security_flags += "; secure"
+        if domain is not None:
+            self.security_flags += f"; domain={domain}"
+
+    async def get_signer(self) -> itsdangerous.Signer:
+        if self._signer is None:
+            key = await get_session_key()
+            self._signer = itsdangerous.TimestampSigner(key)
+        return self._signer
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] not in ("http", "websocket"):  # pragma: no cover
+            await self.app(scope, receive, send)
+            return
+
+        signer = await self.get_signer()
+
+        connection = HTTPConnection(scope)
+        initial_session_was_empty = True
+
+        if self.session_cookie in connection.cookies:
+            data = connection.cookies[self.session_cookie].encode("utf-8")
+            try:
+                data = signer.unsign(data)
+                scope["session"] = json.loads(b64decode(data))
+                initial_session_was_empty = False
+            except BadSignature:
+                scope["session"] = {}
+        else:
+            scope["session"] = {}
+
+        async def send_wrapper(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                if scope["session"]:
+                    # We have session data to persist.
+                    data = b64encode(json.dumps(scope["session"]).encode("utf-8"))
+                    data = signer.sign(data)
+                    headers = MutableHeaders(scope=message)
+                    header_value = "{session_cookie}={data}; path={path}; {max_age}{security_flags}".format(  # noqa: E501
+                        session_cookie=self.session_cookie,
+                        data=data.decode("utf-8"),
+                        path=self.path,
+                        max_age=f"Max-Age={self.max_age}; " if self.max_age else "",
+                        security_flags=self.security_flags,
+                    )
+                    headers.append("Set-Cookie", header_value)
+                elif not initial_session_was_empty:
+                    # The session has been cleared.
+                    headers = MutableHeaders(scope=message)
+                    header_value = "{session_cookie}={data}; path={path}; {expires}{security_flags}".format(  # noqa: E501
+                        session_cookie=self.session_cookie,
+                        data="null",
+                        path=self.path,
+                        expires="expires=Thu, 01 Jan 1970 00:00:00 GMT; ",
+                        security_flags=self.security_flags,
+                    )
+                    headers.append("Set-Cookie", header_value)
+            await send(message)
+
+        await self.app(scope, receive, send_wrapper)