Update shield.py

shield.py

@@ -1,10 +1,11 @@
-import os, requests, mysql.connector
+import requests, mysql.connector
 from flask import Flask, request, Response
 
 app = Flask(__name__)
 
 # --- CONFIGURATION ---
-OPENAI_KEY = "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"
+NEWAPI_INTERNAL = "http://127.0.0.1:3000"
+OPENAI_MOD_KEY = "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"
 
 TIDB_CONFIG = {
     "host": "gateway01.eu-central-1.prod.aws.tidbcloud.com",
@@ -14,63 +15,74 @@ TIDB_CONFIG = {
     "database": "test",
     "autocommit": True,
     "use_pure": True,
-    "ssl_verify_cert": False
+    "ssl_ca": "/etc/ssl/certs/ca-certificates.crt",
+    "ssl_verify_cert": True
 }
 
-def log_to_tidb(user_id, prompt):
+def log_violation(user_id, prompt):
     try:
         conn = mysql.connector.connect(**TIDB_CONFIG)
         cursor = conn.cursor()
-        query = "INSERT INTO safety_violations (user_id, prompt_content) VALUES (%s, %s)"
-        cursor.execute(query, (str(user_id), str(prompt)[:1000]))
+        cursor.execute("INSERT INTO safety_violations (user_id, prompt_content) VALUES (%s, %s)", 
+                      (str(user_id), str(prompt)[:1000]))
         conn.commit()
-        print(f"✅ LOGGED TO TiDB: {user_id}")
-        cursor.close()
+        print(f"✅ LOGGED VIOLATION: {user_id}")
         conn.close()
     except Exception as e:
-        print(f"❌ TiDB ERROR: {e}")
+        print(f"❌ DATABASE ERROR: {e}")
 
-# --- SERVE THE UI ---
+# 1. SERVE UI
 @app.route('/')
 def home():
-    # Reads the index.html file we created
     with open('index.html', 'r') as f:
         return f.read()
 
-# --- API ENDPOINT ---
+# 2. PROXY AUTH & MODEL LISTS DIRECTLY (Don't Block These)
+@app.route('/api/<path:subpath>', methods=['GET', 'POST', 'PUT', 'DELETE'])
+def proxy_api(subpath):
+    return forward_request(f"{NEWAPI_INTERNAL}/api/{subpath}")
+
+@app.route('/v1/models', methods=['GET'])
+def proxy_models():
+    return forward_request(f"{NEWAPI_INTERNAL}/v1/models")
+
+# 3. INTERCEPT & PROTECT CHAT
 @app.route('/v1/chat/completions', methods=['POST'])
-def handle_request():
+def protect_chat():
     data = request.json
-    messages = data.get('messages', [])
-    user_input = " ".join([m.get('content', '') for m in messages])
-    auth_header = request.headers.get('Authorization', 'Anonymous')
-
-    # 1. SAFETY CHECK
+    # Extract text from latest message
+    msgs = data.get('messages', [])
+    last_msg = msgs[-1].get('content', '') if msgs else ""
+    
+    # Check Safety (using your OpenAI Key for moderation)
     try:
-        mod_res = requests.post(
-            "https://api.openai.com/v1/moderations",
-            headers={"Authorization": f"Bearer {OPENAI_KEY}"},
-            json={"input": user_input}
-        ).json()
-
-        if mod_res.get('results', [{}])[0].get('categories', {}).get('sexual/minors'):
-            print(f"!!! BLOCKING CSAM REQUEST FROM {auth_header} !!!")
-            log_to_tidb(auth_header, user_input)
-            return {"error": {"message": "Policy Violation: Content blocked by Shield.", "type": "safety_error"}}, 403
+        mod = requests.post("https://api.openai.com/v1/moderations", 
+            headers={"Authorization": f"Bearer {OPENAI_MOD_KEY}"}, 
+            json={"input": last_msg}, timeout=3).json()
+        
+        if mod.get('results', [{}])[0].get('categories', {}).get('sexual/minors'):
+            user_token = request.headers.get('Authorization', 'Anon')
+            log_violation(user_token, last_msg)
+            return {"error": {"message": "Policy Violation: Blocked by Shield.", "type": "safety_error"}}, 403
     except Exception as e:
-        print(f"Moderation Error: {e}")
+        print(f"Moderation Warning: {e}")
+
+    # If Safe -> Forward to NewAPI (using the USER'S token from the UI)
+    return forward_request(f"{NEWAPI_INTERNAL}/v1/chat/completions")
 
-    # 2. FORWARD TO OPENAI (Using your key)
-    # This powers the chat response in the UI
+def forward_request(url):
     try:
-        resp = requests.post(
-            "https://api.openai.com/v1/chat/completions",
-            headers={"Authorization": f"Bearer {OPENAI_KEY}"},
-            json=data
+        resp = requests.request(
+            method=request.method,
+            url=url,
+            headers={k: v for k, v in request.headers if k.lower() != 'host'},
+            data=request.get_data(),
+            allow_redirects=False
         )
         return Response(resp.content, resp.status_code, resp.headers.items())
     except Exception as e:
-        return {"error": {"message": str(e)}}, 500
+        return {"error": str(e)}, 502
 
 if __name__ == '__main__':
+    # Start Shield on 7860
     app.run(host='0.0.0.0', port=7860)