Update app.py

app.py

@@ -1,38 +1,55 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 from transformers import pipeline
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Tuple
+import re
+
+# Assuming that `InferenceClient` is initialized properly
+client = InferenceClient()
+
+def is_valid_input(input_str: str) -> bool:
+    """
+    Validates the input using regex to prevent malicious patterns.
+    
+    Args:
+        input_str (str): The user's input string to be validated.
+    
+    Returns:
+        bool: True if input is safe; False otherwise.
+    """
+    pattern = r'^[A-Za-z0-9\s,.!?-]*$'  # Allows alphanumeric and some punctuation
+    return bool(re.match(pattern, input_str))
 
 def generate_attack(
     prompt: str,
-    history: List[Any],
+    history: List[Tuple[str, str]],
 ) -> List[str]:
     """
     Simulates a Blackhat AI scenario by generating attack strategies and potential impacts.
     
     Args:
         prompt (str): The user's input to the simulator.
-        history (List[Dict]): The user's message history with timestamps.
+        history (List[Tuple[str, str]]): The user's message history with timestamps, where each tuple contains (user_message, assistant_response).
     
     Returns:
         List[str]: A list of attack responses from the AI.
     """
+    if not is_valid_input(prompt):
+        raise ValueError("Invalid input detected. Please use only alphanumeric characters and allowed punctuation.")
+    
     messages = [
-        {
-            "role": "system",
-            "content": f"Responding to {prompt}..."
-        }
+        {"role": "system", "content": f"Responding to {prompt}..."}
     ]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
 
     messages.append({"role": "user", "content": prompt})
+    
     response = ""
-
     for message in client.chat_completion(
         messages,
         max_tokens=100,  # limit the length of responses
@@ -41,40 +58,39 @@ def generate_attack(
         top_p=None,      # disable top-p filtering
     ):
         token = message.choices[0].delta.content
-
         response += token
         yield response
 
 def simulate_attack(
     prompt: str,
-    history: List[Any],
+    history: List[Tuple[str, str]],
 ) -> List[str]:
     """
     Simulates a Blackhat AI scenario by generating attack strategies and potential impacts.
     
     Args:
         prompt (str): The user's input to the simulator.
-        history (List[Dict]): The user's message history with timestamps.
+        history (List[Tuple[str, str]]): The user's message history with timestamps, where each tuple contains (user_message, assistant_response).
     
     Returns:
         List[str]: A list of attack responses from the AI.
     """
+    if not is_valid_input(prompt):
+        raise ValueError("Invalid input detected. Please use only alphanumeric characters and allowed punctuation.")
+    
     messages = [
-        {
-            "role": "system",
-            "content": f"Simulating a Blackhat AI scenario..."
-        }
+        {"role": "system", "content": "Simulating a Blackhat AI scenario..."}
     ]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
 
     messages.append({"role": "user", "content": prompt})
+    
     response = ""
-
     for message in client.chat_completion(
         messages,
         max_tokens=100,  # limit the length of responses
@@ -83,26 +99,44 @@ def simulate_attack(
         top_p=None,      # disable top-p filtering
     ):
         token = message.choices[0].delta.content
-
         response += token
         yield response
 
-# Define the Gradio ChatInterface with security-focused configuration
-demo = gr.ChatInterface(
-    generate_attack,
-    additional_inputs=[
-        gr.Textbox(
-            value=(
-                "You are an AI simulator for cybersecurity training, designed to generate attack scenarios, analyze their impacts, and suggest countermeasures."
-            ),
-            label="System message",
-        ),
-        # ... (other inputs)
+# Define the Gradio ChatInterface with file upload and security-focused configuration
+def handle_file_upload(file: Any) -> str:
+    """
+    Handles file uploads by extracting contents or performing actions based on the file type.
+    
+    Args:
+        file (Any): The uploaded file.
+    
+    Returns:
+        str: A response or summary of the file content.
+    """
+    if file is None:
+        return "No file uploaded."
+    
+    file_type = file.name.split('.')[-1].lower()
+    if file_type in ['txt', 'json', 'csv']:
+        content = file.read().decode("utf-8")
+        return f"File uploaded: {file.name}, {len(content)} characters."
+    else:
+        return "Unsupported file type. Please upload a .txt, .json, or .csv file."
+
+demo = gr.Interface(
+    fn=generate_attack,
+    inputs=[
+        gr.Textbox(label="User Prompt", placeholder="Enter attack scenario...", lines=2),
+        gr.File(label="Upload File", file_types=["txt", "json", "csv"], file_count="single"),
+        gr.Textbox(value="You are an AI simulator for cybersecurity training, designed to generate attack scenarios, analyze their impacts, and suggest countermeasures.", label="System Message")
     ],
+    outputs="text",
     title="Blackhat AI Simulator",
     description=(
-        "This simulator generates adversarial scenarios, analyzes attack vectors, and provides ethical countermeasures. Use responsibly for cybersecurity training and awareness."
+        "This simulator generates adversarial scenarios, analyzes attack vectors, and provides ethical countermeasures. "
+        "Use responsibly for cybersecurity training and awareness."
     ),
+    theme="dark"
 )
 
 if __name__ == "__main__":