app.py

import os
import gradio as gr
import requests

# Load Groq API Key from environment
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

headers = {
    "Authorization": f"Bearer {GROQ_API_KEY}",
    "Content-Type": "application/json"
}

def initialize_messages():
    return [{
        "role": "system",
        "content": "You are a cybersecurity vulnerability scanner bot. Your role is to assist users by scanning their websites, applications, or code for security vulnerabilities. You will help identify potential threats such as SQL injection, Cross-Site Scripting (XSS), open ports, and other common vulnerabilities. Provide actionable advice on how to patch or mitigate these vulnerabilities and ensure users’ systems are more secure. Offer practical tips for securing data and applications and respond with detailed explanations and step-by-step guidance."
    }]

messages_prmt = initialize_messages()

def customLLMBot(user_input, history):
    global messages_prmt

    messages_prmt.append({"role": "user", "content": user_input})

    data = {
        "model": "llama3-8b-8192",
        "messages": messages_prmt
    }

    # Sending request to the Groq API to get the model response
    response = requests.post("https://api.groq.com/openai/v1/chat/completions", headers=headers, json=data)
    result = response.json()

    # Print the response to inspect the structure
    print("API Response:", result)

    # Check if 'choices' key exists
    if 'choices' in result and len(result['choices']) > 0:
        LLM_reply = result['choices'][0]['message']['content']
    else:
        # If 'choices' is not found, handle it gracefully
        LLM_reply = "Sorry, I couldn't generate a response. Please try again."

    # Append the assistant's response
    messages_prmt.append({"role": "assistant", "content": LLM_reply})

    return LLM_reply

# Gradio Interface setup
iface = gr.ChatInterface(
    fn=customLLMBot,
    chatbot=gr.Chatbot(height=300, type="messages"),  # Use 'messages' for compatibility with OpenAI-style dictionaries
    textbox=gr.Textbox(placeholder="Ask me a question about vulnerabilities"),
    title="Vulnerability Scanner ChatBot",
    description="Chat bot for identifying and resolving common security vulnerabilities like SQL injection, XSS, and more. Ask me about your website or app security, and I'll provide actionable advice to help secure your systems.",
    theme="soft",
    examples=[
        "How can I secure my login form?", 
        "What is SQL injection and how do I prevent it?", 
        "How do I fix an XSS vulnerability?", 
        "Can you scan my website for vulnerabilities?", 
        "What are common web application vulnerabilities?"
    ],
    submit_btn="Send"
)

# Launch the Gradio interface
iface.launch(share=True)