feat: add oauth2 to swagger-ui

rubenfonseca · rubenfonseca · commit 9e980eb6232a · 2024-04-17T11:36:36.000+02:00
diff --git a/aws_lambda_powertools/event_handler/api_gateway.py b/aws_lambda_powertools/event_handler/api_gateway.py
@@ -88,6 +88,7 @@
         Tag,
     )
     from aws_lambda_powertools.event_handler.openapi.params import Dependant
+    from aws_lambda_powertools.event_handler.openapi.swagger_ui.oauth2 import OAuth2Config
     from aws_lambda_powertools.event_handler.openapi.types import (
         TypeModelOrEnum,
     )
@@ -1671,6 +1672,9 @@ def enable_swagger(
         swagger_base_url: Optional[str] = None,
         middlewares: Optional[List[Callable[..., Response]]] = None,
         compress: bool = False,
+        security_schemes: Optional[Dict[str, "SecurityScheme"]] = None,
+        security: Optional[List[Dict[str, List[str]]]] = None,
+        oauth2Config: Optional["OAuth2Config"] = None,
     ):
         """
         Returns the OpenAPI schema as a JSON serializable dict
@@ -1705,6 +1709,12 @@ def enable_swagger(
             List of middlewares to be used for the swagger route.
         compress: bool, default = False
             Whether or not to enable gzip compression swagger route.
+        security_schemes: Dict[str, "SecurityScheme"], optional
+            A declaration of the security schemes available to be used in the specification.
+        security: List[Dict[str, List[str]]], optional
+            A declaration of which security mechanisms are applied globally across the API.
+        oauth2Config: OAuth2Config, optional
+            The OAuth2 configuration for the Swagger UI.
         """
         from aws_lambda_powertools.event_handler.openapi.compat import model_json
         from aws_lambda_powertools.event_handler.openapi.models import Server
@@ -1736,6 +1746,8 @@ def swagger_handler():
                 terms_of_service=terms_of_service,
                 contact=contact,
                 license_info=license_info,
+                security_schemes=security_schemes,
+                security=security,
             )
 
             # The .replace('</', '<\\/') part is necessary to prevent a potential issue where the JSON string contains
@@ -1765,6 +1777,7 @@ def swagger_handler():
                 swagger_js,
                 swagger_css,
                 swagger_base_url,
+                oauth2Config,
             )
 
             return Response(
diff --git a/aws_lambda_powertools/event_handler/openapi/swagger_ui/html.py b/aws_lambda_powertools/event_handler/openapi/swagger_ui/html.py
@@ -1,4 +1,16 @@
-def generate_swagger_html(spec: str, path: str, swagger_js: str, swagger_css: str, swagger_base_url: str) -> str:
+from typing import Optional
+
+from aws_lambda_powertools.event_handler.openapi.swagger_ui.oauth2 import OAuth2Config
+
+
+def generate_swagger_html(
+    spec: str,
+    path: str,
+    swagger_js: str,
+    swagger_css: str,
+    swagger_base_url: str,
+    oauth2: Optional[OAuth2Config],
+) -> str:
     """
     Generate Swagger UI HTML page
 
@@ -8,10 +20,14 @@ def generate_swagger_html(spec: str, path: str, swagger_js: str, swagger_css: st
         The OpenAPI spec
     path: str
         The path to the Swagger documentation
-    js_url: str
+    swagger_js: str
         The URL to the Swagger UI JavaScript file
-    css_url: str
+    swagger_css: str
         The URL to the Swagger UI CSS file
+    swagger_base_url: str
+        The base URL for Swagger UI
+    oauth2: OAuth2Config, optional
+        The OAuth2 configuration.
     """
 
     # If Swagger base URL is present, generate HTML content with linked CSS and JavaScript files
@@ -23,6 +39,9 @@ def generate_swagger_html(spec: str, path: str, swagger_js: str, swagger_css: st
         swagger_css_content = f"<style>{swagger_css}</style>"
         swagger_js_content = f"<script>{swagger_js}</script>"
 
+    # Prepare oauth2 config
+    oauth2_content = f"ui.initOAuth({oauth2.json(exclude_none=True, exclude_unset=True)});" if oauth2 else ""
+
     return f"""
 <!DOCTYPE html>
 <html>
@@ -65,6 +84,7 @@ def generate_swagger_html(spec: str, path: str, swagger_js: str, swagger_css: st
 
   var ui = SwaggerUIBundle(swaggerUIOptions)
   ui.specActions.updateUrl('{path}?format=json');
+  {oauth2_content}
 </script>
 </html>
             """.strip()
diff --git a/aws_lambda_powertools/event_handler/openapi/swagger_ui/oauth2.py b/aws_lambda_powertools/event_handler/openapi/swagger_ui/oauth2.py
@@ -0,0 +1,27 @@
+from typing import Dict, Sequence
+
+from pydantic import BaseModel, Field
+
+from aws_lambda_powertools.event_handler.openapi.pydantic_loader import PYDANTIC_V2
+
+
+# Based on https://swagger.io/docs/open-source-tools/swagger-ui/usage/oauth2/
+class OAuth2Config(BaseModel):
+    clientId: str = Field(alias="client_id")
+    realm: str
+    appName: str = Field(alias="app_name")
+    scopes: Sequence[str] = Field(default=[])
+    additionalQueryStringParams: Dict[str, str] = Field(alias="additional_query_string_params", default={})
+    useBasicAuthenticationWithAccessCodeGrant: bool = Field(
+        alias="use_basic_authentication_with_access_code_grant",
+        default=False,
+    )
+    usePkceWithAuthorizationCodeGrant: bool = Field(alias="use_pkce_with_authorization_code_grant", default=False)
+
+    if PYDANTIC_V2:
+        model_config = {"extra": "allow"}
+    else:
+
+        class Config:
+            extra = "allow"
+            allow_population_by_field_name = True
