#!/usr/bin/env python3 """ Bug Bounty Security Chatbot Specialized in Network Security and Web Application Testing Uses fine-tuned language models for security analysis and guidance """ import gradio as gr import torch from transformers import ( AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM, pipeline, BitsAndBytesConfig ) import json import re import os from typing import List, Dict, Optional, Tuple import logging # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class BugBountyChatbot: def __init__(self, model_path: str = None, model_type: str = "classification"): """ Initialize the Bug Bounty Chatbot Args: model_path: Path to the fine-tuned model model_type: Type of model ("classification" or "generation") """ self.model_path = model_path self.model_type = model_type self.tokenizer = None self.model = None self.pipeline = None # Security testing categories and methodologies self.security_categories = { "web_app": [ "SQL Injection", "XSS (Cross-Site Scripting)", "CSRF (Cross-Site Request Forgery)", "Authentication Bypass", "Authorization Flaws", "File Upload Vulnerabilities", "Directory Traversal", "Server-Side Request Forgery (SSRF)", "XML External Entity (XXE)", "Insecure Direct Object References", "Security Misconfiguration" ], "network": [ "Port Scanning", "Service Enumeration", "Network Sniffing", "Man-in-the-Middle", "DNS Spoofing", "ARP Poisoning", "Network Segmentation Bypass", "Wireless Security Testing", "VPN Vulnerabilities", "Firewall Bypass" ], "infrastructure": [ "Server Misconfiguration", "Default Credentials", "Privilege Escalation", "Container Security", "Cloud Security", "API Security", "Database Security" ] } # Common tools and techniques self.security_tools = { "reconnaissance": ["nmap", "masscan", "sublist3r", "amass", "theHarvester"], "web_testing": ["burp_suite", "owasp_zap", "sqlmap", "nikto", "dirb"], "network_testing": ["wireshark", "tcpdump", "netcat", "metasploit", "nmap"], "exploitation": ["metasploit", "exploit_db", "custom_scripts", "burp_suite"] } # Load model if path is provided if model_path and os.path.exists(model_path): self.load_model() def load_model(self): """Load the fine-tuned model and tokenizer""" try: logger.info(f"Loading model from {self.model_path}") if self.model_type == "classification": self.tokenizer = AutoTokenizer.from_pretrained(self.model_path) self.model = AutoModelForSequenceClassification.from_pretrained( self.model_path, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32 ) self.pipeline = pipeline( "text-classification", model=self.model, tokenizer=self.tokenizer, device=0 if torch.cuda.is_available() else -1 ) else: # For generation models (like CodeGemma) self.tokenizer = AutoTokenizer.from_pretrained(self.model_path) self.model = AutoModelForCausalLM.from_pretrained( self.model_path, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto" if torch.cuda.is_available() else None ) self.pipeline = pipeline( "text-generation", model=self.model, tokenizer=self.tokenizer, device=0 if torch.cuda.is_available() else -1 ) logger.info("Model loaded successfully") except Exception as e: logger.error(f"Error loading model: {e}") self.model = None self.tokenizer = None self.pipeline = None def analyze_security_query(self, query: str) -> Dict: """ Analyze a security-related query and provide structured response Args: query: User's security question or request Returns: Dictionary with analysis results """ analysis = { "category": "general", "vulnerability_types": [], "tools_suggested": [], "methodology": [], "risk_level": "medium", "response": "" } query_lower = query.lower() # Categorize the query if any(term in query_lower for term in ["web", "website", "application", "app", "http", "https"]): analysis["category"] = "web_app" analysis["vulnerability_types"] = self.security_categories["web_app"] analysis["tools_suggested"] = self.security_tools["web_testing"] elif any(term in query_lower for term in ["network", "port", "scan", "tcp", "udp", "ip"]): analysis["category"] = "network" analysis["vulnerability_types"] = self.security_categories["network"] analysis["tools_suggested"] = self.security_tools["network_testing"] elif any(term in query_lower for term in ["server", "infrastructure", "cloud", "container"]): analysis["category"] = "infrastructure" analysis["vulnerability_types"] = self.security_categories["infrastructure"] analysis["tools_suggested"] = self.security_tools["exploitation"] # Determine risk level based on keywords high_risk_terms = ["exploit", "bypass", "injection", "privilege", "escalation"] if any(term in query_lower for term in high_risk_terms): analysis["risk_level"] = "high" return analysis def generate_security_response(self, query: str, analysis: Dict) -> str: """ Generate a comprehensive security response based on analysis Args: query: Original user query analysis: Analysis results from analyze_security_query Returns: Formatted response string """ response_parts = [] # Header with category and risk level risk_emoji = {"low": "🟢", "medium": "🟡", "high": "🔴"} response_parts.append( f"## {risk_emoji.get(analysis['risk_level'], '🟡')} Security Analysis - {analysis['category'].title()}" ) # Main response based on query type if "how to" in query.lower() or "method" in query.lower(): response_parts.append("### Methodology:") response_parts.append("1. **Reconnaissance Phase**") response_parts.append(" - Gather information about the target") response_parts.append(" - Identify attack surface") response_parts.append(" - Map network topology") response_parts.append("\n2. **Scanning Phase**") response_parts.append(" - Port scanning and service enumeration") response_parts.append(" - Vulnerability scanning") response_parts.append(" - Web application scanning") response_parts.append("\n3. **Exploitation Phase**") response_parts.append(" - Attempt to exploit identified vulnerabilities") response_parts.append(" - Document findings") response_parts.append(" - Maintain access if required") elif "tool" in query.lower() or "scan" in query.lower(): response_parts.append("### Recommended Tools:") for tool in analysis["tools_suggested"][:5]: # Limit to top 5 response_parts.append(f"- **{tool.replace('_', ' ').title()}**") elif "vulnerability" in query.lower() or "exploit" in query.lower(): response_parts.append("### Common Vulnerabilities:") for vuln in analysis["vulnerability_types"][:5]: # Limit to top 5 response_parts.append(f"- {vuln}") else: # General security guidance response_parts.append("### Security Guidance:") response_parts.append("Based on your query, here are key security considerations:") if analysis["category"] == "web_app": response_parts.append("- Focus on OWASP Top 10 vulnerabilities") response_parts.append("- Test authentication and authorization mechanisms") response_parts.append("- Validate all input parameters") response_parts.append("- Check for insecure direct object references") elif analysis["category"] == "network": response_parts.append("- Perform comprehensive port scanning") response_parts.append("- Analyze network traffic patterns") response_parts.append("- Test network segmentation") response_parts.append("- Verify firewall rules and configurations") elif analysis["category"] == "infrastructure": response_parts.append("- Review server configurations") response_parts.append("- Check for default credentials") response_parts.append("- Analyze privilege levels") response_parts.append("- Test container and cloud security") # Add model-based response if available if self.pipeline and self.model_type == "generation": try: # Create a prompt for the model prompt = f"""<|system|> You are a cybersecurity expert specializing in bug bounty hunting and penetration testing. Provide detailed, actionable security guidance. <|user|> {query} <|assistant|>""" model_response = self.pipeline( prompt, max_length=512, num_return_sequences=1, temperature=0.7, do_sample=True, pad_token_id=self.tokenizer.eos_token_id ) if model_response and len(model_response) > 0: generated_text = model_response[0]['generated_text'] # Extract only the assistant's response if "<|assistant|>" in generated_text: assistant_response = generated_text.split("<|assistant|>")[-1].strip() response_parts.append(f"\n### AI-Generated Insights:\n{assistant_response}") except Exception as e: logger.error(f"Error generating model response: {e}") # Add disclaimer response_parts.append("\n---") response_parts.append("⚠️ **Disclaimer**: This information is for educational and authorized testing purposes only.") response_parts.append("Always ensure you have proper authorization before testing any systems.") return "\n".join(response_parts) def chat(self, message: str, history: List[List[str]]) -> Tuple[str, List[List[str]]]: """ Main chat function for Gradio interface Args: message: User's message history: Chat history Returns: Tuple of (response, updated_history) """ if not message.strip(): return "Please enter a security-related question or request.", history # Analyze the query analysis = self.analyze_security_query(message) # Generate response response = self.generate_security_response(message, analysis) # Update history history.append([message, response]) return "", history def create_chatbot_interface(): """Create and configure the Gradio interface""" # Initialize chatbot with CodeGemma 7B model from Hugging Face Hub chatbot = BugBountyChatbot( model_path="BenjaminKaindu0506/codegemma-7b-bugbounty", model_type="generation", base_model="unsloth/codegemma-7b" ) # Custom CSS for better styling css = """ .gradio-container { font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; } .chat-message { padding: 10px; margin: 5px 0; border-radius: 10px; } .user-message { background-color: #e3f2fd; margin-left: 20%; } .bot-message { background-color: #f5f5f5; margin-right: 20%; } """ # Create Gradio interface with gr.Blocks(css=css, title="Bug Bounty Security Chatbot") as interface: gr.Markdown(""" # 🛡️ Bug Bounty Security Chatbot **Specialized in Network Security and Web Application Testing** This AI-powered chatbot provides expert guidance on: - 🔍 **Reconnaissance techniques** - 🌐 **Web application security testing** - 🔗 **Network security analysis** - ⚡ **Vulnerability assessment** - 🛠️ **Security tool recommendations** Ask me about security testing methodologies, tools, vulnerabilities, or specific attack techniques! """) with gr.Row(): with gr.Column(scale=3): chatbot_interface = gr.Chatbot( label="Security Chat", height=600, show_label=True, container=True, bubble_full_width=False ) with gr.Row(): msg_input = gr.Textbox( placeholder="Ask about security testing, vulnerabilities, tools, or methodologies...", label="Your Security Question", lines=2, scale=4 ) send_btn = gr.Button("Send", variant="primary", scale=1) with gr.Row(): clear_btn = gr.Button("Clear Chat", variant="secondary") example_btn = gr.Button("Load Examples", variant="secondary") with gr.Column(scale=1): gr.Markdown("### 🎯 Quick Examples") examples = [ "How to test for SQL injection vulnerabilities?", "What tools should I use for network reconnaissance?", "How to perform web application security testing?", "What are common authentication bypass techniques?", "How to scan for open ports and services?", "What is the OWASP Top 10 and how to test for them?", "How to perform privilege escalation testing?", "What are the steps for a complete penetration test?" ] example_buttons = [] for example in examples: btn = gr.Button(example, size="sm", variant="outline") example_buttons.append(btn) gr.Markdown("### 🔧 Security Categories") gr.Markdown(""" - **Web Applications**: XSS, SQLi, CSRF, Auth bypass - **Network Security**: Port scanning, traffic analysis - **Infrastructure**: Server configs, privilege escalation - **Cloud Security**: Container security, API testing """) # Event handlers def user_input(message, history): return chatbot.chat(message, history) def load_examples(): return examples # Connect events send_btn.click( user_input, inputs=[msg_input, chatbot_interface], outputs=[msg_input, chatbot_interface] ) msg_input.submit( user_input, inputs=[msg_input, chatbot_interface], outputs=[msg_input, chatbot_interface] ) clear_btn.click( lambda: ([], ""), outputs=[chatbot_interface, msg_input] ) # Example button clicks for i, btn in enumerate(example_buttons): btn.click( lambda x=examples[i]: (x, ""), outputs=[msg_input, chatbot_interface] ).then( user_input, inputs=[msg_input, chatbot_interface], outputs=[msg_input, chatbot_interface] ) return interface def main(): """Main function to run the chatbot""" print("🛡️ Initializing Bug Bounty Security Chatbot...") # Create and launch the interface interface = create_chatbot_interface() print("🚀 Starting chatbot interface...") interface.launch( server_name="0.0.0.0", server_port=7860, share=True, # Enable public sharing for Hugging Face Spaces show_error=True ) if __name__ == "__main__": main()