Update shield.py

shield.py

@@ -6,7 +6,6 @@ app = Flask(__name__)
 # --- CONFIGURATION ---
 OPENAI_KEY = "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"
 
-# The data you provided, formatted for a secure Python connection
 TIDB_CONFIG = {
     "host": "gateway01.eu-central-1.prod.aws.tidbcloud.com",
     "port": 4000,
@@ -15,86 +14,63 @@ TIDB_CONFIG = {
     "database": "test",
     "autocommit": True,
     "use_pure": True,
-    "ssl_ca": "/etc/ssl/certs/ca-certificates.crt", # Standard path for Debian/HF
-    "ssl_verify_cert": True
+    "ssl_verify_cert": False
 }
 
 def log_to_tidb(user_id, prompt):
-    """Attempts to log the violation with detailed error reporting."""
-    conn = None
     try:
-        print(f"DEBUG: Attempting to log violation for user {user_id}...")
         conn = mysql.connector.connect(**TIDB_CONFIG)
         cursor = conn.cursor()
-        
         query = "INSERT INTO safety_violations (user_id, prompt_content) VALUES (%s, %s)"
-        values = (str(user_id), str(prompt)[:1000])
-        
-        cursor.execute(query, values)
+        cursor.execute(query, (str(user_id), str(prompt)[:1000]))
         conn.commit()
-        
-        print("✅ SUCCESS: Violation logged to TiDB.")
+        print(f"✅ LOGGED TO TiDB: {user_id}")
         cursor.close()
+        conn.close()
     except Exception as e:
-        print(f"❌ DATABASE ERROR: {str(e)}")
-    finally:
-        if conn and conn.is_connected():
-            conn.close()
+        print(f"❌ TiDB ERROR: {e}")
 
+# --- SERVE THE UI ---
+@app.route('/')
+def home():
+    # Reads the index.html file we created
+    with open('index.html', 'r') as f:
+        return f.read()
+
+# --- API ENDPOINT ---
 @app.route('/v1/chat/completions', methods=['POST'])
-def protect_and_proxy():
+def handle_request():
     data = request.json
     messages = data.get('messages', [])
-    full_text = " ".join([m.get('content', '') for m in messages])
+    user_input = " ".join([m.get('content', '') for m in messages])
+    auth_header = request.headers.get('Authorization', 'Anonymous')
 
-    # 1. Moderation Check via OpenAI
+    # 1. SAFETY CHECK
     try:
-        res = requests.post(
+        mod_res = requests.post(
             "https://api.openai.com/v1/moderations",
             headers={"Authorization": f"Bearer {OPENAI_KEY}"},
-            json={"input": full_text},
-            timeout=10
-        )
-        mod_data = res.json()
-        results = mod_data.get('results', [{}])[0]
-    except Exception as e:
-        print(f"Moderation API Error: {e}")
-        results = {}
+            json={"input": user_input}
+        ).json()
 
-    # 2. Block for CSAM (sexual/minors)
-    if results.get('categories', {}).get('sexual/minors'):
-        user_auth = request.headers.get('Authorization', 'Anonymous')
-        print(f"!!! CSAM DETECTED: {user_auth} !!!")
-        
-        log_to_tidb(user_auth, full_text)
-        
-        return {"error": {"message": "Policy Violation: CSAM is strictly prohibited.", "type": "safety_error"}}, 403
+        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
+    except Exception as e:
+        print(f"Moderation Error: {e}")
 
-    # 3. Forward safe traffic to NewAPI on port 3000
+    # 2. FORWARD TO OPENAI (Using your key)
+    # This powers the chat response in the UI
     try:
-        resp = requests.post("http://127.0.0.1:3000/v1/chat/completions", 
-                             json=data, headers=dict(request.headers), timeout=120)
+        resp = requests.post(
+            "https://api.openai.com/v1/chat/completions",
+            headers={"Authorization": f"Bearer {OPENAI_KEY}"},
+            json=data
+        )
         return Response(resp.content, resp.status_code, resp.headers.items())
     except Exception as e:
-        print(f"Proxy Error: {e}")
-        return {"error": {"message": "Internal Proxy Error", "type": "server_error"}}, 500
-
-@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE'])
-def catch_all(path):
-    resp = requests.request(method=request.method, url=f"http://127.0.0.1:3000/{path}",
-                            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())
+        return {"error": {"message": str(e)}}, 500
 
 if __name__ == '__main__':
-    # Startup Test: Try to connect to TiDB once to verify settings
-    print("Starting Shield... Testing TiDB connection...")
-    try:
-        test_conn = mysql.connector.connect(**TIDB_CONFIG)
-        print("✅ TiDB Connection Test: SUCCESSFUL")
-        test_conn.close()
-    except Exception as e:
-        print(f"⚠️ TiDB Connection Test: FAILED - {e}")
-        print("Tip: Check if your TiDB IP Access List allows all connections (0.0.0.0/0)")
-
-    app.run(host='0.0.0.0', port=7860)
+    app.run(host='0.0.0.0', port=7860)