Update shield.py

shield.py

@@ -6,6 +6,7 @@ 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,
@@ -13,18 +14,19 @@ TIDB_CONFIG = {
     "password": "Bxg8rpU27gyH60E0",
     "database": "test",
     "autocommit": True,
-    "ssl_verify_cert": False,  # Bypasses the certificate path issue in HF
-    "use_pure": True           # More stable in Docker environments
+    "use_pure": True,
+    "ssl_ca": "/etc/ssl/certs/ca-certificates.crt", # Standard path for Debian/HF
+    "ssl_verify_cert": True
 }
 
 def log_to_tidb(user_id, prompt):
-    """Force logs the violation into TiDB."""
+    """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()
         
-        # SQL matches the columns you created in your DESCRIBE screenshot
         query = "INSERT INTO safety_violations (user_id, prompt_content) VALUES (%s, %s)"
         values = (str(user_id), str(prompt)[:1000])
         
@@ -33,10 +35,11 @@ def log_to_tidb(user_id, prompt):
         
         print("✅ SUCCESS: Violation logged to TiDB.")
         cursor.close()
-        conn.close()
     except Exception as e:
-        # If it still fails, this will show up in your 'Logs' tab on Hugging Face
         print(f"❌ DATABASE ERROR: {str(e)}")
+    finally:
+        if conn and conn.is_connected():
+            conn.close()
 
 @app.route('/v1/chat/completions', methods=['POST'])
 def protect_and_proxy():
@@ -44,33 +47,38 @@ def protect_and_proxy():
     messages = data.get('messages', [])
     full_text = " ".join([m.get('content', '') for m in messages])
 
-    # 1. Moderation Check (FREE)
-    res = requests.post(
-        "https://api.openai.com/v1/moderations",
-        headers={"Authorization": f"Bearer {OPENAI_KEY}"},
-        json={"input": full_text}
-    ).json()
-    
-    # 2. Block only for 'sexual/minors'
-    results = res.json().get('results', [{}])[0]
+    # 1. Moderation Check via OpenAI
+    try:
+        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 = {}
+
+    # 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} !!!")
         
-        # Trigger the log function
         log_to_tidb(user_auth, full_text)
         
         return {"error": {"message": "Policy Violation: CSAM is strictly prohibited.", "type": "safety_error"}}, 403
 
-    # 3. If safe, pass to the real NewAPI on internal port 3000
+    # 3. Forward safe traffic to NewAPI on port 3000
     try:
         resp = requests.post("http://127.0.0.1:3000/v1/chat/completions", 
-                             json=data, headers=dict(request.headers), timeout=60)
+                             json=data, headers=dict(request.headers), timeout=120)
         return Response(resp.content, resp.status_code, resp.headers.items())
     except Exception as e:
-        return {"error": {"message": f"Proxy error: {str(e)}", "type": "internal_error"}}, 500
+        print(f"Proxy Error: {e}")
+        return {"error": {"message": "Internal Proxy Error", "type": "server_error"}}, 500
 
-# Proxy all other dashboard/admin traffic
 @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}",
@@ -79,4 +87,14 @@ def catch_all(path):
     return Response(resp.content, resp.status_code, resp.headers.items())
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+    # 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)