From ac8b7472d7cba2d7d2e6480932eb2615b12d5603 Mon Sep 17 00:00:00 2001
From: huang qiwei <qiwei_huang@qq.com>
Date: Mon, 5 Aug 2024 12:51:53 +0000
Subject: [PATCH] Add auth middleware and fix a api-key bug

---
 .env.dev            |  4 ++++
 .vscode/launch.json | 24 ++++++++++++++++++++++++
 app.py              | 36 ++++++++++++++++++++++++------------
 requirements.txt    |  3 ++-
 4 files changed, 54 insertions(+), 13 deletions(-)
 create mode 100644 .env.dev
 create mode 100644 .vscode/launch.json

diff --git a/.env.dev b/.env.dev
new file mode 100644
index 0000000..b5cf776
--- /dev/null
+++ b/.env.dev
@@ -0,0 +1,4 @@
+RESOURCE_MAPPER={"gpt-4o": "azureopenai"}
+MODEL_MAPPER={"gpt-4o": "gpt-4o"}
+KEYS_MAPPER={"azureopenai": "somepasskey"}
+AUTHORIZATION=["somepasskey"]
\ No newline at end of file
diff --git a/.vscode/launch.json b/.vscode/launch.json
new file mode 100644
index 0000000..c8927bd
--- /dev/null
+++ b/.vscode/launch.json
@@ -0,0 +1,24 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: Flask",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "flask",
+            "env": {
+                "FLASK_APP": "app.py",
+                "FLASK_DEBUG": "1"
+            },
+            "envFile": "${workspaceFolder}/.env",
+            "args": [
+                "run"
+            ],
+            "jinja": true,
+            "autoStartBrowser": false
+        }
+    ]
+}
\ No newline at end of file
diff --git a/app.py b/app.py
index 7127a21..40d8a59 100644
--- a/app.py
+++ b/app.py
@@ -5,6 +5,8 @@
 import requests
 import json
 
+
+from flask_httpauth import HTTPTokenAuth
 from opentelemetry.instrumentation.flask import FlaskInstrumentor
 
 app = Flask(__name__)
@@ -13,7 +15,12 @@
 
 CORS(app)
 
-default_resource = "default_resource"
+auth = HTTPTokenAuth(scheme='Bearer')
+
+# format: ["token1","token2"]
+# load list to json
+authorization = json.loads(os.getenv('AUTHORIZATION'))
+tokens = {token for token in authorization}
 
 #  format of the resource_mapper is {"deployment name": "openai resource name"}, e.g. {"gpt-4": "azureopenai1"}
 resource_mapper = json.loads(os.getenv('RESOURCE_MAPPER'))
@@ -25,7 +32,20 @@
 resource_keys = json.loads(os.getenv('KEYS_MAPPER'))
 
 
+@auth.verify_token
+def verify_token(token):
+    if token in tokens:
+        return "openAIUser"
+    return None
+
+
+@auth.error_handler
+def unauthorized():
+    return jsonify({'error': 'Unauthorized access'}), 401
+
+
 @app.route('/<path:path>', methods=['OPTIONS', 'POST'])
+@auth.login_required
 def handler(path):
     if request.method == 'OPTIONS':
         return '', 204
@@ -34,12 +54,8 @@ def handler(path):
         return 'Bad Request', 400
 
     body_bytes = request.get_data()
-    auth = request.headers.get('Authorization')
-
-    if "IloveJXY" not in auth:
-        return 'Unauthorized', 401
 
-    deployment = "gpt-4"
+    deployment = "gpt-4o"
     api_version = "2024-02-15-preview"
 
     if path.startswith("//"):
@@ -90,7 +106,7 @@ def handler(path):
 
     headers = {'api-key': resource_keys[resource]}
     for key, value in request.headers.items():
-        if key.lower() != 'authorization' and key.lower() != 'host':
+        if key.lower() != 'authorization' and key.lower() != 'host' and key.lower() != "api-key":
             headers[key] = value
 
     # Stream the request to the target URL
@@ -112,6 +128,7 @@ def generate():
 
 
 @app.route('/v1/models', methods=['GET'])
+@auth.login_required
 def get_models():
     # Example data
     response = {
@@ -142,9 +159,4 @@ def get_models():
 
 
 if __name__ == '__main__':
-    if os.getenv('MODEL_MAPPER') is None or os.getenv('RESOURCE_MAPPER') is None:
-        raise ValueError(
-            "MODEL_MAPPER and RESOURCE_MAPPER environment variables must be set")
-    if os.getenv('KEYS_MAPPER') is None:
-        raise ValueError("KEYS_MAPPER environment variable must be set")
     app.run(host='0.0.0.0', port=8000, debug=True)
diff --git a/requirements.txt b/requirements.txt
index 2b69ed0..4fb8460 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,4 +3,5 @@ flask_cors
 requests
 openai
 opentelemetry-instrumentation-flask
-gunicorn
\ No newline at end of file
+gunicorn
+Flask-HTTPAuth
\ No newline at end of file