From ec47515ab07f3a3267ee317527b5ca64f108497f Mon Sep 17 00:00:00 2001
From: Walison Filipe <walisonfilipe@hotmail.com>
Date: Wed, 28 Feb 2024 01:03:54 -0300
Subject: [PATCH] Extend openapi.json with security schemes

---
 dialog_api_key/plugin.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/dialog_api_key/plugin.py b/dialog_api_key/plugin.py
index 872ea2d..6d298e1 100644
--- a/dialog_api_key/plugin.py
+++ b/dialog_api_key/plugin.py
@@ -2,6 +2,7 @@
 
 from decouple import Csv, config
 from fastapi import FastAPI, Request
+from fastapi.openapi.utils import get_openapi
 from fastapi.responses import JSONResponse
 
 
@@ -41,4 +42,35 @@ async def api_key_middleware(request: Request, call_next):
 
 
 def register_plugin(app: FastAPI):
+    def custom_openapi():
+        if app.openapi_schema:
+            return app.openapi_schema
+
+        schema = get_openapi(
+            title=app.title,
+            version=app.version,
+            description=app.description,
+            summary=app.summary,
+            routes=app.routes,
+        )
+        schema["components"]["securitySchemes"] = {
+            "ApiKeyAuth": {
+                "type": "apiKey",
+                "in": "header",
+                "name": settings.header,
+            }
+        }
+        schema["security"] = [
+            {"ApiKeyAuth": []},
+        ]
+        for path in settings.except_paths:
+            try:
+                schema["paths"][path]["security"] = []
+            except KeyError:
+                continue
+
+        app.openapi_schema = schema
+        return schema
+
+    app.openapi = custom_openapi
     app.middleware("http")(api_key_middleware)