Create shield.py

shield.py

@@ -0,0 +1,65 @@
+import os, requests, mysql.connector
+from flask import Flask, request, Response
+
+app = Flask(__name__)
+
+# YOUR OPENAI PROJECT KEY (Already inserted)
+OPENAI_KEY = "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"
+
+# TiDB Credentials from Environment Variables
+TIDB_CONFIG = {
+    "host": os.getenv("TIDB_HOST"),
+    "port": 4000,
+    "user": os.getenv("TIDB_USER"),
+    "password": os.getenv("TIDB_PASSWORD"),
+    "database": os.getenv("TIDB_DB_NAME", "test"),
+    "ssl_verify_cert": True,
+    "ssl_ca": "/etc/ssl/certs/ca-certificates.crt" # Standard for Debian/HF
+}
+
+def log_to_tidb(user, prompt):
+    try:
+        conn = mysql.connector.connect(**TIDB_CONFIG)
+        cursor = conn.cursor()
+        cursor.execute("INSERT INTO safety_violations (user_id, prompt_content) VALUES (%s, %s)", (user, prompt[:1000]))
+        conn.commit()
+        cursor.close()
+        conn.close()
+    except Exception as e:
+        print(f"TiDB Error: {e}")
+
+@app.route('/v1/chat/completions', methods=['POST'])
+def protect_and_proxy():
+    data = request.json
+    messages = data.get('messages', [])
+    text_to_check = " ".join([m.get('content', '') for m in messages])
+
+    # 1. OpenAI Moderation Check (FREE)
+    res = requests.post(
+        "https://api.openai.com/v1/moderations",
+        headers={"Authorization": f"Bearer {OPENAI_KEY}"},
+        json={"input": text_to_check}
+    ).json()
+    
+    # 2. Block only for 'sexual/minors' (CSAM)
+    results = res['results'][0]
+    if results['categories'].get('sexual/minors'):
+        user_auth = request.headers.get('Authorization', 'Anonymous')
+        log_to_tidb(user_auth, text_to_check)
+        return {"error": {"message": "Policy Violation: CSAM is strictly prohibited.", "type": "safety_error"}}, 403
+
+    # 3. Pass through to NewAPI running on internal port 3000
+    resp = requests.post("http://127.0.0.1:3000/v1/chat/completions", 
+                         json=data, headers=dict(request.headers))
+    return Response(resp.content, resp.status_code, resp.headers.items())
+
+# Proxy all other routes (login, admin, etc.) to the real API
+@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())
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=7860) # HF required port