feat: enable api key check

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import base64
 import requests
 import uuid
+import os
 
 USER_ID = str(uuid.uuid4())
 SERVER_URL = "https://lavague.mithrilsecurity.io"
@@ -17,7 +18,8 @@ title = """
 
 def exec_code_req(url, query):
     headers = {
-        "X-User-ID": USER_ID  # Include the X-User-ID header for authentication
+        "X-User-ID": USER_ID,  # Include the X-User-ID header for authentication
+        "X-API-Key": os.environ['API_KEY'] 
     }
     try:
         response = requests.post(SERVER_URL + "/execute_req", json={"url": url, "query": query}, headers=headers)
@@ -28,23 +30,10 @@ def exec_code_req(url, query):
     except requests.RequestException as e:
         return {"error": str(e)}
 
-def get_html(url):
+def get_screenshot_req(url):
     headers = {
-        "X-User-ID": USER_ID  # Include the X-User-ID header for authentication
-    }
-    try:
-        response = requests.post(SERVER_URL + "/get_html", json={"url": url}, headers=headers)
-        if response.status_code == 200:
-            return response.json()
-        else:
-            return {"error": f"Failed with status code {response.status_code}"}
-    except requests.RequestException as e:
-        return {"error": str(e)}
-
-
-def send_request(url):
-    headers = {
-        "X-User-ID": USER_ID  # Include the X-User-ID header for authentication
+        "X-User-ID": USER_ID,  # Include the X-User-ID header for authentication
+        "X-API-Key": os.environ['API_KEY'] 
     }
     try:
         response = requests.get(SERVER_URL + "/screenshot", params={"url": url}, headers=headers)
@@ -56,7 +45,7 @@ def send_request(url):
         return {"error": str(e)}
 
 def process_url(url):
-    r = send_request(url)
+    r = get_screenshot_req(url)
     f = open("screenshot.png", "wb")
     scr = base64.b64decode(r["result"])
     f.write(scr)
@@ -79,7 +68,7 @@ def exec_code(code, source_nodes, full_code, url, query):
     return output, code, html, status, full_code, url, source_nodes
 
 def update_image_display(img, url):
-    r = send_request(url)
+    r = get_screenshot_req(url)
     f = open("screenshot.png", "wb")
     scr = base64.b64decode(r["result"])
     f.write(scr)