app.py

import gradio as gr
from huggingface_hub import InferenceClient
from transformers import pipeline
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[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[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}..."}
    ]
    
    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
        stream=True,
        temperature=0.5,  # adjust the temperature to improve response quality
        top_p=None,      # disable top-p filtering
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

def simulate_attack(
    prompt: str,
    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[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": "Simulating a Blackhat AI scenario..."}
    ]
    
    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
        stream=True,
        temperature=0.5,  # adjust the temperature to improve response quality
        top_p=None,      # disable top-p filtering
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

# 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."
    ),
    theme="dark"
)

if __name__ == "__main__":
    demo.launch()