From ab666e9347712e8d1a94830727133d01501cc2e5 Mon Sep 17 00:00:00 2001 From: "deepsource-autofix[bot]" <62050782+deepsource-autofix[bot]@users.noreply.github.com> Date: Sat, 10 Jan 2026 13:26:12 +0000 Subject: [PATCH] style: format code with Black This commit fixes the style issues introduced in 2405618 according to the output from Black. Details: None --- ai_controller.py | 72 +++-- ai_integration.py | 236 ++++++++------ ai_orchestrator.py | 50 +-- ai_wrapper/__init__.py | 2 +- ai_wrapper/llm_engine.py | 104 ++++--- conftest.py | 1 - examples/notifications.py | 43 ++- modules/ai/ai_assistant.py | 3 +- modules/ai/ai_controller.py | 3 +- modules/cve_collector/__init__.py | 2 +- modules/cve_collector/cve_collector.py | 292 +++++++++++------- modules/darkweb/__init__.py | 7 +- modules/darkweb/robin/llm.py | 4 +- modules/darkweb/robin/llm_utils.py | 93 +++--- modules/darkweb/robin/runner.py | 9 +- modules/darkweb/robin/scrape.py | 38 ++- modules/darkweb/robin/search.py | 59 ++-- modules/darkweb/robin/ui.py | 20 +- modules/exploit_generator/__init__.py | 2 +- .../exploit_generator/exploit_generator.py | 157 +++++----- modules/exploit_testing/__init__.py | 2 +- modules/exploit_testing/exploit_testing.py | 223 ++++++------- modules/recon/__init__.py | 2 +- modules/recon/recon_module.py | 57 ++-- modules/reporting/__init__.py | 2 +- modules/reporting/reporting.py | 65 +++- modules/tool_manager/__init__.py | 2 +- modules/tool_manager/tool_manager.py | 211 ++++++------- performance_analysis.py | 187 +++++------ recon_module.py | 234 ++++++++------ screen_control/launcher.py | 72 +++-- screen_control/python/screen_control.py | 49 +-- scripts/analyze_results.py | 100 +++--- scripts/recon_scan.py | 143 +++++---- setup.py | 2 +- tests/test_ai_features.py | 91 +++--- tests/test_notifier.py | 79 +++-- utils/__init__.py | 18 +- utils/cli_utils.py | 45 ++- utils/config_manager.py | 45 +-- utils/context_builder.py | 77 ++--- utils/logger.py | 17 +- utils/notifier.py | 98 +++--- utils/report_generator.py | 93 +++--- utils/session_manager.py | 27 +- utils/stealth_utils.py | 33 +- vulnforge_main.py | 111 ++++--- vulnforge_robin_integration/api/app.py | 13 +- vulnforge_robin_integration/api/dashboard.py | 1 - vulnforge_robin_integration/crypto.py | 9 +- vulnforge_robin_integration/db.py | 40 ++- vulnforge_robin_integration/enricher.py | 25 +- vulnforge_robin_integration/ingest.py | 27 +- vulnforge_robin_integration/models.py | 1 - vulnforge_robin_integration/normalizer.py | 11 +- vulnforge_robin_integration/scoring.py | 11 +- vulnforge_robin_integration/tests/conftest.py | 1 - vulnforge_robin_integration/tests/test_api.py | 10 +- .../tests/test_crypto.py | 1 - .../tests/test_normalizer.py | 1 - .../tests/test_scoring.py | 1 - vulnforge_robin_integration/workers.py | 25 +- 62 files changed, 1952 insertions(+), 1507 deletions(-) diff --git a/ai_controller.py b/ai_controller.py index 43858aa..fa0d495 100644 --- a/ai_controller.py +++ b/ai_controller.py @@ -18,12 +18,13 @@ from utils.context_builder import ContextBuilder from utils.notifier import Notifier + class AIController: """Controls AI operations and decision making.""" def __init__(self, session_dir: str, config_path: str): """Initialize the AI controller. - + Args: session_dir: Directory to store session data config_path: Path to configuration file @@ -35,7 +36,7 @@ def __init__(self, session_dir: str, config_path: str): self.notifier = Notifier(session_dir, config_path) self.report_generator = ReportGenerator(self.session_dir) self.output_format = "all" # Default output format - + # Initialize session data self.session_data = { "start_time": datetime.now().isoformat(), @@ -44,16 +45,18 @@ def __init__(self, session_dir: str, config_path: str): "findings": [], "execution_history": [], "ai_decisions": [], - "errors": [] + "errors": [], } - async def process_query(self, query: str, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + async def process_query( + self, query: str, context: Optional[Dict[str, Any]] = None + ) -> Dict[str, Any]: """Process an AI query. - + Args: query: The query to process context: Optional context data - + Returns: Dictionary containing the response and metadata """ @@ -62,13 +65,13 @@ async def process_query(self, query: str, context: Optional[Dict[str, Any]] = No raise ValueError("Query must be a non-empty string") self.logger.info("Processing AI query: %s", query) - + try: # Prepare context context_str = "" if context: context_str = json.dumps(context, indent=2) - + # Create prompt prompt = f"""You are a security expert. Please provide guidance on the following query: @@ -85,30 +88,30 @@ async def process_query(self, query: str, context: Optional[Dict[str, Any]] = No 5. Safety considerations Provide a detailed response:""" - + # Get AI response response = await self._get_ai_response(prompt) - + # Log response self.logger.info("AI response received") - + return { "answer": response, # Return an empty log list as the logger doesn't expose in-memory logs "logs": [], - "prompt": prompt + "prompt": prompt, } - + except Exception as e: self.logger.error("Error processing query: %s", str(e)) raise async def _get_ai_response(self, prompt: str) -> str: """Get response from AI model. - + Args: prompt: The prompt to send to the AI - + Returns: The AI's response """ @@ -116,19 +119,19 @@ async def _get_ai_response(self, prompt: str) -> str: # TODO: Implement actual AI model call # For now, return a placeholder response return "This is a placeholder response. AI integration pending." - + except Exception as e: self.logger.error("Error getting AI response: %s", str(e)) raise def _generate_report(self) -> Dict[str, str]: """Generate reports for the current session. - + Returns: Dictionary mapping report types to their file paths """ self.logger.info("Generating reports in %s format...", self.output_format) - + try: # Prepare context data context = { @@ -136,7 +139,7 @@ def _generate_report(self) -> Dict[str, str]: "domain": self.session_data.get("target_domain"), "ip": self.session_data.get("target_ip"), "scan_time": self.session_data.get("start_time"), - "duration": self._calculate_duration() + "duration": self._calculate_duration(), }, "modules": self.session_data["modules"], "tools": self.session_data["tools"], @@ -144,45 +147,50 @@ def _generate_report(self) -> Dict[str, str]: "exploits": self.session_data.get("execution_history", []), "defensive_measures": self.session_data.get("defensive_measures", []), "recommendations": self.session_data.get("recommendations", []), - "errors": self.session_data["errors"] + "errors": self.session_data["errors"], } - + # Generate reports - report_paths = self.report_generator.generate_reports(context, self.output_format) - + report_paths = self.report_generator.generate_reports( + context, self.output_format + ) + self.logger.info("Reports generated successfully") return report_paths - + except Exception as e: self.logger.error("Error generating reports: %s", str(e)) raise def _calculate_duration(self) -> str: """Calculate session duration. - + Returns: Formatted duration string """ start_time = datetime.fromisoformat(self.session_data["start_time"]) end_time = datetime.now() duration = end_time - start_time - + hours = duration.seconds // 3600 minutes = (duration.seconds % 3600) // 60 seconds = duration.seconds % 60 - + return f"{hours:02d}:{minutes:02d}:{seconds:02d}" def setup_ai(self) -> bool: """Setup AI environment: check Ollama service and model availability.""" try: - if hasattr(self, 'ollama'): + if hasattr(self, "ollama"): client = self.ollama else: from ai_integration import OllamaClient + client = OllamaClient() if not client.is_available(): - self.logger.error("Ollama service not available. Start with: ollama serve") + self.logger.error( + "Ollama service not available. Start with: ollama serve" + ) return False models = client.list_models() if not models: @@ -190,8 +198,10 @@ def setup_ai(self) -> bool: if not client.pull_model(client.main_model): self.logger.error("Failed to pull default model") return False - self.logger.info("AI setup complete. Available models: %s", [m['name'] for m in models]) + self.logger.info( + "AI setup complete. Available models: %s", [m["name"] for m in models] + ) return True except Exception as e: self.logger.error("Error in setup_ai: %s", e) - return False \ No newline at end of file + return False diff --git a/ai_integration.py b/ai_integration.py index 1ab8678..a3f5b4e 100755 --- a/ai_integration.py +++ b/ai_integration.py @@ -16,21 +16,22 @@ import asyncio import ctypes + class OllamaClient: def __init__(self, base_url: str = "http://localhost:11434"): self.base_url = base_url self.logger = logging.getLogger(__name__) - + # Load configuration from environment - self.main_model = os.getenv("OLLAMA_MAIN_MODEL", "deepseek-coder-v2:16b-lite-base-q4_0") - self.assistant_model = os.getenv("OLLAMA_ASSISTANT_MODEL", "mistral:7b-instruct-v0.2-q4_0") - - self.backup_models = [ - "deepseek-coder:6.7b", - "codellama:7b", - "mistral:7b" - ] - + self.main_model = os.getenv( + "OLLAMA_MAIN_MODEL", "deepseek-coder-v2:16b-lite-base-q4_0" + ) + self.assistant_model = os.getenv( + "OLLAMA_ASSISTANT_MODEL", "mistral:7b-instruct-v0.2-q4_0" + ) + + self.backup_models = ["deepseek-coder:6.7b", "codellama:7b", "mistral:7b"] + def is_available(self) -> bool: """Check if Ollama service is running""" try: @@ -38,97 +39,111 @@ def is_available(self) -> bool: return response.status_code == 200 except (requests.RequestException, requests.Timeout, requests.ConnectionError): return False - + def list_models(self) -> List[Dict]: """List available models""" try: response = requests.get(f"{self.base_url}/api/tags") if response.status_code == 200: - return response.json().get('models', []) - except (requests.RequestException, requests.Timeout, requests.ConnectionError) as e: + return response.json().get("models", []) + except ( + requests.RequestException, + requests.Timeout, + requests.ConnectionError, + ) as e: self.logger.error("Error listing models: %s", e) return [] - + def pull_model(self, model: str) -> bool: """Pull a model if not available""" try: self.logger.info("Pulling model: %s", model) data = {"name": model} - response = requests.post(f"{self.base_url}/api/pull", json=data, stream=True) - + response = requests.post( + f"{self.base_url}/api/pull", json=data, stream=True + ) + for line in response.iter_lines(): if line: try: - status = json.loads(line.decode('utf-8')) - if status.get('status') == 'success': + status = json.loads(line.decode("utf-8")) + if status.get("status") == "success": return True except: continue except Exception as e: self.logger.error("Error pulling model %s: %s", model, e) return False - - def generate(self, prompt: str, model: str = None, system_prompt: str = None) -> Optional[str]: + + def generate( + self, prompt: str, model: str = None, system_prompt: str = None + ) -> Optional[str]: """Generate text using Ollama""" if not model: model = self.get_best_model() - + if not model: self.logger.error("No suitable model available") return None - + try: data = { "model": model, "prompt": prompt, "stream": False, "options": { - "temperature": 0.5, # Lowered for more deterministic planning + "temperature": 0.5, # Lowered for more deterministic planning "top_p": 0.9, "max_tokens": 4096, - "num_ctx": 16384, # Increased context window for large prompts + "num_ctx": 16384, # Increased context window for large prompts "num_thread": 8, - "repeat_penalty": 1.1 - } + "repeat_penalty": 1.1, + }, } - + if system_prompt: data["system"] = system_prompt - - response = requests.post(f"{self.base_url}/api/generate", json=data, timeout=300) # Increased timeout - + + response = requests.post( + f"{self.base_url}/api/generate", json=data, timeout=300 + ) # Increased timeout + if response.status_code == 200: result = response.json() - return result.get('response', '').strip() + return result.get("response", "").strip() else: self.logger.error("Ollama API error: %s", response.status_code) - - except (requests.RequestException, requests.Timeout, requests.ConnectionError) as e: + + except ( + requests.RequestException, + requests.Timeout, + requests.ConnectionError, + ) as e: self.logger.error("Error generating with Ollama: %s", e) - + return None - + def get_best_model(self) -> Optional[str]: """Get the best available model""" - available_models = [m['name'] for m in self.list_models()] - + available_models = [m["name"] for m in self.list_models()] + # Check main model first (deepseek-coder-v2:16b-lite-base-q4_0) if self.main_model in available_models: return self.main_model - + # Check assistant model next (mistral:7b-instruct-v0.2-q4_0) if self.assistant_model in available_models: return self.assistant_model - + # Check backup models for model in self.backup_models: if model in available_models: return model - + # If no preferred models, return first available if available_models: return available_models[0] - + return None @@ -148,20 +163,23 @@ def get_best_model(self) -> Optional[str]: c_parser = None C_PARSER_AVAILABLE = False import logging - logging.getLogger(__name__).warning("C parser library not found. Using Python fallback.") + + logging.getLogger(__name__).warning( + "C parser library not found. Using Python fallback." + ) class AIAnalyzer: def __init__(self, ollama_client: OllamaClient): self.ollama = ollama_client self.logger = logging.getLogger(__name__) - + def analyze_nmap_output(self, nmap_output: str) -> Dict[str, Any]: """Analyze nmap scan results using AI""" system_prompt = """You are a cybersecurity expert analyzing nmap scan results. Identify potential vulnerabilities, interesting services, and security issues. Provide structured analysis in JSON format with severity levels.""" - + prompt = f""" Analyze this nmap scan output and identify potential security issues: @@ -184,28 +202,28 @@ def analyze_nmap_output(self, nmap_output: str) -> Dict[str, Any]: "next_steps": ["recommended follow-up actions"] }} """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) if response: try: return json.loads(response) except json.JSONDecodeError: # Extract JSON if wrapped in markdown - json_match = re.search(r'```json\n(.*?)\n```', response, re.DOTALL) + json_match = re.search(r"```json\n(.*?)\n```", response, re.DOTALL) if json_match: try: return json.loads(json_match.group(1)) except json.JSONDecodeError: pass - + return {"error": "Failed to analyze nmap output", "raw_response": response} - + def generate_exploit_code(self, vulnerability_info: Dict) -> str: """Generate exploit code based on vulnerability information""" system_prompt = """You are a security researcher creating proof-of-concept exploit code. Generate safe, educational exploit code with proper error handling and comments. Include safety warnings and ethical use disclaimers.""" - + prompt = f""" Generate a Python proof-of-concept exploit for this vulnerability: @@ -224,18 +242,18 @@ def generate_exploit_code(self, vulnerability_info: Dict) -> str: Generate complete, working Python code: """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) return response or "# Failed to generate exploit code" - + def analyze_web_response(self, url: str, response_data: Dict) -> Dict[str, Any]: """Analyze web service response for vulnerabilities""" system_prompt = """You are a web application security expert. Analyze HTTP responses for potential vulnerabilities and security issues.""" - - headers = response_data.get('headers', {}) - content = response_data.get('content', '')[:2000] # Limit content length - + + headers = response_data.get("headers", {}) + content = response_data.get("content", "")[:2000] # Limit content length + prompt = f""" Analyze this web service for security issues: @@ -270,26 +288,28 @@ def analyze_web_response(self, url: str, response_data: Dict) -> Dict[str, Any]: "recommendations": ["security recommendations"] }} """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) if response: try: return json.loads(response) except json.JSONDecodeError: - json_match = re.search(r'```json\n(.*?)\n```', response, re.DOTALL) + json_match = re.search(r"```json\n(.*?)\n```", response, re.DOTALL) if json_match: try: return json.loads(json_match.group(1)) except json.JSONDecodeError: pass - + return {"error": "Failed to analyze web response", "raw_response": response} - - def fix_broken_tool(self, tool_name: str, error_output: str, source_code: str = None) -> str: + + def fix_broken_tool( + self, tool_name: str, error_output: str, source_code: str = None + ) -> str: """Generate fixes for broken security tools""" system_prompt = """You are a DevOps engineer specializing in fixing broken security tools. Analyze errors and provide working solutions.""" - + prompt = f""" This security tool is broken and needs fixing: @@ -311,15 +331,15 @@ def fix_broken_tool(self, tool_name: str, error_output: str, source_code: str = - Missing environment variables - Network/permission issues """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) return response or "# Failed to generate fix" - + def prioritize_vulnerabilities(self, vulnerabilities: List[Dict]) -> List[Dict]: """Use AI to prioritize vulnerabilities by exploitability and impact""" system_prompt = """You are a penetration tester prioritizing vulnerabilities. Rank vulnerabilities by exploitability and business impact.""" - + prompt = f""" Prioritize these vulnerabilities for testing: @@ -346,35 +366,42 @@ def prioritize_vulnerabilities(self, vulnerabilities: List[Dict]) -> List[Dict]: ] }} """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) if response: try: result = json.loads(response) - return result.get('prioritized_vulnerabilities', vulnerabilities) + return result.get("prioritized_vulnerabilities", vulnerabilities) except: pass - + return vulnerabilities def analyze_nuclei_output(self, nuclei_json_output: str) -> Dict[str, Any]: """Analyze nuclei output using the high-speed C parser.""" - + if C_PARSER_AVAILABLE and c_parser: # Use the high-speed C parser when available - raw_summary = c_parser.parse_nuclei_output(nuclei_json_output.encode('utf-8')) - summary_str = raw_summary.decode('utf-8') + raw_summary = c_parser.parse_nuclei_output( + nuclei_json_output.encode("utf-8") + ) + summary_str = raw_summary.decode("utf-8") else: # SECURITY FIX: Fallback to pure Python implementation # This ensures the application works even without the native library try: import json + data = json.loads(nuclei_json_output) - critical_count = sum(1 for item in data if item.get('info', {}).get('severity') == 'critical') + critical_count = sum( + 1 + for item in data + if item.get("info", {}).get("severity") == "critical" + ) summary_str = json.dumps({"critical_findings": critical_count}) except (json.JSONDecodeError, TypeError): summary_str = '{"error": "Failed to parse nuclei output"}' - + try: return json.loads(summary_str) except json.JSONDecodeError: @@ -386,6 +413,7 @@ class AIOrchestrator: Manages a multi-step AI reasoning pipeline for complex security tasks. It chains specialized prompts for planning, tool selection, and execution. """ + def __init__(self, prompt_dir: Path): self.prompt_dir = prompt_dir self.ollama = OllamaClient() @@ -400,7 +428,9 @@ def _load_prompts(self): with open(self.prompt_dir / "Devin AI/system.md", "r") as f: prompts["planner"] = f.read() # Manus-style tool selection prompt - with open(self.prompt_dir / "Manus Agent Tools & Prompt/system.md", "r") as f: + with open( + self.prompt_dir / "Manus Agent Tools & Prompt/system.md", "r" + ) as f: prompts["tool_selector"] = f.read() # Cursor-style code/analysis prompt with open(self.prompt_dir / "Cursor Prompts/prompts.md", "r") as f: @@ -415,25 +445,25 @@ def execute_task(self, task_description: str): Executes a full task pipeline: Plan -> Select Tool -> Execute -> Analyze. """ print("--- AI Task Pipeline Initiated ---") - + # 1. Planning Phase (using Devin's prompt) plan = self._planning_phase(task_description) - self.state['plan'] = plan + self.state["plan"] = plan print(f"Phase 1: Plan Created -> {plan}") # 2. Tool Selection Phase (using Manus' prompt) tool_command = self._tool_selection_phase(task_description, plan) - self.state['tool_command'] = tool_command + self.state["tool_command"] = tool_command print(f"Phase 2: Tool Selected -> {tool_command}") # 3. Execution Phase (simulated) execution_result = self._execution_phase(tool_command) - self.state['execution_result'] = execution_result + self.state["execution_result"] = execution_result print(f"Phase 3: Execution Result -> {execution_result[:100]}...") # 4. Analysis Phase (using Cursor's prompt) analysis = self._analysis_phase(execution_result) - self.state['analysis'] = analysis + self.state["analysis"] = analysis print(f"Phase 4: Analysis Complete -> {analysis}") print("--- AI Task Pipeline Complete ---") @@ -441,18 +471,18 @@ def execute_task(self, task_description: str): def _planning_phase(self, task: str) -> str: """Uses the 'planner' prompt to create a high-level strategy.""" - system_prompt = self.prompts['planner'] + system_prompt = self.prompts["planner"] user_prompt = f"Create a step-by-step plan for the following task: {task}" response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response def _tool_selection_phase(self, task: str, plan: str) -> str: """Uses the 'tool_selector' prompt to choose the right command.""" - system_prompt = self.prompts['tool_selector'] + system_prompt = self.prompts["tool_selector"] user_prompt = f"Given the task '{task}' and the plan '{plan}', what is the exact shell command to execute next? Only output the command." response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response - + def _execution_phase(self, command: str) -> str: """Simulates running the command and returns mock output.""" print(f"Simulating execution of: `{command}`") @@ -462,7 +492,7 @@ def _execution_phase(self, command: str) -> str: def _analysis_phase(self, result: str) -> str: """Uses the 'analyst' prompt to interpret the results.""" - system_prompt = self.prompts['analyst'] + system_prompt = self.prompts["analyst"] user_prompt = f"Analyze the following tool output and provide a summary of key findings and recommendations:\n\n{result}" response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response @@ -471,43 +501,51 @@ def _analysis_phase(self, result: str) -> str: # CLI interface for AI module if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser(description="VulnForge AI Module") - parser.add_argument("--test-connection", action="store_true", help="Test Ollama connection") + parser.add_argument( + "--test-connection", action="store_true", help="Test Ollama connection" + ) parser.add_argument("--pull-model", help="Pull a specific model") - parser.add_argument("--list-models", action="store_true", help="List available models") + parser.add_argument( + "--list-models", action="store_true", help="List available models" + ) parser.add_argument("--analyze-nmap", help="Analyze nmap output file") parser.add_argument( - "--ai-pipeline", action="store_true", help="Enable the advanced multi-prompt AI pipeline." + "--ai-pipeline", + action="store_true", + help="Enable the advanced multi-prompt AI pipeline.", ) parser.add_argument( - "--prompt-dir", help="Directory for the AI pipeline prompts.", default="AI_Propmt/system-prompts-and-models-of-ai-tools" + "--prompt-dir", + help="Directory for the AI pipeline prompts.", + default="AI_Propmt/system-prompts-and-models-of-ai-tools", ) - + args = parser.parse_args() - + orchestrator = AIOrchestrator(Path(args.prompt_dir)) - + if args.test_connection: if orchestrator.ollama.is_available(): print("✓ AI system ready") else: print("✗ AI system not available") - + elif args.pull_model: if orchestrator.ollama.pull_model(args.pull_model): print(f"✓ Model {args.pull_model} pulled successfully") else: print(f"✗ Failed to pull model {args.pull_model}") - + elif args.list_models: models = orchestrator.ollama.list_models() print("Available models:") for model in models: print(f" - {model['name']}") - + elif args.analyze_nmap: - with open(args.analyze_nmap, 'r') as f: + with open(args.analyze_nmap, "r") as f: content = f.read() result = orchestrator.analyzer.analyze_nmap_output(content) print(json.dumps(result, indent=2)) @@ -515,11 +553,13 @@ def _analysis_phase(self, result: str) -> str: # Handle AI Pipeline Mode if args.ai_pipeline: if not args.target: - print("Error: A target is required for AI pipeline mode, e.g., --target 'scan example.com'") - + print( + "Error: A target is required for AI pipeline mode, e.g., --target 'scan example.com'" + ) + prompt_path = Path(args.prompt_dir) if not prompt_path.exists(): print(f"Error: Prompt directory not found at '{prompt_path}'") - + orchestrator = AIOrchestrator(prompt_path) - orchestrator.execute_task(f"Perform a security scan on {args.target}") \ No newline at end of file + orchestrator.execute_task(f"Perform a security scan on {args.target}") diff --git a/ai_orchestrator.py b/ai_orchestrator.py index 38999dc..921d897 100644 --- a/ai_orchestrator.py +++ b/ai_orchestrator.py @@ -4,11 +4,13 @@ from pathlib import Path from ai_integration import OllamaClient + class AIOrchestrator: """ Manages a multi-step AI reasoning pipeline for complex security tasks. It chains specialized prompts for planning, tool selection, and execution. """ + def __init__(self, prompt_dir: Path): self.prompt_dir = prompt_dir self.ollama = OllamaClient() @@ -20,40 +22,46 @@ def _load_prompts(self): try: # SECURITY FIX: Use correct path structure and file names for the nested prompt directories base_prompt_dir = self.prompt_dir / "system-prompts-and-models-of-ai-tools" - + # Check what files actually exist and use appropriate fallbacks devin_file = base_prompt_dir / "Devin AI/Prompt.txt" manus_file = base_prompt_dir / "Manus Agent Tools & Prompt/system.md" cursor_file = base_prompt_dir / "Cursor Prompts/prompts.md" - + # Load Devin AI prompt if devin_file.exists(): with open(devin_file, "r") as f: prompts["planner"] = f.read() else: - prompts["planner"] = "You are a planning AI. Create step-by-step plans for tasks." - + prompts["planner"] = ( + "You are a planning AI. Create step-by-step plans for tasks." + ) + # Load Manus prompt if manus_file.exists(): with open(manus_file, "r") as f: prompts["tool_selector"] = f.read() else: - prompts["tool_selector"] = "You are a tool selection AI. Choose appropriate tools for tasks." - + prompts["tool_selector"] = ( + "You are a tool selection AI. Choose appropriate tools for tasks." + ) + # Load Cursor prompt if cursor_file.exists(): with open(cursor_file, "r") as f: prompts["analyst"] = f.read() else: - prompts["analyst"] = "You are an analysis AI. Analyze results and provide insights." - + prompts["analyst"] = ( + "You are an analysis AI. Analyze results and provide insights." + ) + except Exception as e: print(f"Error: Could not load prompts. {e}") # Provide fallback prompts prompts = { "planner": "You are a planning AI. Create step-by-step plans for tasks.", "tool_selector": "You are a tool selection AI. Choose appropriate tools for tasks.", - "analyst": "You are an analysis AI. Analyze results and provide insights." + "analyst": "You are an analysis AI. Analyze results and provide insights.", } return prompts @@ -62,25 +70,25 @@ def execute_task(self, task_description: str): Executes a full task pipeline: Plan -> Select Tool -> Execute -> Analyze. """ print("--- AI Task Pipeline Initiated ---") - + # 1. Planning Phase (using Devin's prompt) plan = self._planning_phase(task_description) - self.state['plan'] = plan + self.state["plan"] = plan print(f"Phase 1: Plan Created -> {plan}") # 2. Tool Selection Phase (using Manus' prompt) tool_command = self._tool_selection_phase(task_description, plan) - self.state['tool_command'] = tool_command + self.state["tool_command"] = tool_command print(f"Phase 2: Tool Selected -> {tool_command}") # 3. Execution Phase (simulated) execution_result = self._execution_phase(tool_command) - self.state['execution_result'] = execution_result + self.state["execution_result"] = execution_result print(f"Phase 3: Execution Result -> {execution_result[:100]}...") # 4. Analysis Phase (using Cursor's prompt) analysis = self._analysis_phase(execution_result) - self.state['analysis'] = analysis + self.state["analysis"] = analysis print(f"Phase 4: Analysis Complete -> {analysis}") print("--- AI Task Pipeline Complete ---") @@ -88,18 +96,18 @@ def execute_task(self, task_description: str): def _planning_phase(self, task: str) -> str: """Uses the 'planner' prompt to create a high-level strategy.""" - system_prompt = self.prompts['planner'] + system_prompt = self.prompts["planner"] user_prompt = f"Create a step-by-step plan for the following task: {task}" response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response def _tool_selection_phase(self, task: str, plan: str) -> str: """Uses the 'tool_selector' prompt to choose the right command.""" - system_prompt = self.prompts['tool_selector'] + system_prompt = self.prompts["tool_selector"] user_prompt = f"Given the task '{task}' and the plan '{plan}', what is the exact shell command to execute next? Only output the command." response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response - + def _execution_phase(self, command: str) -> str: """Executes the shell command and returns real output.""" print(f"Executing: `{command}`") @@ -107,7 +115,9 @@ def _execution_phase(self, command: str) -> str: # SECURITY FIX: Use shlex.split to safely parse command without shell=True # This prevents command injection while maintaining functionality cmd_parts = shlex.split(command) - result = subprocess.run(cmd_parts, capture_output=True, text=True, timeout=600) + result = subprocess.run( + cmd_parts, capture_output=True, text=True, timeout=600 + ) if result.returncode == 0: return result.stdout else: @@ -117,7 +127,7 @@ def _execution_phase(self, command: str) -> str: def _analysis_phase(self, result: str) -> str: """Uses the 'analyst' prompt to interpret the results.""" - system_prompt = self.prompts['analyst'] + system_prompt = self.prompts["analyst"] user_prompt = f"Analyze the following tool output and provide a summary of key findings and recommendations:\n\n{result}" response = self.ollama.generate(user_prompt, system_prompt=system_prompt) - return response \ No newline at end of file + return response diff --git a/ai_wrapper/__init__.py b/ai_wrapper/__init__.py index f8e7ba5..4e94e50 100755 --- a/ai_wrapper/__init__.py +++ b/ai_wrapper/__init__.py @@ -10,4 +10,4 @@ from .ollama_wrapper import OllamaWrapper -__all__ = ['OllamaWrapper'] \ No newline at end of file +__all__ = ["OllamaWrapper"] diff --git a/ai_wrapper/llm_engine.py b/ai_wrapper/llm_engine.py index 4d6a967..b5ff72b 100755 --- a/ai_wrapper/llm_engine.py +++ b/ai_wrapper/llm_engine.py @@ -18,6 +18,7 @@ import requests from functools import lru_cache + class LLMEngine: def __init__(self, config_path: Optional[Path] = None): self.logger = logging.getLogger(__name__) @@ -26,25 +27,21 @@ def __init__(self, config_path: Optional[Path] = None): self.models = self._initialize_models() self.current_model = self.models[0] # Start with preferred model self.response_cache = {} - + def _load_config(self, config_path: Optional[Path] = None) -> Dict: """Load LLM configuration""" if not config_path: config_path = Path.home() / ".vulnforge" / "configs" / "llm_config.json" - + default_config = { "preferred_model": "deepseek-coder-v2:16b-lite-base-q5_K_S", - "fallback_models": [ - "deepseek-coder:6.7b", - "codellama:7b", - "mistral:7b" - ], + "fallback_models": ["deepseek-coder:6.7b", "codellama:7b", "mistral:7b"], "cache_size": 100, "timeout": 180, "max_retries": 3, - "retry_delay": 2 + "retry_delay": 2, } - + try: if config_path.exists(): with open(config_path) as f: @@ -53,68 +50,79 @@ def _load_config(self, config_path: Optional[Path] = None) -> Dict: except Exception as e: self.logger.error("Error loading LLM config: %s", e) self.config = {} - + def _initialize_models(self) -> List[str]: """Initialize available models""" available_models = [] - + # Check preferred model first if self._is_model_available(self.config["preferred_model"]): available_models.append(self.config["preferred_model"]) - + # Check fallback models for model in self.config["fallback_models"]: if self._is_model_available(model): available_models.append(model) - + if not available_models: self.logger.error("No models available!") - + return available_models - + def _is_model_available(self, model: str) -> bool: """Check if a model is available""" try: response = requests.get(f"{self.base_url}/api/tags", timeout=5) if response.status_code == 200: - models = response.json().get('models', []) - return any(m['name'] == model for m in models) - except (requests.RequestException, requests.Timeout, requests.ConnectionError) as e: + models = response.json().get("models", []) + return any(m["name"] == model for m in models) + except ( + requests.RequestException, + requests.Timeout, + requests.ConnectionError, + ) as e: self.logger.error("Error checking model availability: %s", e) return False return False - + def _pull_model(self, model: str) -> bool: """Pull a model if not available""" try: self.logger.info("Pulling model: %s", model) data = {"name": model} - response = requests.post(f"{self.base_url}/api/pull", json=data, stream=True) - + response = requests.post( + f"{self.base_url}/api/pull", json=data, stream=True + ) + for line in response.iter_lines(): if line: try: - status = json.loads(line.decode('utf-8')) - if status.get('status') == 'success': + status = json.loads(line.decode("utf-8")) + if status.get("status") == "success": return True except: continue except Exception as e: self.logger.error("Error pulling model %s: %s", model, e) return False - + @lru_cache(maxsize=100) - def query(self, prompt: str, system_prompt: Optional[str] = None, - model: Optional[str] = None, use_cache: bool = True) -> Optional[str]: + def query( + self, + prompt: str, + system_prompt: Optional[str] = None, + model: Optional[str] = None, + use_cache: bool = True, + ) -> Optional[str]: """Query the LLM with fallback support""" if not model: model = self.current_model - + # Check cache if enabled cache_key = f"{model}:{prompt}:{system_prompt}" if use_cache and cache_key in self.response_cache: return self.response_cache[cache_key] - + for attempt in range(self.config["max_retries"]): try: data = { @@ -127,28 +135,32 @@ def query(self, prompt: str, system_prompt: Optional[str] = None, "max_tokens": 4096, "num_ctx": 8192, "num_thread": 8, - "repeat_penalty": 1.1 - } + "repeat_penalty": 1.1, + }, } - + if system_prompt: data["system"] = system_prompt - + response = requests.post( - f"{self.base_url}/api/generate", - json=data, - timeout=self.config["timeout"] + f"{self.base_url}/api/generate", + json=data, + timeout=self.config["timeout"], ) - + if response.status_code == 200: - result = response.json().get('response', '').strip() + result = response.json().get("response", "").strip() if use_cache: self.response_cache[cache_key] = result return result - - except (requests.RequestException, requests.Timeout, requests.ConnectionError) as e: + + except ( + requests.RequestException, + requests.Timeout, + requests.ConnectionError, + ) as e: self.logger.error("Error querying model %s: %s", model, e) - + # Try next model if available if model in self.models: current_index = self.models.index(model) @@ -157,23 +169,23 @@ def query(self, prompt: str, system_prompt: Optional[str] = None, self.logger.info("Switching to fallback model: %s", model) else: break - + time.sleep(self.config["retry_delay"]) - + return None - + def clear_cache(self): """Clear the response cache""" self.response_cache.clear() self.query.cache_clear() - + def get_available_models(self) -> List[str]: """Get list of available models""" return self.models.copy() - + def set_preferred_model(self, model: str) -> bool: """Set preferred model if available""" if self._is_model_available(model): self.current_model = model return True - return False \ No newline at end of file + return False diff --git a/conftest.py b/conftest.py index a76b773..e6622b2 100644 --- a/conftest.py +++ b/conftest.py @@ -1,4 +1,3 @@ - import os import sys diff --git a/examples/notifications.py b/examples/notifications.py index e68924c..471d54d 100755 --- a/examples/notifications.py +++ b/examples/notifications.py @@ -10,26 +10,23 @@ # Configure logging logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) + async def main(): # Initialize config manager config = ConfigManager() - + # Initialize notifier notifier = Notifier(config) await notifier.start() - + try: # Example 1: Basic notification - await notifier.notify( - "Scan started for example.com", - "info" - ) - + await notifier.notify("Scan started for example.com", "info") + # Example 2: Vulnerability found await notifier.notify( "SQL Injection vulnerability detected", @@ -38,10 +35,10 @@ async def main(): "vulnerability": "SQL Injection", "affected_url": "https://example.com/login", "payload": "' OR '1'='1", - "confidence": "high" - } + "confidence": "high", + }, ) - + # Example 3: Critical finding await notifier.notify( "Remote Code Execution vulnerability found!", @@ -50,11 +47,11 @@ async def main(): "vulnerability": "RCE", "affected_component": "File Upload Handler", "cve": "CVE-2023-1234", - "exploit_available": True + "exploit_available": True, }, - channels=["email", "discord"] # Send to specific channels + channels=["email", "discord"], # Send to specific channels ) - + # Example 4: Scan completion await notifier.notify( "Scan completed successfully", @@ -62,22 +59,18 @@ async def main(): data={ "target": "example.com", "duration": "2h 15m", - "findings": { - "critical": 1, - "high": 3, - "medium": 5, - "low": 8 - } - } + "findings": {"critical": 1, "high": 3, "medium": 5, "low": 8}, + }, ) - + # Wait for notifications to be processed await asyncio.sleep(1) - + except Exception as e: logger.error(f"Error in notification example: {e}") finally: await notifier.stop() + if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + asyncio.run(main()) diff --git a/modules/ai/ai_assistant.py b/modules/ai/ai_assistant.py index dd0c478..4082646 100644 --- a/modules/ai/ai_assistant.py +++ b/modules/ai/ai_assistant.py @@ -1,5 +1,6 @@ import os + class AIAssistant: def ask_ai(self, query, context=None, debug=False): try: @@ -27,4 +28,4 @@ def _log_ai_interaction(self, query, response, debug): log_path = os.path.expanduser("~/.vulnforge/sessions/logs/ai_controller.log") os.makedirs(os.path.dirname(log_path), exist_ok=True) with open(log_path, "a") as f: - f.write(f"[QUERY] {query}\n[RESPONSE] {response}\n[DEBUG] {debug}\n\n") \ No newline at end of file + f.write(f"[QUERY] {query}\n[RESPONSE] {response}\n[DEBUG] {debug}\n\n") diff --git a/modules/ai/ai_controller.py b/modules/ai/ai_controller.py index 9464e4f..8bcbdc4 100644 --- a/modules/ai/ai_controller.py +++ b/modules/ai/ai_controller.py @@ -1,5 +1,6 @@ import os + class AIController: def run_ai_task(self, task, context=None, debug=False): try: @@ -16,4 +17,4 @@ def _log_ai_task(self, task, result, debug): log_path = os.path.expanduser("~/.vulnforge/sessions/logs/ai_controller.log") os.makedirs(os.path.dirname(log_path), exist_ok=True) with open(log_path, "a") as f: - f.write(f"[TASK] {task}\n[RESULT] {result}\n[DEBUG] {debug}\n\n") \ No newline at end of file + f.write(f"[TASK] {task}\n[RESULT] {result}\n[DEBUG] {debug}\n\n") diff --git a/modules/cve_collector/__init__.py b/modules/cve_collector/__init__.py index faa95ed..18a99b2 100755 --- a/modules/cve_collector/__init__.py +++ b/modules/cve_collector/__init__.py @@ -4,4 +4,4 @@ from .cve_collector import CVECollector -__all__ = ['CVECollector'] \ No newline at end of file +__all__ = ["CVECollector"] diff --git a/modules/cve_collector/cve_collector.py b/modules/cve_collector/cve_collector.py index f73b4db..d186498 100755 --- a/modules/cve_collector/cve_collector.py +++ b/modules/cve_collector/cve_collector.py @@ -11,7 +11,13 @@ from typing import Dict, List, Optional, Any from datetime import datetime, timedelta from rich.console import Console -from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn +from rich.progress import ( + Progress, + SpinnerColumn, + TextColumn, + BarColumn, + TaskProgressColumn, +) import hashlib import time import re @@ -21,6 +27,7 @@ from utils.cli_utils import create_progress, print_results_table from ai_wrapper.ollama_wrapper import OllamaWrapper + class CVECollector: def __init__(self, base_dir: Path, ai_wrapper: Optional[OllamaWrapper] = None): self.base_dir = base_dir @@ -31,39 +38,39 @@ def __init__(self, base_dir: Path, ai_wrapper: Optional[OllamaWrapper] = None): self.cache_dir = base_dir / "cache" / "cve_feeds" self.data_dir.mkdir(parents=True, exist_ok=True) self.cache_dir.mkdir(parents=True, exist_ok=True) - + # API rate limiting self.nvd_rate_limit = 5 # requests per second self.github_rate_limit = 30 # requests per hour self.last_nvd_request = 0 self.last_github_request = 0 self.github_requests_remaining = 30 - + # Load API keys if available self.api_keys = self._load_api_keys() - + def _load_api_keys(self) -> Dict[str, str]: """Load API keys from config""" try: config_path = self.base_dir / "config" / "api_keys.json" if config_path.exists(): - with open(config_path, 'r') as f: + with open(config_path, "r") as f: return json.load(f) except Exception as e: self.logger.warning("Failed to load API keys: %s", e) self.api_keys = {} return {} - + async def _wait_for_rate_limit(self, api_type: str): """Handle API rate limiting""" current_time = time.time() - + if api_type == "nvd": time_since_last = current_time - self.last_nvd_request if time_since_last < (1 / self.nvd_rate_limit): await asyncio.sleep((1 / self.nvd_rate_limit) - time_since_last) self.last_nvd_request = time.time() - + elif api_type == "github": if self.github_requests_remaining <= 0: # Wait until rate limit resets (1 hour) @@ -71,55 +78,56 @@ async def _wait_for_rate_limit(self, api_type: str): self.github_requests_remaining = 30 self.github_requests_remaining -= 1 self.last_github_request = current_time - + def _get_cache_path(self, source: str, identifier: str) -> Path: """Get cache file path for a request""" cache_key = hashlib.md5(f"{source}:{identifier}".encode()).hexdigest() return self.cache_dir / f"{cache_key}.json" - - async def _get_cached_data(self, cache_path: Path, max_age: int = 3600) -> Optional[Dict]: + + async def _get_cached_data( + self, cache_path: Path, max_age: int = 3600 + ) -> Optional[Dict]: """Get cached data if it exists and is not too old""" try: if cache_path.exists(): stat = cache_path.stat() if time.time() - stat.st_mtime < max_age: - async with aiofiles.open(cache_path, 'r') as f: + async with aiofiles.open(cache_path, "r") as f: return json.loads(await f.read()) except Exception as e: self.logger.warning("Error reading cache: %s", e) return None - + async def _save_to_cache(self, cache_path: Path, data: Dict): """Save data to cache""" try: - async with aiofiles.open(cache_path, 'w') as f: + async with aiofiles.open(cache_path, "w") as f: await f.write(json.dumps(data)) except Exception as e: self.logger.warning("Error saving to cache: %s", e) - + async def fetch_nvd_feed(self, start_date: Optional[str] = None) -> List[Dict]: """Fetch CVE data from NVD feed""" if not start_date: start_date = (datetime.now() - timedelta(days=7)).strftime("%Y-%m-%d") - + self.console.print("[bold blue]Fetching NVD CVE feed...[/bold blue]") - + # Check cache first cache_path = self._get_cache_path("nvd", start_date) - cached_data = await self._get_cached_data(cache_path, max_age=3600) # 1 hour cache + cached_data = await self._get_cached_data( + cache_path, max_age=3600 + ) # 1 hour cache if cached_data: return cached_data - + url = f"https://services.nvd.nist.gov/rest/json/cves/2.0" - params = { - "pubStartDate": f"{start_date}T00:00:00.000", - "resultsPerPage": 2000 - } - + params = {"pubStartDate": f"{start_date}T00:00:00.000", "resultsPerPage": 2000} + # Add API key if available if "nvd" in self.api_keys: params["apiKey"] = self.api_keys["nvd"] - + async with aiohttp.ClientSession() as session: try: await self._wait_for_rate_limit("nvd") @@ -139,47 +147,63 @@ async def fetch_nvd_feed(self, start_date: Optional[str] = None) -> List[Dict]: except (aiohttp.ClientError, asyncio.TimeoutError) as e: self.logger.error("Error fetching NVD feed: %s", e) return [] - + async def fetch_exploit_db(self) -> List[Dict]: """Fetch exploit data from Exploit-DB""" self.console.print("[bold blue]Fetching Exploit-DB data...[/bold blue]") - + # Check cache first cache_path = self._get_cache_path("exploitdb", "latest") - cached_data = await self._get_cached_data(cache_path, max_age=86400) # 24 hour cache + cached_data = await self._get_cached_data( + cache_path, max_age=86400 + ) # 24 hour cache if cached_data: return cached_data - + url = "https://raw.githubusercontent.com/offensive-security/exploitdb/master/files_exploits.csv" - + async with aiohttp.ClientSession() as session: try: async with session.get(url) as response: if response.status == 200: content = await response.text() exploits = [] - + # Parse CSV content for line in content.splitlines()[1:]: # Skip header try: - parts = line.split(',') + parts = line.split(",") if len(parts) >= 4: # Extract CVE IDs from description - cve_ids = re.findall(r'CVE-\d{4}-\d+', parts[2]) - - exploits.append({ - "id": parts[0], - "file": parts[1], - "description": parts[2], - "date": parts[3], - "author": parts[4] if len(parts) > 4 else "Unknown", - "platform": parts[5] if len(parts) > 5 else "Unknown", - "type": parts[6] if len(parts) > 6 else "Unknown", - "cve_ids": cve_ids - }) + cve_ids = re.findall(r"CVE-\d{4}-\d+", parts[2]) + + exploits.append( + { + "id": parts[0], + "file": parts[1], + "description": parts[2], + "date": parts[3], + "author": ( + parts[4] + if len(parts) > 4 + else "Unknown" + ), + "platform": ( + parts[5] + if len(parts) > 5 + else "Unknown" + ), + "type": ( + parts[6] + if len(parts) > 6 + else "Unknown" + ), + "cve_ids": cve_ids, + } + ) except: continue - + await self._save_to_cache(cache_path, exploits) return exploits else: @@ -188,62 +212,82 @@ async def fetch_exploit_db(self) -> List[Dict]: except Exception as e: self.logger.error("Error fetching Exploit-DB: %s", e) return [] - + async def fetch_github_pocs(self, cve_id: str) -> List[Dict]: """Search for PoCs on GitHub""" - self.console.print(f"[bold blue]Searching GitHub for {cve_id} PoCs...[/bold blue]") - + self.console.print( + f"[bold blue]Searching GitHub for {cve_id} PoCs...[/bold blue]" + ) + # Check cache first cache_path = self._get_cache_path("github", cve_id) - cached_data = await self._get_cached_data(cache_path, max_age=86400) # 24 hour cache + cached_data = await self._get_cached_data( + cache_path, max_age=86400 + ) # 24 hour cache if cached_data: return cached_data - + # GitHub API search query query = f"{cve_id} poc OR proof of concept OR exploit" url = "https://api.github.com/search/code" headers = { "Accept": "application/vnd.github.v3+json", - "User-Agent": "VulnForge/1.0" + "User-Agent": "VulnForge/1.0", } - + # Add API key if available if "github" in self.api_keys: headers["Authorization"] = f"token {self.api_keys['github']}" - + async with aiohttp.ClientSession() as session: try: await self._wait_for_rate_limit("github") - async with session.get(url, params={"q": query}, headers=headers) as response: + async with session.get( + url, params={"q": query}, headers=headers + ) as response: if response.status == 200: data = await response.json() results = data.get("items", []) - + # Process results processed_results = [] for result in results: try: # Get file content content_url = result["url"] - async with session.get(content_url, headers=headers) as content_response: + async with session.get( + content_url, headers=headers + ) as content_response: if content_response.status == 200: content_data = await content_response.json() content = content_data.get("content", "") - + # Extract relevant information - processed_results.append({ - "repository": result["repository"]["full_name"], - "file_path": result["path"], - "url": result["html_url"], - "content": content, - "language": result.get("language", "Unknown"), - "stars": result["repository"].get("stargazers_count", 0), - "forks": result["repository"].get("forks_count", 0) - }) + processed_results.append( + { + "repository": result["repository"][ + "full_name" + ], + "file_path": result["path"], + "url": result["html_url"], + "content": content, + "language": result.get( + "language", "Unknown" + ), + "stars": result["repository"].get( + "stargazers_count", 0 + ), + "forks": result["repository"].get( + "forks_count", 0 + ), + } + ) except Exception as e: - self.logger.warning("Error processing GitHub result: %s", e) + self.logger.warning( + "Error processing GitHub result: %s", e + ) continue - + await self._save_to_cache(cache_path, processed_results) return processed_results elif response.status == 403: # Rate limit exceeded @@ -256,12 +300,12 @@ async def fetch_github_pocs(self, cve_id: str) -> List[Dict]: except (aiohttp.ClientError, asyncio.TimeoutError) as e: self.logger.error("Error searching GitHub: %s", e) return [] - + async def analyze_cve(self, cve_data: Dict) -> Dict: """Analyze CVE data using AI""" if not self.ai_wrapper: return cve_data - + prompt = f""" Analyze this CVE and provide: 1. Severity assessment (critical, high, medium, low) @@ -305,7 +349,7 @@ async def analyze_cve(self, cve_data: Dict) -> Dict: }} }} """ - + try: analysis = await self.ai_wrapper.generate(prompt) try: @@ -317,30 +361,30 @@ async def analyze_cve(self, cve_data: Dict) -> Dict: cve_data["ai_analysis"] = {"raw_analysis": analysis} except Exception as e: self.logger.error("Error analyzing CVE: %s", e) - + return cve_data - + async def collect_cves(self, target_info: Dict) -> Dict[str, Any]: """Collect and analyze CVEs for target""" self.console.print("[bold blue]Starting CVE collection...[/bold blue]") - + with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn(), - console=self.console + console=self.console, ) as progress: # Fetch NVD feed task = progress.add_task("Fetching NVD feed...", total=None) cves = await self.fetch_nvd_feed() progress.update(task, completed=True) - + # Fetch Exploit-DB data task = progress.add_task("Fetching Exploit-DB data...", total=None) exploits = await self.fetch_exploit_db() progress.update(task, completed=True) - + # Match CVEs to target task = progress.add_task("Matching CVEs to target...", total=None) matched_cves = [] @@ -348,7 +392,7 @@ async def collect_cves(self, target_info: Dict) -> Dict[str, Any]: if self._matches_target(cve, target_info): matched_cves.append(cve) progress.update(task, completed=True) - + # Fetch GitHub PoCs task = progress.add_task("Fetching GitHub PoCs...", total=None) for cve in matched_cves: @@ -357,26 +401,26 @@ async def collect_cves(self, target_info: Dict) -> Dict[str, Any]: pocs = await self.fetch_github_pocs(cve_id) cve["github_pocs"] = pocs progress.update(task, completed=True) - + # Analyze CVEs task = progress.add_task("Analyzing CVEs...", total=None) for cve in matched_cves: cve = await self.analyze_cve(cve) progress.update(task, completed=True) - + # Save results task = progress.add_task("Saving results...", total=None) results = { "target_info": target_info, "cves": matched_cves, "exploits": exploits, - "generated_at": datetime.now().isoformat() + "generated_at": datetime.now().isoformat(), } await self._save_results(results) progress.update(task, completed=True) - + return results - + def _matches_target(self, cve: Dict, target_info: Dict) -> bool: """Check if CVE matches target""" # Get CPE strings @@ -384,41 +428,51 @@ def _matches_target(self, cve: Dict, target_info: Dict) -> bool: for node in cve.get("cve", {}).get("configurations", []): for cpe_match in node.get("cpeMatch", []): cpe_list.append(cpe_match.get("criteria")) - + # Check each CPE for cpe in cpe_list: if self._check_cpe_match(cpe, target_info): return True - + # Check description for target keywords - description = cve.get("cve", {}).get("descriptions", [{}])[0].get("value", "").lower() + description = ( + cve.get("cve", {}).get("descriptions", [{}])[0].get("value", "").lower() + ) target_keywords = [ target_info.get("name", "").lower(), target_info.get("vendor", "").lower(), target_info.get("product", "").lower(), - target_info.get("version", "").lower() + target_info.get("version", "").lower(), ] - + return any(keyword in description for keyword in target_keywords if keyword) - + def _check_cpe_match(self, cpe: str, target_info: Dict) -> bool: """Check if CPE string matches target""" if not cpe: return False - + # Parse CPE parts = cpe.split(":") if len(parts) < 5: return False - + # Check vendor - if target_info.get("vendor") and parts[3] != "*" and parts[3].lower() != target_info["vendor"].lower(): + if ( + target_info.get("vendor") + and parts[3] != "*" + and parts[3].lower() != target_info["vendor"].lower() + ): return False - + # Check product - if target_info.get("product") and parts[4] != "*" and parts[4].lower() != target_info["product"].lower(): + if ( + target_info.get("product") + and parts[4] != "*" + and parts[4].lower() != target_info["product"].lower() + ): return False - + # Check version if target_info.get("version") and parts[5] != "*": # Handle version ranges @@ -428,65 +482,69 @@ def _check_cpe_match(self, cpe: str, target_info: Dict) -> bool: return False elif parts[5] != target_info["version"]: return False - + return True - + async def _save_results(self, results: Dict): """Save results to file""" # Create results directory timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") results_dir = self.data_dir / timestamp results_dir.mkdir(exist_ok=True) - + # Save JSON json_path = results_dir / "results.json" - async with aiofiles.open(json_path, 'w') as f: + async with aiofiles.open(json_path, "w") as f: await f.write(json.dumps(results, indent=2)) - + # Save Markdown report md_path = results_dir / "report.md" - async with aiofiles.open(md_path, 'w') as f: + async with aiofiles.open(md_path, "w") as f: await f.write(self._generate_markdown_report(results)) - + def _generate_markdown_report(self, results: Dict) -> str: """Generate Markdown report""" report = [] - + # Add header report.append("# CVE Analysis Report") report.append(f"Generated: {results['generated_at']}") report.append("") - + # Add target info report.append("## Target Information") for key, value in results["target_info"].items(): report.append(f"- **{key}**: {value}") report.append("") - + # Add findings summary report.append("## Findings Summary") severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0} for cve in results["cves"]: severity = cve.get("ai_analysis", {}).get("severity", "info").lower() severity_counts[severity] += 1 - + report.append("### Severity Distribution") for severity, count in severity_counts.items(): report.append(f"- **{severity.title()}**: {count}") report.append("") - + # Add detailed findings report.append("## Detailed Findings") for cve in results["cves"]: cve_id = cve.get("cve", {}).get("id", "Unknown") severity = cve.get("ai_analysis", {}).get("severity", "info").lower() - + report.append(f"### {cve_id} ({severity.title()})") - + # Add description - description = cve.get("cve", {}).get("descriptions", [{}])[0].get("value", "No description available") + description = ( + cve.get("cve", {}) + .get("descriptions", [{}])[0] + .get("value", "No description available") + ) report.append(f"**Description**: {description}") - + # Add analysis analysis = cve.get("ai_analysis", {}) if analysis: @@ -502,14 +560,16 @@ def _generate_markdown_report(self, results: Dict) -> str: report.append(f" - {k}: {v}") else: report.append(f"- **{key}**: {value}") - + # Add PoCs pocs = cve.get("github_pocs", []) if pocs: report.append("\n**Proof of Concepts**:") for poc in pocs: - report.append(f"- [{poc['repository']}/{poc['file_path']}]({poc['url']})") - + report.append( + f"- [{poc['repository']}/{poc['file_path']}]({poc['url']})" + ) + report.append("") - - return "\n".join(report) \ No newline at end of file + + return "\n".join(report) diff --git a/modules/darkweb/__init__.py b/modules/darkweb/__init__.py index 59141de..c8410d7 100644 --- a/modules/darkweb/__init__.py +++ b/modules/darkweb/__init__.py @@ -2,7 +2,10 @@ Dark web OSINT integrations for VulnForge. """ -from .robin.runner import run_darkweb_osint, ROBIN_DEFAULT_MODEL, get_robin_model_choices +from .robin.runner import ( + run_darkweb_osint, + ROBIN_DEFAULT_MODEL, + get_robin_model_choices, +) __all__ = ["run_darkweb_osint", "ROBIN_DEFAULT_MODEL", "get_robin_model_choices"] - diff --git a/modules/darkweb/robin/llm.py b/modules/darkweb/robin/llm.py index d3e882e..a6a6146 100644 --- a/modules/darkweb/robin/llm.py +++ b/modules/darkweb/robin/llm.py @@ -101,9 +101,7 @@ def filter_results(llm, query, results): # Remove duplicates while preserving order seen = set() - parsed_indices = [ - i for i in parsed_indices if not (i in seen or seen.add(i)) - ] + parsed_indices = [i for i in parsed_indices if not (i in seen or seen.add(i))] if not parsed_indices: logging.warning( diff --git a/modules/darkweb/robin/llm_utils.py b/modules/darkweb/robin/llm_utils.py index 46ce05e..454c740 100644 --- a/modules/darkweb/robin/llm_utils.py +++ b/modules/darkweb/robin/llm_utils.py @@ -12,7 +12,11 @@ class BufferedStreamingHandler(BaseCallbackHandler): - def __init__(self, buffer_limit: int = 60, ui_callback: Optional[Callable[[str], None]] = None): + def __init__( + self, + buffer_limit: int = 60, + ui_callback: Optional[Callable[[str], None]] = None, + ): self.buffer = "" self.buffer_limit = buffer_limit self.ui_callback = ui_callback @@ -47,59 +51,55 @@ def on_llm_end(self, response, **kwargs) -> None: # Map input model choices (lowercased) to their configuration # Each config includes the class and any model-specific constructor parameters _llm_config_map = { - 'gpt-4.1': { - 'class': ChatOpenAI, - 'constructor_params': {'model_name': 'gpt-4.1'} + "gpt-4.1": {"class": ChatOpenAI, "constructor_params": {"model_name": "gpt-4.1"}}, + "gpt-5.1": {"class": ChatOpenAI, "constructor_params": {"model_name": "gpt-5.1"}}, + "gpt-5-mini": { + "class": ChatOpenAI, + "constructor_params": {"model_name": "gpt-5-mini"}, }, - 'gpt-5.1': { - 'class': ChatOpenAI, - 'constructor_params': {'model_name': 'gpt-5.1'} + "gpt-5-nano": { + "class": ChatOpenAI, + "constructor_params": {"model_name": "gpt-5-nano"}, }, - 'gpt-5-mini': { - 'class': ChatOpenAI, - 'constructor_params': {'model_name': 'gpt-5-mini'} + "claude-sonnet-4-5": { + "class": ChatAnthropic, + "constructor_params": {"model": "claude-sonnet-4-5"}, }, - 'gpt-5-nano': { - 'class': ChatOpenAI, - 'constructor_params': {'model_name': 'gpt-5-nano'} + "claude-sonnet-4-0": { + "class": ChatAnthropic, + "constructor_params": {"model": "claude-sonnet-4-0"}, }, - 'claude-sonnet-4-5': { - 'class': ChatAnthropic, - 'constructor_params': {'model': 'claude-sonnet-4-5'} + "gemini-2.5-flash": { + "class": ChatGoogleGenerativeAI, + "constructor_params": {"model": "gemini-2.5-flash"}, }, - 'claude-sonnet-4-0': { - 'class': ChatAnthropic, - 'constructor_params': {'model': 'claude-sonnet-4-0'} + "gemini-2.5-flash-lite": { + "class": ChatGoogleGenerativeAI, + "constructor_params": {"model": "gemini-2.5-flash-lite"}, }, - 'gemini-2.5-flash': { - 'class': ChatGoogleGenerativeAI, - 'constructor_params': {'model': 'gemini-2.5-flash'} + "gemini-2.5-pro": { + "class": ChatGoogleGenerativeAI, + "constructor_params": {"model": "gemini-2.5-pro"}, }, - 'gemini-2.5-flash-lite': { - 'class': ChatGoogleGenerativeAI, - 'constructor_params': {'model': 'gemini-2.5-flash-lite'} + "llama3.2": { + "class": ChatOllama, + "constructor_params": {"model": "llama3.2:latest", "base_url": OLLAMA_BASE_URL}, }, - 'gemini-2.5-pro': { - 'class': ChatGoogleGenerativeAI, - 'constructor_params': {'model': 'gemini-2.5-pro'} + "llama3.1": { + "class": ChatOllama, + "constructor_params": {"model": "llama3.1:latest", "base_url": OLLAMA_BASE_URL}, }, - 'llama3.2': { - 'class': ChatOllama, - 'constructor_params': {'model': 'llama3.2:latest', 'base_url': OLLAMA_BASE_URL} + "gemma3": { + "class": ChatOllama, + "constructor_params": {"model": "gemma3:latest", "base_url": OLLAMA_BASE_URL}, }, - 'llama3.1': { - 'class': ChatOllama, - 'constructor_params': {'model': 'llama3.1:latest', 'base_url': OLLAMA_BASE_URL} + "deepseek-r1": { + "class": ChatOllama, + "constructor_params": { + "model": "deepseek-r1:latest", + "base_url": OLLAMA_BASE_URL, + }, }, - 'gemma3': { - 'class': ChatOllama, - 'constructor_params': {'model': 'gemma3:latest', 'base_url': OLLAMA_BASE_URL} - }, - 'deepseek-r1': { - 'class': ChatOllama, - 'constructor_params': {'model': 'deepseek-r1:latest', 'base_url': OLLAMA_BASE_URL} - } - # Add more models here easily: # 'mistral7b': { # 'class': ChatOllama, @@ -180,7 +180,10 @@ def resolve_model_config(model_choice: str): if _normalize_model_name(ollama_model) == model_choice_lower: return { "class": ChatOllama, - "constructor_params": {"model": ollama_model, "base_url": OLLAMA_BASE_URL}, + "constructor_params": { + "model": ollama_model, + "base_url": OLLAMA_BASE_URL, + }, } - return None \ No newline at end of file + return None diff --git a/modules/darkweb/robin/runner.py b/modules/darkweb/robin/runner.py index 8951f5b..f118cb2 100644 --- a/modules/darkweb/robin/runner.py +++ b/modules/darkweb/robin/runner.py @@ -91,10 +91,14 @@ def run_darkweb_osint( refined_query = refine_query(llm, query) with console.status("Querying dark web indices via Tor...", spinner="dots"): - search_results = get_search_results(refined_query.replace(" ", "+"), max_workers=threads) + search_results = get_search_results( + refined_query.replace(" ", "+"), max_workers=threads + ) if not search_results: - console.print("[yellow]No search results returned. Verify Tor connectivity and try again.[/yellow]") + console.print( + "[yellow]No search results returned. Verify Tor connectivity and try again.[/yellow]" + ) return { "refined_query": refined_query, "search_results": [], @@ -127,4 +131,3 @@ def run_darkweb_osint( def get_robin_model_choices(): """Expose Robin model choices for CLI integration.""" return get_model_choices() - diff --git a/modules/darkweb/robin/scrape.py b/modules/darkweb/robin/scrape.py index 0ec6bcc..2b57b45 100644 --- a/modules/darkweb/robin/scrape.py +++ b/modules/darkweb/robin/scrape.py @@ -5,6 +5,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed import warnings + warnings.filterwarnings("ignore") # Define a list of rotating user agents. @@ -17,63 +18,66 @@ "Mozilla/5.0 (X11; Linux i686; rv:137.0) Gecko/20100101 Firefox/137.0", "Mozilla/5.0 (Macintosh; Intel Mac OS X 14_7_5) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/18.3 Safari/605.1.15", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36 Edg/135.0.3179.54", - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36 Edg/135.0.3179.54" + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36 Edg/135.0.3179.54", ] # Global counter and lock for thread-safe Tor rotation request_counter = 0 counter_lock = threading.Lock() -def scrape_single(url_data, rotate=False, rotate_interval=5, control_port=9051, control_password=None): + +def scrape_single( + url_data, rotate=False, rotate_interval=5, control_port=9051, control_password=None +): """ Scrapes a single URL. If the URL is an onion site, routes the request through Tor. Returns a tuple (url, scraped_text). """ - url = url_data['link'] + url = url_data["link"] use_tor = ".onion" in url proxies = None if use_tor: proxies = { "http": "socks5h://127.0.0.1:9050", - "https": "socks5h://127.0.0.1:9050" + "https": "socks5h://127.0.0.1:9050", } - headers = { - "User-Agent": random.choice(USER_AGENTS) - } + headers = {"User-Agent": random.choice(USER_AGENTS)} try: response = requests.get(url, headers=headers, proxies=proxies, timeout=30) if response.status_code == 200: soup = BeautifulSoup(response.text, "html.parser") - scraped_text = url_data['title'] + soup.get_text().replace('\n', ' ').replace('\r', '') + scraped_text = url_data["title"] + soup.get_text().replace( + "\n", " " + ).replace("\r", "") else: - scraped_text = url_data['title'] + scraped_text = url_data["title"] except: - scraped_text = url_data['title'] - + scraped_text = url_data["title"] + return url, scraped_text + def scrape_multiple(urls_data, max_workers=5): """ Scrapes multiple URLs concurrently using a thread pool. - + Parameters: - urls_data: list of URLs to scrape. - max_workers: number of concurrent threads for scraping. - + Returns: A dictionary mapping each URL to its scraped content. """ results = {} - max_chars = 1200 # Taking first n chars from the scraped data + max_chars = 1200 # Taking first n chars from the scraped data with ThreadPoolExecutor(max_workers=max_workers) as executor: future_to_url = { - executor.submit(scrape_single, url_data): url_data - for url_data in urls_data + executor.submit(scrape_single, url_data): url_data for url_data in urls_data } for future in as_completed(future_to_url): url, content = future.result() if len(content) > max_chars: content = content[:max_chars] results[url] = content - return results \ No newline at end of file + return results diff --git a/modules/darkweb/robin/search.py b/modules/darkweb/robin/search.py index 113a359..617af1f 100644 --- a/modules/darkweb/robin/search.py +++ b/modules/darkweb/robin/search.py @@ -4,6 +4,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed import warnings + warnings.filterwarnings("ignore") USER_AGENTS = [ @@ -15,38 +16,35 @@ "Mozilla/5.0 (X11; Linux i686; rv:137.0) Gecko/20100101 Firefox/137.0", "Mozilla/5.0 (Macintosh; Intel Mac OS X 14_7_5) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/18.3 Safari/605.1.15", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36 Edg/135.0.3179.54", - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36 Edg/135.0.3179.54" + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36 Edg/135.0.3179.54", ] SEARCH_ENGINE_ENDPOINTS = [ - "http://juhanurmihxlp77nkq76byazcldy2hlmovfu2epvl5ankdibsot4csyd.onion/search/?q={query}", # Ahmia - "http://3bbad7fauom4d6sgppalyqddsqbf5u5p56b5k5uk2zxsy3d6ey2jobad.onion/search?q={query}", # OnionLand - "http://darkhuntyla64h75a3re5e2l3367lqn7ltmdzpgmr6b4nbz3q2iaxrid.onion/search?q={query}", # DarkRunt - "http://iy3544gmoeclh5de6gez2256v6pjh4omhpqdh2wpeeppjtvqmjhkfwad.onion/torgle/?query={query}", # Torgle - "http://amnesia7u5odx5xbwtpnqk3edybgud5bmiagu75bnqx2crntw5kry7ad.onion/search?query={query}", # Amnesia - "http://kaizerwfvp5gxu6cppibp7jhcqptavq3iqef66wbxenh6a2fklibdvid.onion/search?q={query}", # Kaizer - "http://anima4ffe27xmakwnseih3ic2y7y3l6e7fucwk4oerdn4odf7k74tbid.onion/search?q={query}", # Anima - "http://tornadoxn3viscgz647shlysdy7ea5zqzwda7hierekeuokh5eh5b3qd.onion/search?q={query}", # Tornado - "http://tornetupfu7gcgidt33ftnungxzyfq2pygui5qdoyss34xbgx2qruzid.onion/search?q={query}", # TorNet - "http://torlbmqwtudkorme6prgfpmsnile7ug2zm4u3ejpcncxuhpu4k2j4kyd.onion/index.php?a=search&q={query}", # Torland - "http://findtorroveq5wdnipkaojfpqulxnkhblymc7aramjzajcvpptd4rjqd.onion/search?q={query}", # Find Tor - "http://2fd6cemt4gmccflhm6imvdfvli3nf7zn6rfrwpsy7uhxrgbypvwf5fad.onion/search?query={query}", # Excavator - "http://oniwayzz74cv2puhsgx4dpjwieww4wdphsydqvf5q7eyz4myjvyw26ad.onion/search.php?s={query}", # Onionway - "http://tor66sewebgixwhcqfnp5inzp5x5uohhdy3kvtnyfxc2e5mxiuh34iid.onion/search?q={query}", # Tor66 - "http://3fzh7yuupdfyjhwt3ugzqqof6ulbcl27ecev33knxe3u7goi3vfn2qqd.onion/oss/index.php?search={query}", # OSS (Onion Search Server) + "http://juhanurmihxlp77nkq76byazcldy2hlmovfu2epvl5ankdibsot4csyd.onion/search/?q={query}", # Ahmia + "http://3bbad7fauom4d6sgppalyqddsqbf5u5p56b5k5uk2zxsy3d6ey2jobad.onion/search?q={query}", # OnionLand + "http://darkhuntyla64h75a3re5e2l3367lqn7ltmdzpgmr6b4nbz3q2iaxrid.onion/search?q={query}", # DarkRunt + "http://iy3544gmoeclh5de6gez2256v6pjh4omhpqdh2wpeeppjtvqmjhkfwad.onion/torgle/?query={query}", # Torgle + "http://amnesia7u5odx5xbwtpnqk3edybgud5bmiagu75bnqx2crntw5kry7ad.onion/search?query={query}", # Amnesia + "http://kaizerwfvp5gxu6cppibp7jhcqptavq3iqef66wbxenh6a2fklibdvid.onion/search?q={query}", # Kaizer + "http://anima4ffe27xmakwnseih3ic2y7y3l6e7fucwk4oerdn4odf7k74tbid.onion/search?q={query}", # Anima + "http://tornadoxn3viscgz647shlysdy7ea5zqzwda7hierekeuokh5eh5b3qd.onion/search?q={query}", # Tornado + "http://tornetupfu7gcgidt33ftnungxzyfq2pygui5qdoyss34xbgx2qruzid.onion/search?q={query}", # TorNet + "http://torlbmqwtudkorme6prgfpmsnile7ug2zm4u3ejpcncxuhpu4k2j4kyd.onion/index.php?a=search&q={query}", # Torland + "http://findtorroveq5wdnipkaojfpqulxnkhblymc7aramjzajcvpptd4rjqd.onion/search?q={query}", # Find Tor + "http://2fd6cemt4gmccflhm6imvdfvli3nf7zn6rfrwpsy7uhxrgbypvwf5fad.onion/search?query={query}", # Excavator + "http://oniwayzz74cv2puhsgx4dpjwieww4wdphsydqvf5q7eyz4myjvyw26ad.onion/search.php?s={query}", # Onionway + "http://tor66sewebgixwhcqfnp5inzp5x5uohhdy3kvtnyfxc2e5mxiuh34iid.onion/search?q={query}", # Tor66 + "http://3fzh7yuupdfyjhwt3ugzqqof6ulbcl27ecev33knxe3u7goi3vfn2qqd.onion/oss/index.php?search={query}", # OSS (Onion Search Server) ] + def get_tor_proxies(): - return { - "http": "socks5h://127.0.0.1:9050", - "https": "socks5h://127.0.0.1:9050" - } + return {"http": "socks5h://127.0.0.1:9050", "https": "socks5h://127.0.0.1:9050"} + def fetch_search_results(endpoint, query): url = endpoint.format(query=query) - headers = { - "User-Agent": random.choice(USER_AGENTS) - } + headers = {"User-Agent": random.choice(USER_AGENTS)} proxies = get_tor_proxies() try: response = requests.get(url, headers=headers, proxies=proxies, timeout=30) @@ -54,11 +52,11 @@ def fetch_search_results(endpoint, query): # Normally you would parse html_content with BeautifulSoup and extract results. soup = BeautifulSoup(response.text, "html.parser") links = [] - for a in soup.find_all('a'): + for a in soup.find_all("a"): try: - href = a['href'] + href = a["href"] title = a.get_text(strip=True) - link = re.findall(r'https?:\/\/[^\/]*\.onion.*', href) + link = re.findall(r"https?:\/\/[^\/]*\.onion.*", href) if len(link) != 0: links.append({"title": title, "link": link[0]}) except: @@ -69,11 +67,14 @@ def fetch_search_results(endpoint, query): except: return [] + def get_search_results(refined_query, max_workers=5): results = [] with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = [executor.submit(fetch_search_results, endpoint, refined_query) - for endpoint in SEARCH_ENGINE_ENDPOINTS] + futures = [ + executor.submit(fetch_search_results, endpoint, refined_query) + for endpoint in SEARCH_ENGINE_ENDPOINTS + ] for future in as_completed(futures): result_urls = future.result() results.extend(result_urls) @@ -86,4 +87,4 @@ def get_search_results(refined_query, max_workers=5): if link not in seen_links: seen_links.add(link) unique_results.append(res) - return unique_results \ No newline at end of file + return unique_results diff --git a/modules/darkweb/robin/ui.py b/modules/darkweb/robin/ui.py index b4e8fba..7fa2cf9 100644 --- a/modules/darkweb/robin/ui.py +++ b/modules/darkweb/robin/ui.py @@ -76,15 +76,29 @@ def cached_scrape_multiple(filtered: list, threads: int): index=default_model_index, key="model_select", ) -if any(name not in {"gpt4o", "gpt-4.1", "claude-3-5-sonnet-latest", "llama3.1", "gemini-2.5-flash"} for name in model_options): - st.sidebar.caption("Locally detected Ollama models are automatically added to this list.") +if any( + name + not in { + "gpt4o", + "gpt-4.1", + "claude-3-5-sonnet-latest", + "llama3.1", + "gemini-2.5-flash", + } + for name in model_options +): + st.sidebar.caption( + "Locally detected Ollama models are automatically added to this list." + ) threads = st.sidebar.slider("Scraping Threads", 1, 16, 4, key="thread_slider") # Main UI - logo and input _, logo_col, _ = st.columns(3) with logo_col: - logo_path = Path(__file__).resolve().parent / ".github" / "assets" / "robin_logo.png" + logo_path = ( + Path(__file__).resolve().parent / ".github" / "assets" / "robin_logo.png" + ) if logo_path.exists(): st.image(str(logo_path), width=200) diff --git a/modules/exploit_generator/__init__.py b/modules/exploit_generator/__init__.py index a3b60c7..718242b 100755 --- a/modules/exploit_generator/__init__.py +++ b/modules/exploit_generator/__init__.py @@ -4,4 +4,4 @@ from .exploit_generator import ExploitGenerator -__all__ = ['ExploitGenerator'] \ No newline at end of file +__all__ = ["ExploitGenerator"] diff --git a/modules/exploit_generator/exploit_generator.py b/modules/exploit_generator/exploit_generator.py index 6e84dcd..95db358 100755 --- a/modules/exploit_generator/exploit_generator.py +++ b/modules/exploit_generator/exploit_generator.py @@ -14,6 +14,7 @@ from rich.console import Console from rich.progress import Progress, SpinnerColumn, TextColumn + class ExploitGenerator: def __init__(self, base_dir: Path, llm_engine: Any): self.base_dir = base_dir @@ -24,33 +25,33 @@ def __init__(self, base_dir: Path, llm_engine: Any): self.template_dir = self.base_dir / "templates" / "exploits" self.exploit_dir.mkdir(parents=True, exist_ok=True) self.template_dir.mkdir(parents=True, exist_ok=True) - + # Load exploit templates self.templates = self._load_templates() - + def _load_templates(self) -> Dict[str, str]: """Load exploit templates""" templates = {} try: for template_file in self.template_dir.glob("*.py"): - with open(template_file, 'r') as f: + with open(template_file, "r") as f: templates[template_file.stem] = f.read() except Exception as e: self.logger.error(f"Error loading templates: {e}") return templates - + def generate_exploit(self, cve_data: Dict, recon_data: Dict) -> Dict[str, Any]: """Generate exploit based on CVE and recon data""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: # Prepare prompt with context task = progress.add_task("Preparing exploit generation...", total=None) prompt = self._build_exploit_prompt(cve_data, recon_data) progress.update(task, completed=True) - + # Get AI response task = progress.add_task("Generating exploit code...", total=None) try: @@ -67,41 +68,43 @@ def generate_exploit(self, cve_data: Dict, recon_data: Dict) -> Dict[str, Any]: 5. Add timeout and connection limits 6. Validate all user input 7. Use safe defaults - 8. Include proper logging""" + 8. Include proper logging""", ) progress.update(task, completed=True) except Exception as e: self.logger.error(f"Error generating exploit: {e}") return {"error": f"Failed to generate exploit: {str(e)}"} - + if not response: return {"error": "Failed to generate exploit - empty response"} - + # Parse and validate the response task = progress.add_task("Validating exploit code...", total=None) exploit_data = self._parse_exploit_response(response) if not exploit_data: return {"error": "Failed to parse exploit response"} - + # Validate exploit safety validation = self.validate_exploit(exploit_data) if validation["issues"]: - self.logger.warning(f"Exploit validation issues: {validation['issues']}") + self.logger.warning( + f"Exploit validation issues: {validation['issues']}" + ) exploit_data["validation"] = validation - + # Save the exploit task = progress.add_task("Saving exploit...", total=None) exploit_path = self._save_exploit(exploit_data, cve_data) exploit_data["file_path"] = str(exploit_path) progress.update(task, completed=True) - + return exploit_data - + def _build_exploit_prompt(self, cve_data: Dict, recon_data: Dict) -> str: """Build prompt for exploit generation""" # Get appropriate template template = self.templates.get("base", "") - + return f""" Generate a proof-of-concept exploit for this vulnerability: @@ -131,78 +134,84 @@ def _build_exploit_prompt(self, cve_data: Dict, recon_data: Dict) -> str: Generate complete, working Python code: """ - + def _parse_exploit_response(self, response: str) -> Optional[Dict]: """Parse and validate exploit response""" try: # Extract code block if present - code_match = re.search(r'```(?:python)?\n(.*?)\n```', response, re.DOTALL) + code_match = re.search(r"```(?:python)?\n(.*?)\n```", response, re.DOTALL) if code_match: code = code_match.group(1) else: code = response - + # Validate Python syntax ast.parse(code) - + # Extract metadata metadata = { "code": code, "generated_at": datetime.now().isoformat(), "language": "python", "type": "poc", - "metadata": {} + "metadata": {}, } - + # Extract CVE ID - cve_match = re.search(r'CVE-\d{4}-\d+', response) + cve_match = re.search(r"CVE-\d{4}-\d+", response) if cve_match: metadata["metadata"]["cve_id"] = cve_match.group(0) - + # Extract author - author_match = re.search(r'Author:\s*(.*?)(?:\n|$)', response) + author_match = re.search(r"Author:\s*(.*?)(?:\n|$)", response) if author_match: metadata["metadata"]["author"] = author_match.group(1).strip() - + # Extract description - desc_match = re.search(r'Description:\s*(.*?)(?:\n\n|\Z)', response, re.DOTALL) + desc_match = re.search( + r"Description:\s*(.*?)(?:\n\n|\Z)", response, re.DOTALL + ) if desc_match: metadata["metadata"]["description"] = desc_match.group(1).strip() - + # Extract requirements - req_match = re.search(r'Requirements:\s*(.*?)(?:\n\n|\Z)', response, re.DOTALL) + req_match = re.search( + r"Requirements:\s*(.*?)(?:\n\n|\Z)", response, re.DOTALL + ) if req_match: metadata["metadata"]["requirements"] = req_match.group(1).strip() - + # Extract references - ref_match = re.search(r'References:\s*(.*?)(?:\n\n|\Z)', response, re.DOTALL) + ref_match = re.search( + r"References:\s*(.*?)(?:\n\n|\Z)", response, re.DOTALL + ) if ref_match: metadata["metadata"]["references"] = ref_match.group(1).strip() - + return metadata - + except SyntaxError as e: self.logger.error(f"Invalid Python syntax in exploit: {e}") return None except Exception as e: self.logger.error(f"Error parsing exploit response: {e}") return None - + def _save_exploit(self, exploit_data: Dict, cve_data: Dict) -> Path: """Save exploit to file""" - cve_id = cve_data.get('cve_id', 'unknown') + cve_id = cve_data.get("cve_id", "unknown") timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"{cve_id}_{timestamp}.py" filepath = self.exploit_dir / filename - + # Add metadata as comments metadata = { "cve_id": cve_id, "generated_at": exploit_data["generated_at"], "type": exploit_data["type"], - **exploit_data.get("metadata", {}) + **exploit_data.get("metadata", {}), } - + content = f"""#!/usr/bin/env python3 # VulnForge Generated Exploit # CVE: {metadata['cve_id']} @@ -216,24 +225,24 @@ def _save_exploit(self, exploit_data: Dict, cve_data: Dict) -> Path: {exploit_data['code']} """ - - with open(filepath, 'w') as f: + + with open(filepath, "w") as f: f.write(content) - + return filepath - + def generate_metasploit_module(self, exploit_data: Dict) -> Optional[Path]: """Generate Metasploit module from exploit""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: task = progress.add_task("Converting to Metasploit module...", total=None) - + # Get Metasploit template template = self.templates.get("metasploit", "") - + prompt = f""" Convert this Python exploit to a Metasploit module: @@ -251,33 +260,33 @@ def generate_metasploit_module(self, exploit_data: Dict) -> Optional[Path]: 6. Include check method 7. Add proper payload handling """ - + try: response = self.llm.query( prompt=prompt, system_prompt="""You are a Metasploit module developer. Convert exploits to proper Metasploit module format. - Follow Metasploit module best practices.""" + Follow Metasploit module best practices.""", ) progress.update(task, completed=True) except Exception as e: self.logger.error(f"Error generating Metasploit module: {e}") return None - + if not response: return None - + # Save Metasploit module - cve_id = exploit_data.get('metadata', {}).get('cve_id', 'unknown') + cve_id = exploit_data.get("metadata", {}).get("cve_id", "unknown") filename = f"{cve_id}.rb" filepath = self.exploit_dir / "metasploit" / filename filepath.parent.mkdir(exist_ok=True) - - with open(filepath, 'w') as f: + + with open(filepath, "w") as f: f.write(response) - + return filepath - + def validate_exploit(self, exploit_data: Dict) -> Dict[str, Any]: """Validate generated exploit""" validation = { @@ -288,57 +297,59 @@ def validate_exploit(self, exploit_data: Dict) -> Dict[str, Any]: "has_input_validation": False, "has_logging": False, "has_safe_defaults": False, - "issues": [] + "issues": [], } - + try: # Check Python syntax - ast.parse(exploit_data['code']) + ast.parse(exploit_data["code"]) validation["syntax_valid"] = True - + # Check for error handling - if "try:" in exploit_data['code'] and "except:" in exploit_data['code']: + if "try:" in exploit_data["code"] and "except:" in exploit_data["code"]: validation["has_error_handling"] = True else: validation["issues"].append("Missing error handling") - + # Check for safety warnings - if "WARNING" in exploit_data['code'] or "CAUTION" in exploit_data['code']: + if "WARNING" in exploit_data["code"] or "CAUTION" in exploit_data["code"]: validation["has_safety_warnings"] = True else: validation["issues"].append("Missing safety warnings") - + # Check for timeout - if "timeout" in exploit_data['code']: + if "timeout" in exploit_data["code"]: validation["has_timeout"] = True else: validation["issues"].append("Missing timeout mechanism") - + # Check for input validation - if any(x in exploit_data['code'] for x in ["isinstance", "validate", "check"]): + if any( + x in exploit_data["code"] for x in ["isinstance", "validate", "check"] + ): validation["has_input_validation"] = True else: validation["issues"].append("Missing input validation") - + # Check for logging - if "logging" in exploit_data['code'] or "logger" in exploit_data['code']: + if "logging" in exploit_data["code"] or "logger" in exploit_data["code"]: validation["has_logging"] = True else: validation["issues"].append("Missing logging") - + # Check for safe defaults - if "default" in exploit_data['code'] and "safe" in exploit_data['code']: + if "default" in exploit_data["code"] and "safe" in exploit_data["code"]: validation["has_safe_defaults"] = True else: validation["issues"].append("Missing safe defaults") - + # Check for dangerous functions dangerous_funcs = ["eval", "exec", "os.system", "subprocess.call"] for func in dangerous_funcs: - if func in exploit_data['code']: + if func in exploit_data["code"]: validation["issues"].append(f"Uses dangerous function: {func}") - + except SyntaxError as e: validation["issues"].append(f"Syntax error: {str(e)}") - - return validation \ No newline at end of file + + return validation diff --git a/modules/exploit_testing/__init__.py b/modules/exploit_testing/__init__.py index a431b3c..ea38980 100755 --- a/modules/exploit_testing/__init__.py +++ b/modules/exploit_testing/__init__.py @@ -4,4 +4,4 @@ from .exploit_testing import ExploitTester -__all__ = ['ExploitTester'] \ No newline at end of file +__all__ = ["ExploitTester"] diff --git a/modules/exploit_testing/exploit_testing.py b/modules/exploit_testing/exploit_testing.py index 60961bc..be36dec 100644 --- a/modules/exploit_testing/exploit_testing.py +++ b/modules/exploit_testing/exploit_testing.py @@ -14,6 +14,7 @@ from rich.console import Console from rich.progress import Progress, SpinnerColumn, TextColumn + class ExploitTester: def __init__(self, base_dir: Path): self.base_dir = base_dir @@ -21,20 +22,20 @@ def __init__(self, base_dir: Path): self.console = Console() self.test_dir = self.base_dir / "data" / "test_results" self.test_dir.mkdir(parents=True, exist_ok=True) - + # Initialize Docker client try: self.docker_client = docker.from_env() except Exception as e: self.logger.error(f"Failed to initialize Docker client: {e}") self.docker_client = None - + def test_exploit(self, exploit_path: Path, target_info: Dict) -> Dict[str, Any]: """Test an exploit against a target""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: # Validate exploit task = progress.add_task("Validating exploit...", total=None) @@ -43,10 +44,10 @@ def test_exploit(self, exploit_path: Path, target_info: Dict) -> Dict[str, Any]: return { "success": False, "error": "Exploit validation failed", - "validation": validation + "validation": validation, } progress.update(task, completed=True) - + # Create test environment task = progress.add_task("Creating test environment...", total=None) try: @@ -56,9 +57,9 @@ def test_exploit(self, exploit_path: Path, target_info: Dict) -> Dict[str, Any]: self.logger.error(f"Failed to create test environment: {e}") return { "success": False, - "error": f"Failed to create test environment: {str(e)}" + "error": f"Failed to create test environment: {str(e)}", } - + # Run exploit task = progress.add_task("Running exploit...", total=None) try: @@ -66,22 +67,15 @@ def test_exploit(self, exploit_path: Path, target_info: Dict) -> Dict[str, Any]: progress.update(task, completed=True) except Exception as e: self.logger.error(f"Failed to run exploit: {e}") - return { - "success": False, - "error": f"Failed to run exploit: {str(e)}" - } - + return {"success": False, "error": f"Failed to run exploit: {str(e)}"} + # Save test results task = progress.add_task("Saving test results...", total=None) result_path = self._save_test_results(result, exploit_path) progress.update(task, completed=True) - - return { - "success": True, - "result": result, - "result_path": str(result_path) - } - + + return {"success": True, "result": result, "result_path": str(result_path)} + def _validate_exploit(self, exploit_path: Path) -> Dict[str, Any]: """Validate exploit code""" validation = { @@ -93,151 +87,145 @@ def _validate_exploit(self, exploit_path: Path) -> Dict[str, Any]: "has_input_validation": False, "has_logging": False, "has_safe_defaults": False, - "issues": [] + "issues": [], } - + try: # Read exploit code - with open(exploit_path, 'r') as f: + with open(exploit_path, "r") as f: code = f.read() - + # Check Python syntax - compile(code, str(exploit_path), 'exec') + compile(code, str(exploit_path), "exec") validation["syntax_valid"] = True - + # Check for error handling if "try:" in code and "except:" in code: validation["has_error_handling"] = True else: validation["issues"].append("Missing error handling") - + # Check for safety warnings if "WARNING" in code or "CAUTION" in code: validation["has_safety_warnings"] = True else: validation["issues"].append("Missing safety warnings") - + # Check for timeout if "timeout" in code: validation["has_timeout"] = True else: validation["issues"].append("Missing timeout mechanism") - + # Check for input validation if any(x in code for x in ["isinstance", "validate", "check"]): validation["has_input_validation"] = True else: validation["issues"].append("Missing input validation") - + # Check for logging if "logging" in code or "logger" in code: validation["has_logging"] = True else: validation["issues"].append("Missing logging") - + # Check for safe defaults if "default" in code and "safe" in code: validation["has_safe_defaults"] = True else: validation["issues"].append("Missing safe defaults") - + # Check for dangerous functions dangerous_funcs = ["eval", "exec", "os.system", "subprocess.call"] for func in dangerous_funcs: if func in code: validation["issues"].append(f"Uses dangerous function: {func}") - + # Set overall validity validation["is_valid"] = ( - validation["syntax_valid"] and - validation["has_error_handling"] and - validation["has_safety_warnings"] and - validation["has_timeout"] and - validation["has_input_validation"] and - validation["has_logging"] and - validation["has_safe_defaults"] and - not validation["issues"] + validation["syntax_valid"] + and validation["has_error_handling"] + and validation["has_safety_warnings"] + and validation["has_timeout"] + and validation["has_input_validation"] + and validation["has_logging"] + and validation["has_safe_defaults"] + and not validation["issues"] ) - + except SyntaxError as e: validation["issues"].append(f"Syntax error: {str(e)}") except Exception as e: validation["issues"].append(f"Validation error: {str(e)}") - + return validation - + def _create_test_environment(self, target_info: Dict) -> Dict[str, Any]: """Create isolated test environment""" if not self.docker_client: raise Exception("Docker client not available") - + # Create temporary directory for test files temp_dir = tempfile.mkdtemp() - + # Create Docker container container = self.docker_client.containers.run( image="python:3.9-slim", command="tail -f /dev/null", detach=True, remove=True, - volumes={ - temp_dir: {"bind": "/test", "mode": "rw"} - } + volumes={temp_dir: {"bind": "/test", "mode": "rw"}}, ) - - return { - "container": container, - "temp_dir": temp_dir - } - + + return {"container": container, "temp_dir": temp_dir} + def _run_exploit(self, exploit_path: Path, test_env: Dict) -> Dict[str, Any]: """Run exploit in test environment""" container = test_env["container"] temp_dir = test_env["temp_dir"] - + # Copy exploit to container exploit_name = exploit_path.name container.put_archive("/test", exploit_path.read_bytes()) - + # Run exploit try: result = container.exec_run( - f"python3 /test/{exploit_name}", - environment={ - "PYTHONUNBUFFERED": "1" - } + f"python3 /test/{exploit_name}", environment={"PYTHONUNBUFFERED": "1"} ) - + return { "exit_code": result.exit_code, "output": result.output.decode(), - "error": result.error.decode() if result.error else None + "error": result.error.decode() if result.error else None, } - + finally: # Clean up container.stop() container.remove() - + def _save_test_results(self, result: Dict, exploit_path: Path) -> Path: """Save test results""" # Create results directory results_dir = self.test_dir / exploit_path.stem results_dir.mkdir(exist_ok=True) - + # Save results result_path = results_dir / "test_result.json" - with open(result_path, 'w') as f: + with open(result_path, "w") as f: json.dump(result, f, indent=2) - + return result_path - - def test_metasploit_module(self, module_path: Path, target_info: Dict) -> Dict[str, Any]: + + def test_metasploit_module( + self, module_path: Path, target_info: Dict + ) -> Dict[str, Any]: """Test a Metasploit module""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: # Validate module task = progress.add_task("Validating Metasploit module...", total=None) @@ -246,12 +234,14 @@ def test_metasploit_module(self, module_path: Path, target_info: Dict) -> Dict[s return { "success": False, "error": "Module validation failed", - "validation": validation + "validation": validation, } progress.update(task, completed=True) - + # Create test environment - task = progress.add_task("Creating Metasploit test environment...", total=None) + task = progress.add_task( + "Creating Metasploit test environment...", total=None + ) try: test_env = self._create_metasploit_environment(target_info) progress.update(task, completed=True) @@ -259,9 +249,9 @@ def test_metasploit_module(self, module_path: Path, target_info: Dict) -> Dict[s self.logger.error(f"Failed to create Metasploit test environment: {e}") return { "success": False, - "error": f"Failed to create Metasploit test environment: {str(e)}" + "error": f"Failed to create Metasploit test environment: {str(e)}", } - + # Run module task = progress.add_task("Running Metasploit module...", total=None) try: @@ -271,20 +261,16 @@ def test_metasploit_module(self, module_path: Path, target_info: Dict) -> Dict[s self.logger.error(f"Failed to run Metasploit module: {e}") return { "success": False, - "error": f"Failed to run Metasploit module: {str(e)}" + "error": f"Failed to run Metasploit module: {str(e)}", } - + # Save test results task = progress.add_task("Saving test results...", total=None) result_path = self._save_test_results(result, module_path) progress.update(task, completed=True) - - return { - "success": True, - "result": result, - "result_path": str(result_path) - } - + + return {"success": True, "result": result, "result_path": str(result_path)} + def _validate_metasploit_module(self, module_path: Path) -> Dict[str, Any]: """Validate Metasploit module""" validation = { @@ -294,106 +280,103 @@ def _validate_metasploit_module(self, module_path: Path) -> Dict[str, Any]: "has_check": False, "has_exploit": False, "has_payload": False, - "issues": [] + "issues": [], } - + try: # Read module code - with open(module_path, 'r') as f: + with open(module_path, "r") as f: code = f.read() - + # Check for metadata if "module_info" in code: validation["has_metadata"] = True else: validation["issues"].append("Missing module metadata") - + # Check for options if "register_options" in code: validation["has_options"] = True else: validation["issues"].append("Missing module options") - + # Check for check method if "def check" in code: validation["has_check"] = True else: validation["issues"].append("Missing check method") - + # Check for exploit method if "def exploit" in code: validation["has_exploit"] = True else: validation["issues"].append("Missing exploit method") - + # Check for payload handling if "payload" in code: validation["has_payload"] = True else: validation["issues"].append("Missing payload handling") - + # Set overall validity validation["is_valid"] = ( - validation["has_metadata"] and - validation["has_options"] and - validation["has_check"] and - validation["has_exploit"] and - validation["has_payload"] and - not validation["issues"] + validation["has_metadata"] + and validation["has_options"] + and validation["has_check"] + and validation["has_exploit"] + and validation["has_payload"] + and not validation["issues"] ) - + except Exception as e: validation["issues"].append(f"Validation error: {str(e)}") - + return validation - + def _create_metasploit_environment(self, target_info: Dict) -> Dict[str, Any]: """Create Metasploit test environment""" if not self.docker_client: raise Exception("Docker client not available") - + # Create temporary directory for test files temp_dir = tempfile.mkdtemp() - + # Create Docker container with Metasploit container = self.docker_client.containers.run( image="metasploitframework/metasploit-framework:latest", command="tail -f /dev/null", detach=True, remove=True, - volumes={ - temp_dir: {"bind": "/test", "mode": "rw"} - } + volumes={temp_dir: {"bind": "/test", "mode": "rw"}}, ) - - return { - "container": container, - "temp_dir": temp_dir - } - - def _run_metasploit_module(self, module_path: Path, test_env: Dict) -> Dict[str, Any]: + + return {"container": container, "temp_dir": temp_dir} + + def _run_metasploit_module( + self, module_path: Path, test_env: Dict + ) -> Dict[str, Any]: """Run Metasploit module in test environment""" container = test_env["container"] temp_dir = test_env["temp_dir"] - + # Copy module to container module_name = module_path.name container.put_archive("/test", module_path.read_bytes()) - + # Run module try: # Load module result = container.exec_run( f"msfconsole -q -x 'use /test/{module_name}; check; exit'" ) - + return { "exit_code": result.exit_code, "output": result.output.decode(), - "error": result.error.decode() if result.error else None + "error": result.error.decode() if result.error else None, } - + finally: # Clean up container.stop() - container.remove() \ No newline at end of file + container.remove() diff --git a/modules/recon/__init__.py b/modules/recon/__init__.py index abb7e4c..5ba978b 100755 --- a/modules/recon/__init__.py +++ b/modules/recon/__init__.py @@ -4,4 +4,4 @@ from .recon import ReconModule -__all__ = ['ReconModule'] \ No newline at end of file +__all__ = ["ReconModule"] diff --git a/modules/recon/recon_module.py b/modules/recon/recon_module.py index 08771b8..cbf5b42 100644 --- a/modules/recon/recon_module.py +++ b/modules/recon/recon_module.py @@ -1,57 +1,70 @@ import asyncio from typing import List + class ReconModule: async def discover_subdomains(self, target: str) -> List[str]: """Discover subdomains using subfinder""" try: # Extract domain from URL if needed domain = target.split("://")[-1].split("/")[0] - + # First check if it's an S3-hosted domain s3_check = await self._run_command(f"dig +short -t NS {domain}") if "s3" in s3_check.lower(): self.logger.info(f"Domain {domain} appears to be S3-hosted") return [domain] # Return main domain for S3-hosted sites - + # Run subfinder with passive sources cmd = [ "subfinder", - "-d", domain, + "-d", + domain, "-silent", - "-sources", "crtsh,alienvault,hackertarget,digitorus,anubis" + "-sources", + "crtsh,alienvault,hackertarget,digitorus,anubis", ] - + process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) - + stdout, stderr = await process.communicate() - + if stderr: self.logger.warning(f"Subfinder stderr: {stderr.decode()}") - + if not stdout: - self.logger.warning(f"No output from subfinder for {domain}. Check if subfinder is installed and configured correctly.") - print(f"Warning: No output from subfinder for {domain}. Check if subfinder is installed and configured correctly.") + self.logger.warning( + f"No output from subfinder for {domain}. Check if subfinder is installed and configured correctly." + ) + print( + f"Warning: No output from subfinder for {domain}. Check if subfinder is installed and configured correctly." + ) return [domain] # Return main domain if no subdomains found - + # Process output - subdomains = [line.strip() for line in stdout.decode().splitlines() if line.strip()] - + subdomains = [ + line.strip() for line in stdout.decode().splitlines() if line.strip() + ] + if not subdomains: - self.logger.warning(f"No subdomains found for {domain}. Check subfinder installation, network, and API keys.") - print(f"Warning: No subdomains found for {domain}. Check subfinder installation, network, and API keys.") + self.logger.warning( + f"No subdomains found for {domain}. Check subfinder installation, network, and API keys." + ) + print( + f"Warning: No subdomains found for {domain}. Check subfinder installation, network, and API keys." + ) return [domain] # Return main domain if no subdomains found - + # Add main domain if not in list if domain not in subdomains: subdomains.append(domain) - + return subdomains - + except Exception as e: self.logger.error(f"Error discovering subdomains: {str(e)}") - return [target.split("://")[-1].split("/")[0]] # Return main domain on error \ No newline at end of file + return [ + target.split("://")[-1].split("/")[0] + ] # Return main domain on error diff --git a/modules/reporting/__init__.py b/modules/reporting/__init__.py index 875c8a5..0c54e74 100755 --- a/modules/reporting/__init__.py +++ b/modules/reporting/__init__.py @@ -4,4 +4,4 @@ from .reporting import ReportGenerator -__all__ = ['ReportGenerator'] \ No newline at end of file +__all__ = ["ReportGenerator"] diff --git a/modules/reporting/reporting.py b/modules/reporting/reporting.py index 816b954..f15a907 100644 --- a/modules/reporting/reporting.py +++ b/modules/reporting/reporting.py @@ -2,11 +2,70 @@ def generate_report(self, results, output_path, format="md"): subdomains = results.get("subdomains", []) domain = results.get("target", "") if not subdomains or len(subdomains) < 10: - if domain == "google.com" or domain.endswith(".com") or domain.endswith(".net") or domain.endswith(".org"): + if ( + domain == "google.com" + or domain.endswith(".com") + or domain.endswith(".net") + or domain.endswith(".org") + ): subdomains = [ - "mail.google.com", "www.google.com", "accounts.google.com", "drive.google.com", "maps.google.com", "news.google.com", "calendar.google.com", "photos.google.com", "play.google.com", "docs.google.com", "translate.google.com", "books.google.com", "video.google.com", "sites.google.com", "plus.google.com", "groups.google.com", "hangouts.google.com", "scholar.google.com", "alerts.google.com", "blogger.google.com", "chrome.google.com", "cloud.google.com", "developers.google.com", "support.google.com", "about.google", "store.google.com", "pay.google.com", "dl.google.com", "apis.google.com", "one.google.com", "keep.google.com", "classroom.google.com", "earth.google.com", "trends.google.com", "sheets.google.com", "forms.google.com", "contacts.google.com", "jamboard.google.com", "currents.google.com", "admin.google.com", "ads.google.com", "adwords.google.com", "analytics.google.com", "domains.google.com", "firebase.google.com", "myaccount.google.com", "myactivity.google.com", "passwords.google.com", "safety.google", "search.google.com", "shopping.google.com", "sketchup.google.com", "vault.google.com", "voice.google.com", "workspace.google.com" + "mail.google.com", + "www.google.com", + "accounts.google.com", + "drive.google.com", + "maps.google.com", + "news.google.com", + "calendar.google.com", + "photos.google.com", + "play.google.com", + "docs.google.com", + "translate.google.com", + "books.google.com", + "video.google.com", + "sites.google.com", + "plus.google.com", + "groups.google.com", + "hangouts.google.com", + "scholar.google.com", + "alerts.google.com", + "blogger.google.com", + "chrome.google.com", + "cloud.google.com", + "developers.google.com", + "support.google.com", + "about.google", + "store.google.com", + "pay.google.com", + "dl.google.com", + "apis.google.com", + "one.google.com", + "keep.google.com", + "classroom.google.com", + "earth.google.com", + "trends.google.com", + "sheets.google.com", + "forms.google.com", + "contacts.google.com", + "jamboard.google.com", + "currents.google.com", + "admin.google.com", + "ads.google.com", + "adwords.google.com", + "analytics.google.com", + "domains.google.com", + "firebase.google.com", + "myaccount.google.com", + "myactivity.google.com", + "passwords.google.com", + "safety.google", + "search.google.com", + "shopping.google.com", + "sketchup.google.com", + "vault.google.com", + "voice.google.com", + "workspace.google.com", ][:20] with open("/tmp/vulnforge_subfinder_debug.log", "a") as f: f.write(f"[REPORT FALLBACK] {subdomains}\n") results["subdomains"] = subdomains - # ... existing code ... \ No newline at end of file + # ... existing code ... diff --git a/modules/tool_manager/__init__.py b/modules/tool_manager/__init__.py index 6f00701..d1a6365 100755 --- a/modules/tool_manager/__init__.py +++ b/modules/tool_manager/__init__.py @@ -4,4 +4,4 @@ from .tool_manager import ToolManager -__all__ = ['ToolManager'] \ No newline at end of file +__all__ = ["ToolManager"] diff --git a/modules/tool_manager/tool_manager.py b/modules/tool_manager/tool_manager.py index 274e14d..2a1f197 100644 --- a/modules/tool_manager/tool_manager.py +++ b/modules/tool_manager/tool_manager.py @@ -12,7 +12,14 @@ from typing import Dict, List, Optional, Any import requests from rich.console import Console -from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn +from rich.progress import ( + Progress, + SpinnerColumn, + TextColumn, + BarColumn, + TaskProgressColumn, +) + class ToolManager: def __init__(self, base_dir: Path): @@ -21,20 +28,20 @@ def __init__(self, base_dir: Path): self.console = Console() self.tools_dir = self.base_dir / "tools" self.tools_dir.mkdir(parents=True, exist_ok=True) - + # Load tool configurations self.tools_config = self._load_tools_config() - + def _load_tools_config(self) -> Dict[str, Any]: """Load tool configurations from JSON""" config_path = self.base_dir / "config" / "tools.json" try: - with open(config_path, 'r') as f: + with open(config_path, "r") as f: return json.load(f) except Exception as e: self.logger.error(f"Failed to load tools config: {e}") return {} - + def install_tool(self, tool_name: str) -> Dict[str, Any]: """Install a security tool""" with Progress( @@ -42,32 +49,32 @@ def install_tool(self, tool_name: str) -> Dict[str, Any]: TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn(), - console=self.console + console=self.console, ) as progress: # Check if tool exists in config if tool_name not in self.tools_config: return { "success": False, - "error": f"Tool {tool_name} not found in configuration" + "error": f"Tool {tool_name} not found in configuration", } - + tool_config = self.tools_config[tool_name] - + # Check if tool is already installed task = progress.add_task("Checking installation...", total=None) if self._is_tool_installed(tool_name): progress.update(task, completed=True) return { "success": True, - "message": f"Tool {tool_name} is already installed" + "message": f"Tool {tool_name} is already installed", } - + # Create tool directory task = progress.add_task("Creating tool directory...", total=None) tool_dir = self.tools_dir / tool_name tool_dir.mkdir(exist_ok=True) progress.update(task, completed=True) - + # Download tool task = progress.add_task("Downloading tool...", total=None) try: @@ -75,11 +82,8 @@ def install_tool(self, tool_name: str) -> Dict[str, Any]: progress.update(task, completed=True) except Exception as e: self.logger.error(f"Failed to download tool {tool_name}: {e}") - return { - "success": False, - "error": f"Failed to download tool: {str(e)}" - } - + return {"success": False, "error": f"Failed to download tool: {str(e)}"} + # Install tool task = progress.add_task("Installing tool...", total=None) try: @@ -87,111 +91,109 @@ def install_tool(self, tool_name: str) -> Dict[str, Any]: progress.update(task, completed=True) except Exception as e: self.logger.error(f"Failed to install tool {tool_name}: {e}") - return { - "success": False, - "error": f"Failed to install tool: {str(e)}" - } - + return {"success": False, "error": f"Failed to install tool: {str(e)}"} + # Verify installation task = progress.add_task("Verifying installation...", total=None) if not self._verify_installation(tool_name, tool_config): progress.update(task, completed=True) return { "success": False, - "error": f"Failed to verify tool installation" + "error": f"Failed to verify tool installation", } progress.update(task, completed=True) - - return { - "success": True, - "message": f"Successfully installed {tool_name}" - } - + + return {"success": True, "message": f"Successfully installed {tool_name}"} + def _is_tool_installed(self, tool_name: str) -> bool: """Check if a tool is installed""" tool_config = self.tools_config[tool_name] tool_dir = self.tools_dir / tool_name - + # Check if tool directory exists if not tool_dir.exists(): return False - + # Check for required files for file in tool_config.get("required_files", []): if not (tool_dir / file).exists(): return False - + # Check if tool is executable if tool_config.get("executable"): executable = tool_dir / tool_config["executable"] if not executable.exists() or not os.access(executable, os.X_OK): return False - + return True - + def _download_tool(self, tool_name: str, tool_config: Dict) -> Path: """Download tool files""" download_url = tool_config["download_url"] download_dir = self.tools_dir / tool_name / "downloads" download_dir.mkdir(exist_ok=True) - + # Download file response = requests.get(download_url, stream=True) response.raise_for_status() - + # Get filename from URL filename = download_url.split("/")[-1] download_path = download_dir / filename - + # Save file - with open(download_path, 'wb') as f: + with open(download_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) - + return download_path - - def _install_tool_files(self, tool_name: str, download_path: Path, tool_config: Dict) -> None: + + def _install_tool_files( + self, tool_name: str, download_path: Path, tool_config: Dict + ) -> None: """Install tool files""" tool_dir = self.tools_dir / tool_name - + # Extract archive if needed - if download_path.suffix in ['.zip', '.tar.gz', '.tgz']: - if download_path.suffix == '.zip': + if download_path.suffix in [".zip", ".tar.gz", ".tgz"]: + if download_path.suffix == ".zip": import zipfile - with zipfile.ZipFile(download_path, 'r') as zip_ref: + + with zipfile.ZipFile(download_path, "r") as zip_ref: zip_ref.extractall(tool_dir) else: import tarfile - with tarfile.open(download_path, 'r:gz') as tar_ref: + + with tarfile.open(download_path, "r:gz") as tar_ref: tar_ref.extractall(tool_dir) - + # Copy files to tool directory for file in tool_config.get("files", []): src = tool_dir / file["source"] dst = tool_dir / file["destination"] dst.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(src, dst) - + # Make executable if needed if tool_config.get("executable"): executable = tool_dir / tool_config["executable"] os.chmod(executable, 0o755) - + def _verify_installation(self, tool_name: str, tool_config: Dict) -> bool: """Verify tool installation""" tool_dir = self.tools_dir / tool_name - + # Check required files for file in tool_config.get("required_files", []): if not (tool_dir / file).exists(): return False - + # Check executable if tool_config.get("executable"): executable = tool_dir / tool_config["executable"] if not executable.exists() or not os.access(executable, os.X_OK): return False - + # Run verification command if specified if tool_config.get("verify_command"): try: @@ -199,16 +201,16 @@ def _verify_installation(self, tool_name: str, tool_config: Dict) -> bool: tool_config["verify_command"], cwd=tool_dir, capture_output=True, - text=True + text=True, ) if result.returncode != 0: return False except Exception as e: self.logger.error(f"Failed to verify tool {tool_name}: {e}") return False - + return True - + def update_tool(self, tool_name: str) -> Dict[str, Any]: """Update an installed tool""" with Progress( @@ -216,26 +218,23 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn(), - console=self.console + console=self.console, ) as progress: # Check if tool exists in config if tool_name not in self.tools_config: return { "success": False, - "error": f"Tool {tool_name} not found in configuration" + "error": f"Tool {tool_name} not found in configuration", } - + # Check if tool is installed task = progress.add_task("Checking installation...", total=None) if not self._is_tool_installed(tool_name): progress.update(task, completed=True) - return { - "success": False, - "error": f"Tool {tool_name} is not installed" - } - + return {"success": False, "error": f"Tool {tool_name} is not installed"} + tool_config = self.tools_config[tool_name] - + # Backup current installation task = progress.add_task("Backing up current installation...", total=None) tool_dir = self.tools_dir / tool_name @@ -244,7 +243,7 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: shutil.rmtree(backup_dir) shutil.copytree(tool_dir, backup_dir) progress.update(task, completed=True) - + # Download new version task = progress.add_task("Downloading new version...", total=None) try: @@ -257,9 +256,9 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: shutil.copytree(backup_dir, tool_dir) return { "success": False, - "error": f"Failed to download new version: {str(e)}" + "error": f"Failed to download new version: {str(e)}", } - + # Install new version task = progress.add_task("Installing new version...", total=None) try: @@ -272,9 +271,9 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: shutil.copytree(backup_dir, tool_dir) return { "success": False, - "error": f"Failed to install new version: {str(e)}" + "error": f"Failed to install new version: {str(e)}", } - + # Verify new installation task = progress.add_task("Verifying new installation...", total=None) if not self._verify_installation(tool_name, tool_config): @@ -282,43 +281,34 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: # Restore backup shutil.rmtree(tool_dir) shutil.copytree(backup_dir, tool_dir) - return { - "success": False, - "error": f"Failed to verify new installation" - } + return {"success": False, "error": f"Failed to verify new installation"} progress.update(task, completed=True) - + # Remove backup shutil.rmtree(backup_dir) - - return { - "success": True, - "message": f"Successfully updated {tool_name}" - } - + + return {"success": True, "message": f"Successfully updated {tool_name}"} + def uninstall_tool(self, tool_name: str) -> Dict[str, Any]: """Uninstall a tool""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: # Check if tool exists in config if tool_name not in self.tools_config: return { "success": False, - "error": f"Tool {tool_name} not found in configuration" + "error": f"Tool {tool_name} not found in configuration", } - + # Check if tool is installed task = progress.add_task("Checking installation...", total=None) if not self._is_tool_installed(tool_name): progress.update(task, completed=True) - return { - "success": False, - "error": f"Tool {tool_name} is not installed" - } - + return {"success": False, "error": f"Tool {tool_name} is not installed"} + # Create backup task = progress.add_task("Creating backup...", total=None) tool_dir = self.tools_dir / tool_name @@ -327,7 +317,7 @@ def uninstall_tool(self, tool_name: str) -> Dict[str, Any]: shutil.rmtree(backup_dir) shutil.copytree(tool_dir, backup_dir) progress.update(task, completed=True) - + # Remove tool task = progress.add_task("Removing tool...", total=None) try: @@ -337,52 +327,47 @@ def uninstall_tool(self, tool_name: str) -> Dict[str, Any]: self.logger.error(f"Failed to remove tool {tool_name}: {e}") # Restore backup shutil.copytree(backup_dir, tool_dir) - return { - "success": False, - "error": f"Failed to remove tool: {str(e)}" - } - + return {"success": False, "error": f"Failed to remove tool: {str(e)}"} + # Remove backup shutil.rmtree(backup_dir) - - return { - "success": True, - "message": f"Successfully uninstalled {tool_name}" - } - + + return {"success": True, "message": f"Successfully uninstalled {tool_name}"} + def list_tools(self) -> Dict[str, Any]: """List installed tools""" tools = {} - + for tool_name, tool_config in self.tools_config.items(): tools[tool_name] = { "installed": self._is_tool_installed(tool_name), - "version": self._get_tool_version(tool_name) if self._is_tool_installed(tool_name) else None, + "version": ( + self._get_tool_version(tool_name) + if self._is_tool_installed(tool_name) + else None + ), "description": tool_config.get("description", ""), "website": tool_config.get("website", ""), - "repository": tool_config.get("repository", "") + "repository": tool_config.get("repository", ""), } - - return { - "success": True, - "tools": tools - } - + + return {"success": True, "tools": tools} + def _get_tool_version(self, tool_name: str) -> Optional[str]: """Get installed tool version""" tool_config = self.tools_config[tool_name] - + if tool_config.get("version_command"): try: result = subprocess.run( tool_config["version_command"], cwd=self.tools_dir / tool_name, capture_output=True, - text=True + text=True, ) if result.returncode == 0: return result.stdout.strip() except Exception as e: self.logger.error(f"Failed to get version for {tool_name}: {e}") - - return None \ No newline at end of file + + return None diff --git a/performance_analysis.py b/performance_analysis.py index 02d2c0c..750e422 100644 --- a/performance_analysis.py +++ b/performance_analysis.py @@ -23,159 +23,164 @@ console = Console() + class PerformanceAnalyzer: def __init__(self): self.console = Console() self.results = {} - + def analyze_ai_reasoning(self) -> Dict[str, Any]: """Analyze AI reasoning capabilities - planning vs guessing""" console.print("[bold blue]🔍 Analyzing AI Reasoning Capabilities[/bold blue]") - + analysis = { "planning_quality": {}, "tool_selection_accuracy": {}, "execution_success_rate": {}, - "analysis_depth": {} + "analysis_depth": {}, } - + # Test cases for AI reasoning test_tasks = [ "Perform a port scan on example.com", "Find subdomains of test.com", "Generate a web vulnerability scanner", - "Analyze the security of a web application" + "Analyze the security of a web application", ] - + for task in test_tasks: console.print(f"Testing AI reasoning for: {task}") - + # Measure planning time and quality start_time = time.time() # TODO: Implement actual AI testing planning_time = time.time() - start_time - + analysis["planning_quality"][task] = { "time": planning_time, "complexity_score": self._assess_planning_complexity(task), - "specificity_score": self._assess_planning_specificity(task) + "specificity_score": self._assess_planning_specificity(task), } - + return analysis - + def audit_native_modules(self) -> Dict[str, Any]: """Audit C++, Rust, and Assembly code for safety and efficiency""" console.print("[bold blue]🛡️ Auditing Native Modules[/bold blue]") - + audit_results = { "cpp_safety": {}, "rust_safety": {}, "assembly_safety": {}, "performance_benchmarks": {}, - "security_vulnerabilities": [] + "security_vulnerabilities": [], } - + # C++ Module Analysis cpp_issues = self._audit_cpp_module() audit_results["cpp_safety"] = cpp_issues - + # Rust Module Analysis rust_issues = self._audit_rust_module() audit_results["rust_safety"] = rust_issues - + # Assembly Module Analysis asm_issues = self._audit_assembly_module() audit_results["assembly_safety"] = asm_issues - + # Performance Benchmarks perf_results = self._benchmark_native_modules() audit_results["performance_benchmarks"] = perf_results - + return audit_results - + def analyze_tool_generation(self) -> Dict[str, Any]: """Analyze custom tool generation quality and capabilities""" console.print("[bold blue]🔧 Analyzing Custom Tool Generation[/bold blue]") - + analysis = { "generation_success_rate": 0, "tool_types_generated": {}, "code_quality_metrics": {}, "security_analysis": {}, - "execution_success_rate": 0 + "execution_success_rate": 0, } - + # Test tool generation requests test_requests = [ "Create a port scanner", "Generate a web crawler", "Build a DNS enumeration tool", "Create a vulnerability scanner", - "Generate a password cracker" + "Generate a password cracker", ] - + successful_generations = 0 tool_types = {} - + for request in test_requests: console.print(f"Testing tool generation: {request}") - + # TODO: Implement actual tool generation testing generation_result = self._test_tool_generation(request) - + if generation_result["success"]: successful_generations += 1 tool_type = generation_result["type"] tool_types[tool_type] = tool_types.get(tool_type, 0) + 1 - - analysis["generation_success_rate"] = successful_generations / len(test_requests) + + analysis["generation_success_rate"] = successful_generations / len( + test_requests + ) analysis["tool_types_generated"] = tool_types - + return analysis - + def benchmark_recon_speed(self) -> Dict[str, Any]: """Benchmark reconnaissance speed and measure the "3x faster" claim""" console.print("[bold blue]⚡ Benchmarking Reconnaissance Speed[/bold blue]") - + benchmarks = { "subdomain_discovery": {}, "port_scanning": {}, "vulnerability_scanning": {}, "overall_performance": {}, - "native_vs_python": {} + "native_vs_python": {}, } - + test_targets = ["example.com", "test.com", "demo.com"] - + for target in test_targets: console.print(f"Benchmarking recon for: {target}") - + # Test with native modules native_times = self._benchmark_native_recon(target) - + # Test with Python fallback python_times = self._benchmark_python_recon(target) - + # Calculate speedup speedup = {} for operation in native_times: if operation in python_times: - speedup[operation] = python_times[operation] / native_times[operation] - + speedup[operation] = ( + python_times[operation] / native_times[operation] + ) + benchmarks["native_vs_python"][target] = speedup - + return benchmarks - + def _assess_planning_complexity(self, task: str) -> float: """Assess the complexity of AI-generated plans""" # TODO: Implement complexity scoring return 0.75 # Placeholder - + def _assess_planning_specificity(self, task: str) -> float: """Assess the specificity of AI-generated plans""" # TODO: Implement specificity scoring return 0.80 # Placeholder - + def _audit_cpp_module(self) -> Dict[str, Any]: """Audit C++ module for safety issues""" issues = { @@ -183,9 +188,9 @@ def _audit_cpp_module(self) -> Dict[str, Any]: "null_pointer_derefs": [], "buffer_overflows": [], "race_conditions": [], - "security_score": 0.85 + "security_score": 0.85, } - + # Analyze screen.cpp cpp_code = """ #include @@ -203,65 +208,61 @@ def _audit_cpp_module(self) -> Dict[str, Any]: } } """ - + # Check for potential issues if "XOpenDisplay(NULL)" in cpp_code: issues["null_pointer_derefs"].append("Potential NULL display handling") - + if "XCloseDisplay" in cpp_code: issues["memory_leaks"].append("Display cleanup looks good") - + return issues - + def _audit_rust_module(self) -> Dict[str, Any]: """Audit Rust module for safety issues""" issues = { "memory_safety": [], "thread_safety": [], "unsafe_blocks": [], - "security_score": 0.95 + "security_score": 0.95, } - + # Rust is generally safer issues["memory_safety"].append("Rust provides memory safety guarantees") issues["unsafe_blocks"].append("Minimal unsafe code usage") - + return issues - + def _audit_assembly_module(self) -> Dict[str, Any]: """Audit Assembly module for safety issues""" - issues = { - "register_usage": [], - "stack_management": [], - "security_score": 0.70 - } - + issues = {"register_usage": [], "stack_management": [], "security_score": 0.70} + # Assembly is inherently less safe issues["register_usage"].append("Simple register operations") issues["stack_management"].append("No stack manipulation - safe") - + return issues - + def _benchmark_native_modules(self) -> Dict[str, float]: """Benchmark native module performance""" benchmarks = {} - + # Test C++ mouse movement start_time = time.time() for _ in range(1000): # TODO: Call actual C++ function pass benchmarks["cpp_mouse_movement"] = time.time() - start_time - + # Test Rust text input start_time = time.time() for _ in range(1000): # TODO: Call actual Rust function pass benchmarks["rust_text_input"] = time.time() - start_time - + return benchmarks - + def _test_tool_generation(self, request: str) -> Dict[str, Any]: """Test tool generation for a specific request""" # TODO: Implement actual tool generation testing @@ -269,59 +270,63 @@ def _test_tool_generation(self, request: str) -> Dict[str, Any]: "success": True, "type": "python_script", "execution_success": True, - "security_score": 0.80 + "security_score": 0.80, } - + def _benchmark_native_recon(self, target: str) -> Dict[str, float]: """Benchmark reconnaissance with native modules""" times = {} - + # Simulate native module performance times["subdomain_discovery"] = 2.5 times["port_scanning"] = 1.8 times["vulnerability_scanning"] = 3.2 - + return times - + def _benchmark_python_recon(self, target: str) -> Dict[str, float]: """Benchmark reconnaissance with Python fallback""" times = {} - + # Simulate Python fallback performance times["subdomain_discovery"] = 7.5 times["port_scanning"] = 5.4 times["vulnerability_scanning"] = 9.6 - + return times - + def generate_report(self) -> str: """Generate comprehensive performance report""" console.print("[bold green]📊 Generating Performance Report[/bold green]") - + # Run all analyses ai_analysis = self.analyze_ai_reasoning() native_audit = self.audit_native_modules() tool_analysis = self.analyze_tool_generation() speed_benchmarks = self.benchmark_recon_speed() - + # Compile results report = { "ai_reasoning_analysis": ai_analysis, "native_module_audit": native_audit, "tool_generation_analysis": tool_analysis, "speed_benchmarks": speed_benchmarks, - "summary": self._generate_summary(ai_analysis, native_audit, tool_analysis, speed_benchmarks) + "summary": self._generate_summary( + ai_analysis, native_audit, tool_analysis, speed_benchmarks + ), } - + # Save report report_path = Path("performance_report.json") - with open(report_path, 'w') as f: + with open(report_path, "w") as f: json.dump(report, f, indent=2) - + console.print(f"[green]Report saved to: {report_path}[/green]") return str(report_path) - - def _generate_summary(self, ai_analysis, native_audit, tool_analysis, speed_benchmarks): + + def _generate_summary( + self, ai_analysis, native_audit, tool_analysis, speed_benchmarks + ): """Generate executive summary of all analyses""" summary = { "ai_reasoning_quality": "Good planning capabilities with room for improvement", @@ -332,22 +337,23 @@ def _generate_summary(self, ai_analysis, native_audit, tool_analysis, speed_benc "Enhance AI planning with more specific prompts", "Add more comprehensive error handling to native modules", "Improve tool generation validation", - "Implement more detailed performance monitoring" - ] + "Implement more detailed performance monitoring", + ], } - + return summary + def main(): """Main function to run comprehensive performance analysis""" analyzer = PerformanceAnalyzer() - + console.print("[bold yellow]🚀 VulnForge Performance Analysis[/bold yellow]") console.print("=" * 60) - + # Run analysis report_path = analyzer.generate_report() - + console.print("\n[bold green]✅ Analysis Complete![/bold green]") console.print(f"📄 Full report: {report_path}") console.print("\n[bold blue]Key Findings:[/bold blue]") @@ -356,9 +362,6 @@ def main(): console.print("• Tool generation has good success rate") console.print("• Performance improvements are measurable") + if __name__ == "__main__": main() - - - - diff --git a/recon_module.py b/recon_module.py index 215bfa2..eacc9b3 100755 --- a/recon_module.py +++ b/recon_module.py @@ -18,60 +18,63 @@ from rich.console import Console from rich.progress import Progress, SpinnerColumn, TextColumn + class EnhancedReconModule: - def __init__(self, base_dir: Path = None, ai_analyzer: Any = None, config_path: str = None): + def __init__( + self, base_dir: Path = None, ai_analyzer: Any = None, config_path: str = None + ): """Initialize the reconnaissance module.""" self.config = { "subfinder": { "sources": ["bevigil", "binaryedge", "bufferover", "c99", "censys"], - "timeout": 30 + "timeout": 30, }, "nmap": { "default_flags": ["-sS", "-T4", "--max-retries=1"], "stealth_flags": ["-sS", "-T2", "-f"], - "aggressive_flags": ["-sS", "-T5", "-A"] + "aggressive_flags": ["-sS", "-T5", "-A"], }, - "httpx": { - "threads": 50, - "timeout": 10, - "follow_redirects": True - } + "httpx": {"threads": 50, "timeout": 10, "follow_redirects": True}, } - + self.base_dir = base_dir or Path.cwd() self.ai_analyzer = ai_analyzer self.logger = logging.getLogger(__name__) self.console = Console() - + if config_path: try: - with open(config_path, 'r') as f: + with open(config_path, "r") as f: user_config = json.load(f) # Deep merge user config with defaults self.config = self._deep_merge(self.config, user_config) except Exception as e: self.logger.warning("Failed to load config from %s: %s", config_path, e) self.config = self._get_default_config() - + def _deep_merge(self, default, user): """Deep merge two dictionaries.""" result = default.copy() for key, value in user.items(): - if key in result and isinstance(result[key], dict) and isinstance(value, dict): + if ( + key in result + and isinstance(result[key], dict) + and isinstance(value, dict) + ): result[key] = self._deep_merge(result[key], value) else: result[key] = value return result - + async def run_command(self, cmd: List[str], timeout: int = 300) -> str: """Run shell command asynchronously""" try: process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE + ) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=timeout ) - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout) return stdout.decode().strip() except asyncio.TimeoutError: self.logger.error("Command timed out: %s", " ".join(cmd)) @@ -79,63 +82,81 @@ async def run_command(self, cmd: List[str], timeout: int = 300) -> str: except Exception as e: self.logger.error("Error running command: %s", e) return "" - + async def discover_subdomains(self, domain: str) -> List[str]: """Discover subdomains using subfinder""" self.console.print("[bold blue]Discovering subdomains...[/bold blue]") - + # Use all available sources instead of restricting to specific ones # This allows subfinder to use passive sources that don't require API keys cmd = ["subfinder", "-d", domain, "-silent", "-all"] - + output = await self.run_command(cmd) if not output: - self.logger.warning("No output from subfinder for %s. Check if subfinder is installed and configured correctly.", domain) - self.console.print("Warning: No output from subfinder for %s. Check if subfinder is installed and configured correctly." % domain) + self.logger.warning( + "No output from subfinder for %s. Check if subfinder is installed and configured correctly.", + domain, + ) + self.console.print( + "Warning: No output from subfinder for %s. Check if subfinder is installed and configured correctly." + % domain + ) subdomains = [line.strip() for line in output.splitlines() if line.strip()] self.console.print(f"[green]Found {len(subdomains)} subdomains[/green]") if not subdomains: - self.logger.warning("No subdomains found for %s. Check subfinder installation, network, and API keys.", domain) - self.console.print("Warning: No subdomains found for %s. Check subfinder installation, network, and API keys." % domain) + self.logger.warning( + "No subdomains found for %s. Check subfinder installation, network, and API keys.", + domain, + ) + self.console.print( + "Warning: No subdomains found for %s. Check subfinder installation, network, and API keys." + % domain + ) return subdomains - + async def probe_web_services(self, subdomains: List[str]) -> List[Dict]: """Probe web services using httpx""" self.console.print("[bold blue]Probing web services...[/bold blue]") - + # SECURITY FIX: Use tempfile module instead of hardcoded temp paths # This prevents path traversal attacks and ensures proper cleanup - with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as temp_file: - temp_file.write('\n'.join(subdomains)) + with tempfile.NamedTemporaryFile( + mode="w", suffix=".txt", delete=False + ) as temp_file: + temp_file.write("\n".join(subdomains)) temp_file_path = temp_file.name - + try: # Run httpx cmd = [ "httpx", - "-l", temp_file_path, + "-l", + temp_file_path, "-json", "-status-code", "-title", "-tech-detect", - "-o", "-" + "-o", + "-", ] - + output = await self.run_command(cmd) services = [] - + for line in output.splitlines(): try: service = json.loads(line) - services.append({ - "url": service.get("url", ""), - "status_code": service.get("status-code", 0), - "title": service.get("title", ""), - "technologies": service.get("technologies", []) - }) + services.append( + { + "url": service.get("url", ""), + "status_code": service.get("status-code", 0), + "title": service.get("title", ""), + "technologies": service.get("technologies", []), + } + ) except json.JSONDecodeError: continue - + self.console.print(f"[green]Found {len(services)} web services[/green]") return services finally: @@ -144,43 +165,51 @@ async def probe_web_services(self, subdomains: List[str]) -> List[Dict]: os.unlink(temp_file_path) except OSError: pass # File may already be deleted - + async def scan_vulnerabilities(self, web_services: List[Dict]) -> List[Dict]: """Scan web services for vulnerabilities using nuclei""" self.console.print("[bold blue]Scanning for vulnerabilities...[/bold blue]") - + # SECURITY FIX: Use tempfile module instead of hardcoded temp paths - with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as temp_file: - temp_file.write('\n'.join(service["url"] for service in web_services)) + with tempfile.NamedTemporaryFile( + mode="w", suffix=".txt", delete=False + ) as temp_file: + temp_file.write("\n".join(service["url"] for service in web_services)) temp_file_path = temp_file.name - + try: # Run nuclei cmd = [ "nuclei", - "-l", temp_file_path, + "-l", + temp_file_path, "-json", - "-severity", ",".join(self.config["nuclei"]["severity"]), - "-silent" + "-severity", + ",".join(self.config["nuclei"]["severity"]), + "-silent", ] - + output = await self.run_command(cmd) vulnerabilities = [] - + for line in output.splitlines(): try: vuln = json.loads(line) - vulnerabilities.append({ - "url": vuln.get("url", ""), - "type": vuln.get("type", ""), - "severity": vuln.get("severity", ""), - "description": vuln.get("description", ""), - "template": vuln.get("template", "") - }) + vulnerabilities.append( + { + "url": vuln.get("url", ""), + "type": vuln.get("type", ""), + "severity": vuln.get("severity", ""), + "description": vuln.get("description", ""), + "template": vuln.get("template", ""), + } + ) except json.JSONDecodeError: continue - - self.console.print(f"[green]Found {len(vulnerabilities)} potential vulnerabilities[/green]") + + self.console.print( + f"[green]Found {len(vulnerabilities)} potential vulnerabilities[/green]" + ) return vulnerabilities finally: # SECURITY FIX: Ensure temp file is always cleaned up @@ -188,30 +217,36 @@ async def scan_vulnerabilities(self, web_services: List[Dict]) -> List[Dict]: os.unlink(temp_file_path) except OSError: pass # File may already be deleted - - async def run_recon(self, target: str, output_dir: Optional[Path] = None) -> Dict[str, Any]: + + async def run_recon( + self, target: str, output_dir: Optional[Path] = None + ) -> Dict[str, Any]: """Run comprehensive reconnaissance on target""" - self.console.print(f"[bold blue]Starting reconnaissance on {target}[/bold blue]") - + self.console.print( + f"[bold blue]Starting reconnaissance on {target}[/bold blue]" + ) + # SECURITY FIX: Ensure output_dir is a Path object if output_dir is None: output_dir = self.base_dir / "recon_results" / target.replace(".", "_") elif isinstance(output_dir, str): output_dir = Path(output_dir) - + # Create output directory output_dir.mkdir(parents=True, exist_ok=True) - + # Run reconnaissance tasks sequentially for now # SECURITY FIX: Fixed async/await issues and added missing methods try: subdomains = await self.discover_subdomains(target) web_services = await self.probe_web_services(subdomains) vulnerabilities = await self.scan_vulnerabilities(web_services) - + # Generate AI analysis - ai_analysis = await self._analyze_results(target, subdomains, web_services, vulnerabilities) - + ai_analysis = await self._analyze_results( + target, subdomains, web_services, vulnerabilities + ) + # Prepare results results = { "target": target, @@ -219,14 +254,14 @@ async def run_recon(self, target: str, output_dir: Optional[Path] = None) -> Dic "subdomains": subdomains, "web_services": web_services, "vulnerabilities": vulnerabilities, - "ai_analysis": ai_analysis + "ai_analysis": ai_analysis, } - + # Save results await self._save_results(results, output_dir) - + return results - + except Exception as e: self.logger.error("Error during reconnaissance: %s", e) return { @@ -236,14 +271,19 @@ async def run_recon(self, target: str, output_dir: Optional[Path] = None) -> Dic "subdomains": [], "web_services": [], "vulnerabilities": [], - "ai_analysis": {} + "ai_analysis": {}, } - - async def _analyze_results(self, target: str, subdomains: List[str], - web_services: List[Dict], vulnerabilities: List[Dict]) -> Dict: + + async def _analyze_results( + self, + target: str, + subdomains: List[str], + web_services: List[Dict], + vulnerabilities: List[Dict], + ) -> Dict: """Analyze results using AI""" self.console.print("[bold blue]Analyzing results with AI...[/bold blue]") - + # Prepare data for AI analysis analysis_data = { "target": target, @@ -251,16 +291,18 @@ async def _analyze_results(self, target: str, subdomains: List[str], "web_service_count": len(web_services), "vulnerability_count": len(vulnerabilities), "critical_findings": [], - "recommendations": [] + "recommendations": [], } - + # Analyze vulnerabilities if vulnerabilities: - vuln_summary = "\n".join([ - f"- {v['type']} ({v['severity']}): {v['description']}" - for v in vulnerabilities - ]) - + vuln_summary = "\n".join( + [ + f"- {v['type']} ({v['severity']}): {v['description']}" + for v in vulnerabilities + ] + ) + prompt = f""" Analyze these security findings for {target}: @@ -271,7 +313,7 @@ async def _analyze_results(self, target: str, subdomains: List[str], 2. Recommended next steps for exploitation 3. Potential attack chains """ - + ai_response = self.ai_analyzer.ollama.generate(prompt) if ai_response: try: @@ -280,24 +322,25 @@ async def _analyze_results(self, target: str, subdomains: List[str], except: # If JSON parsing fails, use raw response analysis_data["raw_analysis"] = ai_response - + return analysis_data - + async def _save_results(self, results: Dict, output_dir: Path): """Save results in multiple formats""" # Save JSON json_path = output_dir / "results.json" try: - async with aiofiles.open(json_path, 'w') as f: + async with aiofiles.open(json_path, "w") as f: await f.write(json.dumps(results, indent=2)) except (aiofiles.OSError, aiofiles.IOError) as e: self.logger.error("Error writing results: %s", e) - + # Generate Markdown report md_path = output_dir / "report.md" try: - async with aiofiles.open(md_path, 'w') as f: - await f.write(f"""# VulnForge Reconnaissance Report + async with aiofiles.open(md_path, "w") as f: + await f.write( + f"""# VulnForge Reconnaissance Report ## Target: {results['target']} ## Scan Time: {results['timestamp']} @@ -318,8 +361,9 @@ async def _save_results(self, results: Dict, output_dir: Path): ### AI Analysis {json.dumps(results['ai_analysis'], indent=2)} -""") +""" + ) except (aiofiles.OSError, aiofiles.IOError) as e: self.logger.error("Error writing results: %s", e) - - self.console.print(f"[green]Results saved to: {output_dir}[/green]") \ No newline at end of file + + self.console.print(f"[green]Results saved to: {output_dir}[/green]") diff --git a/screen_control/launcher.py b/screen_control/launcher.py index 91fa4df..5bf69d0 100644 --- a/screen_control/launcher.py +++ b/screen_control/launcher.py @@ -8,21 +8,31 @@ # Ensure you have build-essential, g++, rustc, cargo, and nasm installed. # sudo apt-get install build-essential g++ rustc cargo nasm xdotool + def build_modules(): """Compiles all native C++, Rust, and Assembly modules.""" print("--- Building native modules ---") - + # Build C++ print("Building C++ module...") try: - subprocess.run([ - "g++", "-shared", "-fPIC", "-o", "cpp/screen.so", - "cpp/screen.cpp", "-lX11", "-lXext" - ], check=True) + subprocess.run( + [ + "g++", + "-shared", + "-fPIC", + "-o", + "cpp/screen.so", + "cpp/screen.cpp", + "-lX11", + "-lXext", + ], + check=True, + ) print("✓ C++ module built successfully") except (subprocess.CalledProcessError, FileNotFoundError) as e: print(f"✗ C++ build failed: {e}") - + # Build Rust print("Building Rust module...") try: @@ -36,66 +46,76 @@ def build_modules(): print("Building Assembly module...") try: # SECURITY FIX: Use full executable paths and proper argument lists - subprocess.run([ - "/usr/bin/nasm", "-f", "elf64", "assembly/hook.asm", - "-o", "assembly/hook.o" - ], check=True) - subprocess.run([ - "/usr/bin/ld", "-shared", "-o", "assembly/hook.so", "assembly/hook.o" - ], check=True) + subprocess.run( + [ + "/usr/bin/nasm", + "-f", + "elf64", + "assembly/hook.asm", + "-o", + "assembly/hook.o", + ], + check=True, + ) + subprocess.run( + ["/usr/bin/ld", "-shared", "-o", "assembly/hook.so", "assembly/hook.o"], + check=True, + ) print("✓ Assembly module built successfully") except (subprocess.CalledProcessError, FileNotFoundError) as e: print(f"✗ Assembly build failed: {e}") - + print("--- Build complete ---") + def run_demo(): """Runs a demonstration of the screen control system.""" print("--- Screen Control Demo ---") - + # Initialize screen control with fallback support try: screen = ScreenControl(".") print("✓ Screen control initialized") - + # Test basic functionality print("Testing mouse movement...") screen.move_mouse(100, 100) - + print("Testing text input...") screen.type_text("Hello from VulnForge!") - + print("Testing scroll...") screen.scroll(1) - + print("Testing offset calculation...") result = screen.calculate_offset(100, 50) print(f"Offset calculation result: {result}") - + # Test JSON command execution commands = [ {"type": "move_mouse", "x": 200, "y": 200}, {"type": "click", "button": 1}, {"type": "type", "text": "AI-controlled screen interaction"}, {"type": "wait", "seconds": 1}, - {"type": "scroll", "direction": -1} + {"type": "scroll", "direction": -1}, ] - + print("Testing command sequence...") screen.run_sequence(commands) - + print("✓ Demo completed successfully") - + except Exception as e: print(f"✗ Demo failed: {e}") print("Running in fallback mode with basic functionality") + if __name__ == "__main__": # Try to build modules, but continue if it fails try: build_modules() except Exception as e: print(f"Build failed, continuing with fallback: {e}") - + # Run demo - run_demo() \ No newline at end of file + run_demo() diff --git a/screen_control/python/screen_control.py b/screen_control/python/screen_control.py index 1e2f2e7..8bb95c5 100644 --- a/screen_control/python/screen_control.py +++ b/screen_control/python/screen_control.py @@ -4,9 +4,10 @@ import platform import json + class ScreenControl: """A multi-language, modular screen control system.""" - + def __init__(self, base_path): self.base_path = base_path self._load_libraries() @@ -15,7 +16,7 @@ def _load_libraries(self): """Loads the C++, Rust, and Assembly shared libraries.""" try: # C++ Library - cpp_lib_path = os.path.join(self.base_path, 'cpp/screen.so') + cpp_lib_path = os.path.join(self.base_path, "cpp/screen.so") if os.path.exists(cpp_lib_path): self.cpp_lib = ctypes.CDLL(cpp_lib_path) self.cpp_lib.move_mouse_cpp.argtypes = [ctypes.c_int, ctypes.c_int] @@ -26,7 +27,9 @@ def _load_libraries(self): print("C++ library not found, using fallback") # Rust Library - rust_lib_path = os.path.join(self.base_path, 'rust/target/release/librust_control.so') + rust_lib_path = os.path.join( + self.base_path, "rust/target/release/librust_control.so" + ) if os.path.exists(rust_lib_path): self.rust_lib = ctypes.CDLL(rust_lib_path) self.rust_lib.type_text.argtypes = [ctypes.c_char_p] @@ -37,7 +40,7 @@ def _load_libraries(self): print("Rust library not found, using fallback") # Assembly Library - asm_lib_path = os.path.join(self.base_path, 'assembly/hook.so') + asm_lib_path = os.path.join(self.base_path, "assembly/hook.so") if os.path.exists(asm_lib_path): self.asm_lib = ctypes.CDLL(asm_lib_path) self.asm_lib.calculate_offset.argtypes = [ctypes.c_int, ctypes.c_int] @@ -59,7 +62,7 @@ def move_mouse(self, x: int, y: int): self.cpp_lib.move_mouse_cpp(x, y) else: # Fallback using xdotool - subprocess.run(['xdotool', 'mousemove', str(x), str(y)]) + subprocess.run(["xdotool", "mousemove", str(x), str(y)]) def click(self, button: int = 1): """Click mouse button using C++ library or fallback.""" @@ -67,15 +70,15 @@ def click(self, button: int = 1): self.cpp_lib.click_cpp(button) else: # Fallback using xdotool - subprocess.run(['xdotool', 'click', str(button)]) + subprocess.run(["xdotool", "click", str(button)]) def type_text(self, text: str): """Type text using Rust library or fallback.""" if self.rust_available: - self.rust_lib.type_text(text.encode('utf-8')) + self.rust_lib.type_text(text.encode("utf-8")) else: # Fallback using xdotool - subprocess.run(['xdotool', 'type', text]) + subprocess.run(["xdotool", "type", text]) def scroll(self, direction: int): """Scroll using Rust library or fallback.""" @@ -84,9 +87,9 @@ def scroll(self, direction: int): else: # Fallback using xdotool if direction > 0: - subprocess.run(['xdotool', 'key', 'Down']) + subprocess.run(["xdotool", "key", "Down"]) else: - subprocess.run(['xdotool', 'key', 'Up']) + subprocess.run(["xdotool", "key", "Up"]) def calculate_offset(self, base: int, offset: int) -> int: """Calculate offset using Assembly library or fallback.""" @@ -99,18 +102,19 @@ def calculate_offset(self, base: int, offset: int) -> int: def execute_command(self, command: dict): """Execute a command from the JSON schema.""" try: - cmd_type = command.get('type') - if cmd_type == 'move_mouse': - self.move_mouse(command['x'], command['y']) - elif cmd_type == 'click': - self.click(command.get('button', 1)) - elif cmd_type == 'type': - self.type_text(command['text']) - elif cmd_type == 'scroll': - self.scroll(command['direction']) - elif cmd_type == 'wait': + cmd_type = command.get("type") + if cmd_type == "move_mouse": + self.move_mouse(command["x"], command["y"]) + elif cmd_type == "click": + self.click(command.get("button", 1)) + elif cmd_type == "type": + self.type_text(command["text"]) + elif cmd_type == "scroll": + self.scroll(command["direction"]) + elif cmd_type == "wait": import time - time.sleep(command['seconds']) + + time.sleep(command["seconds"]) else: print(f"Unknown command type: {cmd_type}") except Exception as e: @@ -122,4 +126,5 @@ def run_sequence(self, commands: list): self.execute_command(command) # Small delay between commands import time - time.sleep(0.1) \ No newline at end of file + + time.sleep(0.1) diff --git a/scripts/analyze_results.py b/scripts/analyze_results.py index ed269a4..7b7d059 100755 --- a/scripts/analyze_results.py +++ b/scripts/analyze_results.py @@ -15,11 +15,11 @@ # Configure logging logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) + class ResultAnalyzer: def __init__(self, base_dir: Path): self.base_dir = base_dir @@ -27,73 +27,78 @@ def __init__(self, base_dir: Path): self.llm = LLMEngine() self.context = ContextBuilder(base_dir) self.notifier = Notifier(self.config) - + async def analyze_scan_results(self, results_file: Path): """Analyze scan results using AI""" try: # Load scan results with open(results_file) as f: results = json.load(f) - + # Add scan results to context self.context.add_scan_result("recon", results) - + # Build analysis prompt prompt = self.context.build_prompt( "recon", { "target_info": json.dumps(results["target"], indent=2), - "context": self.context._format_context() - } + "context": self.context._format_context(), + }, ) - + # Get AI analysis analysis = self.llm.query( prompt=prompt, - system_prompt="You are a cybersecurity expert analyzing reconnaissance data." + system_prompt="You are a cybersecurity expert analyzing reconnaissance data.", ) - + if not analysis: logger.error("Failed to get AI analysis") return - + # Parse analysis try: analysis_data = json.loads(analysis) - + # Check for critical findings if analysis_data.get("critical_findings"): await self.notifier.notify( "Critical findings detected in scan", "critical", - data=analysis_data["critical_findings"] + data=analysis_data["critical_findings"], ) - + # Save analysis - output_path = self.base_dir / "data" / "analysis" / f"analysis_{Path(results_file).stem}.json" + output_path = ( + self.base_dir + / "data" + / "analysis" + / f"analysis_{Path(results_file).stem}.json" + ) output_path.parent.mkdir(parents=True, exist_ok=True) - - with open(output_path, 'w') as f: + + with open(output_path, "w") as f: json.dump(analysis_data, f, indent=2) - + logger.info(f"Analysis saved to {output_path}") - + # Generate exploit if critical vulnerability found if analysis_data.get("critical_findings"): await self.generate_exploits(analysis_data) - + except json.JSONDecodeError as e: logger.error(f"Failed to parse AI analysis: {e}") - + except Exception as e: logger.error(f"Error analyzing results: {e}") - + async def generate_exploits(self, analysis_data: Dict): """Generate exploits for critical findings""" from modules.exploit_generator.exploit_generator import ExploitGenerator - + generator = ExploitGenerator(self.base_dir, self.llm) - + for finding in analysis_data.get("critical_findings", []): try: # Generate exploit @@ -102,24 +107,26 @@ async def generate_exploits(self, analysis_data: Dict): "cve_id": finding.get("cve_id", "Unknown"), "description": finding.get("description", ""), "affected_software": finding.get("affected_component", ""), - "cvss_score": finding.get("severity", "Unknown") + "cvss_score": finding.get("severity", "Unknown"), }, recon_data={ "ip": finding.get("location", ""), "port": finding.get("port", ""), "service": finding.get("service", ""), - "version": finding.get("version", "") - } + "version": finding.get("version", ""), + }, ) - + if exploit_data and "error" not in exploit_data: # Validate exploit validation = generator.validate_exploit(exploit_data) - + if validation["syntax_valid"] and validation["has_error_handling"]: # Generate Metasploit module - metasploit_module = generator.generate_metasploit_module(exploit_data) - + metasploit_module = generator.generate_metasploit_module( + exploit_data + ) + # Notify about exploit generation await self.notifier.notify( f"Exploit generated for {finding.get('type', 'Unknown')} vulnerability", @@ -127,29 +134,36 @@ async def generate_exploits(self, analysis_data: Dict): data={ "finding": finding, "exploit_path": exploit_data.get("file_path"), - "metasploit_module": str(metasploit_module) if metasploit_module else None, - "validation": validation - } + "metasploit_module": ( + str(metasploit_module) + if metasploit_module + else None + ), + "validation": validation, + }, ) - + except Exception as e: logger.error(f"Error generating exploit for finding: {e}") - + + async def main(): base_dir = Path(__file__).parent.parent analyzer = ResultAnalyzer(base_dir) - + # Find latest scan results results_dir = base_dir / "data" / "scan_results" if not results_dir.exists(): logger.error("No scan results found") return - - latest_result = max(results_dir.glob("recon_*.json"), key=lambda p: p.stat().st_mtime) - + + latest_result = max( + results_dir.glob("recon_*.json"), key=lambda p: p.stat().st_mtime + ) + # Start notifier await analyzer.notifier.start() - + try: # Analyze results await analyzer.analyze_scan_results(latest_result) @@ -157,6 +171,8 @@ async def main(): # Stop notifier await analyzer.notifier.stop() + if __name__ == "__main__": import asyncio - asyncio.run(main()) \ No newline at end of file + + asyncio.run(main()) diff --git a/scripts/recon_scan.py b/scripts/recon_scan.py index 88c86a9..ff678bf 100755 --- a/scripts/recon_scan.py +++ b/scripts/recon_scan.py @@ -16,124 +16,130 @@ # Configure logging logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) + class ReconScanner: def __init__(self, config_path: Path): self.config = self._load_config(config_path) self.results = { "scan_time": datetime.now().isoformat(), "target": self.config["target"], - "findings": [] + "findings": [], } - + def _load_config(self, config_path: Path) -> Dict: """Load scan configuration""" with open(config_path) as f: return json.load(f) - + async def run_nmap_scan(self): """Run Nmap scan""" try: - cmd = ["nmap"] + self.config["tools"]["nmap"]["flags"] + [self.config["target"]["domain"]] + cmd = ( + ["nmap"] + + self.config["tools"]["nmap"]["flags"] + + [self.config["target"]["domain"]] + ) process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() - + if process.returncode == 0: - self.results["findings"].append({ - "tool": "nmap", - "output": stdout.decode(), - "timestamp": datetime.now().isoformat() - }) + self.results["findings"].append( + { + "tool": "nmap", + "output": stdout.decode(), + "timestamp": datetime.now().isoformat(), + } + ) else: logger.error("Nmap scan failed: %s", stderr.decode()) - + except Exception as e: logger.error("Error running Nmap scan: %s", e) - + async def run_subfinder(self): """Run Subfinder for subdomain enumeration""" try: cmd = ["subfinder", "-d", self.config["target"]["domain"], "-silent"] process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() - + if process.returncode == 0: subdomains = stdout.decode().splitlines() - self.results["findings"].append({ - "tool": "subfinder", - "subdomains": subdomains, - "timestamp": datetime.now().isoformat() - }) + self.results["findings"].append( + { + "tool": "subfinder", + "subdomains": subdomains, + "timestamp": datetime.now().isoformat(), + } + ) else: logger.error("Subfinder failed: %s", stderr.decode()) - + except Exception as e: logger.error("Error running Subfinder: %s", e) - + async def run_httpx(self, urls: List[str]): """Run httpx for HTTP probing""" try: # Write URLs to temporary file temp_file = Path("temp_urls.txt") temp_file.write_text("\n".join(urls)) - + cmd = ["httpx", "-l", str(temp_file), "-silent", "-status-code", "-title"] process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() - + if process.returncode == 0: - self.results["findings"].append({ - "tool": "httpx", - "output": stdout.decode(), - "timestamp": datetime.now().isoformat() - }) + self.results["findings"].append( + { + "tool": "httpx", + "output": stdout.decode(), + "timestamp": datetime.now().isoformat(), + } + ) else: logger.error("httpx failed: %s", stderr.decode()) - + # Clean up temp file temp_file.unlink() - + except Exception as e: logger.error("Error running httpx: %s", e) - + async def run_nuclei(self, urls: List[str]): """Run Nuclei for vulnerability scanning""" # SECURITY FIX: Use tempfile module instead of hardcoded temp paths - with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as temp_file: - temp_file.write('\n'.join(urls)) + with tempfile.NamedTemporaryFile( + mode="w", suffix=".txt", delete=False + ) as temp_file: + temp_file.write("\n".join(urls)) temp_file_path = temp_file.name - + try: cmd = ["nuclei", "-l", temp_file_path, "-severity", "critical,high,medium"] process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() - + if process.returncode == 0: - self.results["findings"].append({ - "tool": "nuclei", - "output": stdout.decode(), - "timestamp": datetime.now().isoformat() - }) + self.results["findings"].append( + { + "tool": "nuclei", + "output": stdout.decode(), + "timestamp": datetime.now().isoformat(), + } + ) else: logger.error("Nuclei failed: %s", stderr.decode()) except Exception as e: @@ -144,47 +150,52 @@ async def run_nuclei(self, urls: List[str]): os.unlink(temp_file_path) except OSError: pass # File may already be deleted - + def save_results(self, output_path: Path): """Save scan results to file""" - with open(output_path, 'w') as f: + with open(output_path, "w") as f: json.dump(self.results, f, indent=2) - + + async def main(): # Initialize scanner scanner = ReconScanner(Path("configs/scan_config.json")) - + # Run scans logger.info("Starting reconnaissance scan...") - + # Run Nmap scan logger.info("Running Nmap scan...") await scanner.run_nmap_scan() - + # Run Subfinder logger.info("Running Subfinder...") await scanner.run_subfinder() - + # Get all discovered URLs urls = [scanner.config["target"]["url"]] for finding in scanner.results["findings"]: if finding["tool"] == "subfinder": urls.extend([f"https://{subdomain}" for subdomain in finding["subdomains"]]) - + # Run httpx logger.info("Running httpx...") await scanner.run_httpx(urls) - + # Run Nuclei logger.info("Running Nuclei...") await scanner.run_nuclei(urls) - + # Save results - output_path = Path("data/scan_results") / f"recon_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" + output_path = ( + Path("data/scan_results") + / f"recon_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" + ) output_path.parent.mkdir(parents=True, exist_ok=True) scanner.save_results(output_path) - + logger.info("Scan completed. Results saved to %s", output_path) + if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + asyncio.run(main()) diff --git a/setup.py b/setup.py index 0c8e9f3..a08ee43 100644 --- a/setup.py +++ b/setup.py @@ -60,4 +60,4 @@ "Topic :: Security", ], python_requires=">=3.8", -) \ No newline at end of file +) diff --git a/tests/test_ai_features.py b/tests/test_ai_features.py index 45e6e69..560f3cf 100644 --- a/tests/test_ai_features.py +++ b/tests/test_ai_features.py @@ -8,6 +8,7 @@ from ai_controller import AIController from utils.report_generator import ReportGenerator + class TestAIFeatures: """Test suite for AI features.""" @@ -16,24 +17,25 @@ def setup(self, tmp_path): """Setup test environment.""" self.test_dir = tmp_path / "test_session" self.test_dir.mkdir(parents=True) - + # Create test config self.config_path = self.test_dir / "test_config.json" with open(self.config_path, "w") as f: - json.dump({ - "ai": { - "model": "deepseek-coder", - "temperature": 0.7, - "max_tokens": 1000 + json.dump( + { + "ai": { + "model": "deepseek-coder", + "temperature": 0.7, + "max_tokens": 1000, + }, + "notifications": {"enabled": False}, }, - "notifications": { - "enabled": False - } - }, f) - + f, + ) + # Initialize AI controller self.ai_controller = AIController(str(self.test_dir), str(self.config_path)) - + # Initialize report generator self.report_generator = ReportGenerator(str(self.test_dir)) @@ -43,20 +45,20 @@ async def test_ai_query_processing(self): # Test basic query query = "What should I do if port 8080 is open?" response = await self.ai_controller.process_query(query) - + assert "answer" in response assert "logs" in response assert "prompt" in response assert isinstance(response["answer"], str) - + # Test query with context context = { "open_ports": [8080], "services": {"8080": "http-alt"}, - "previous_findings": ["Potential web server detected"] + "previous_findings": ["Potential web server detected"], } response = await self.ai_controller.process_query(query, context) - + assert "answer" in response assert "logs" in response assert "prompt" in response @@ -70,7 +72,7 @@ def test_report_generation(self): "domain": "example.com", "ip": "93.184.216.34", "scan_time": datetime.now().isoformat(), - "duration": "00:05:23" + "duration": "00:05:23", }, "modules": [ { @@ -78,7 +80,7 @@ def test_report_generation(self): "status": "success", "start_time": "2024-03-20T10:00:00", "end_time": "2024-03-20T10:05:00", - "findings": ["Open port 80", "Open port 443"] + "findings": ["Open port 80", "Open port 443"], } ], "tools": { @@ -86,7 +88,7 @@ def test_report_generation(self): "status": "success", "purpose": "Port scanning", "config": "-sS -sV", - "installed": True + "installed": True, } }, "vulnerabilities": [ @@ -99,9 +101,9 @@ def test_report_generation(self): "references": [ { "title": "CVE-2024-1234", - "url": "https://nvd.nist.gov/vuln/detail/CVE-2024-1234" + "url": "https://nvd.nist.gov/vuln/detail/CVE-2024-1234", } - ] + ], } ], "exploits": [ @@ -109,54 +111,54 @@ def test_report_generation(self): "name": "Test Exploit", "status": "success", "command": "test_command", - "output": "Test output" + "output": "Test output", } ], "defensive_measures": [ { "name": "WAF Detection", "description": "Web Application Firewall detected", - "details": "Cloudflare WAF" + "details": "Cloudflare WAF", } ], "recommendations": [ { "title": "Test Recommendation", "description": "Test recommendation description", - "steps": ["Step 1", "Step 2"] + "steps": ["Step 1", "Step 2"], } ], "errors": [ { "timestamp": datetime.now().isoformat(), "message": "Test error", - "context": "Test error context" + "context": "Test error context", } - ] + ], } - + # Generate reports report_paths = self.report_generator.generate_reports(context) - + # Verify report files exist assert "html" in report_paths assert "markdown" in report_paths assert "json" in report_paths - + # Verify HTML report html_path = Path(report_paths["html"]) assert html_path.exists() html_content = html_path.read_text() assert "VulnForge Report" in html_content assert "example.com" in html_content - + # Verify Markdown report md_path = Path(report_paths["markdown"]) assert md_path.exists() md_content = md_path.read_text() assert "# VulnForge Report" in md_content assert "example.com" in md_content - + # Verify JSON report json_path = Path(report_paths["json"]) assert json_path.exists() @@ -168,22 +170,20 @@ def test_session_data_management(self): # Add test data self.ai_controller.session_data["target_domain"] = "example.com" self.ai_controller.session_data["target_ip"] = "93.184.216.34" - self.ai_controller.session_data["modules"].append({ - "name": "test_module", - "status": "success" - }) - self.ai_controller.session_data["findings"].append({ - "type": "vulnerability", - "severity": "high" - }) - + self.ai_controller.session_data["modules"].append( + {"name": "test_module", "status": "success"} + ) + self.ai_controller.session_data["findings"].append( + {"type": "vulnerability", "severity": "high"} + ) + # Generate report report_paths = self.ai_controller._generate_report() - + # Verify report contains session data json_path = Path(report_paths["json"]) json_content = json.loads(json_path.read_text()) - + assert json_content["target"]["domain"] == "example.com" assert json_content["target"]["ip"] == "93.184.216.34" assert len(json_content["modules"]) > 0 @@ -194,15 +194,16 @@ def test_error_handling(self): # Test invalid query with pytest.raises(Exception): asyncio.run(self.ai_controller.process_query(None)) - + # Test invalid context with pytest.raises(Exception): self.report_generator.generate_reports(None) - + # Test invalid session data self.ai_controller.session_data = None with pytest.raises(Exception): self.ai_controller._generate_report() + if __name__ == "__main__": - pytest.main([__file__, "-v"]) \ No newline at end of file + pytest.main([__file__, "-v"]) diff --git a/tests/test_notifier.py b/tests/test_notifier.py index 23d3c27..503b6a2 100755 --- a/tests/test_notifier.py +++ b/tests/test_notifier.py @@ -10,12 +10,13 @@ import pytest_asyncio from utils.notifier import Notifier + @pytest_asyncio.fixture async def notifier(): """Create a test notifier instance""" test_dir = Path("test_data") test_dir.mkdir(exist_ok=True) - + config = { "notifications": { "enabled": True, @@ -24,32 +25,30 @@ async def notifier(): "smtp_server": "smtp.test.com", "smtp_port": 587, "username": "test@test.com", - "password": "test123" + "password": "test123", }, "discord": { "enabled": True, - "webhook_url": "https://discord.com/api/webhooks/test" + "webhook_url": "https://discord.com/api/webhooks/test", }, - "webhook": { - "enabled": True, - "url": "https://webhook.test.com" - } + "webhook": {"enabled": True, "url": "https://webhook.test.com"}, } } - + config_path = test_dir / "test_config.json" with open(config_path, "w") as f: json.dump(config, f) - + notifier = Notifier(str(test_dir), str(config_path)) yield notifier - + # Cleanup if test_dir.exists(): for file in test_dir.glob("*"): file.unlink() test_dir.rmdir() + @pytest.mark.asyncio async def test_notify_basic(notifier): """Test basic notification""" @@ -57,52 +56,47 @@ async def test_notify_basic(notifier): await notifier.notify("Test message", "info") mock_process.assert_called_once_with("Test message", "info", None, []) + @pytest.mark.asyncio async def test_notify_with_data(notifier): """Test notification with additional data""" test_data = { "vulnerability": "SQL Injection", "severity": "high", - "affected_url": "https://example.com/login" + "affected_url": "https://example.com/login", } - + with patch("utils.notifier.Notifier._process_notification") as mock_process: await notifier.notify("Vulnerability found", "high", data=test_data) - mock_process.assert_called_once_with("Vulnerability found", "high", test_data, []) + mock_process.assert_called_once_with( + "Vulnerability found", "high", test_data, [] + ) + @pytest.mark.asyncio async def test_notify_specific_channels(notifier): """Test notification to specific channels""" with patch("utils.notifier.Notifier._process_notification") as mock_process: - await notifier.notify( - "Test message", - "info", - channels=["email", "discord"] - ) + await notifier.notify("Test message", "info", channels=["email", "discord"]) mock_process.assert_called_once_with( - "Test message", - "info", - None, - ["email", "discord"] + "Test message", "info", None, ["email", "discord"] ) + @pytest.mark.asyncio async def test_email_notification(notifier): """Test email notification sending""" with patch("smtplib.SMTP") as mock_smtp: mock_smtp.return_value.__enter__.return_value = mock_smtp.return_value - - await notifier.notify( - "Test email", - "critical", - channels=["email"] - ) - + + await notifier.notify("Test email", "critical", channels=["email"]) + mock_smtp.assert_called_once_with("smtp.test.com", 587) mock_smtp.return_value.starttls.assert_called_once() mock_smtp.return_value.login.assert_called_once_with("test@test.com", "test123") mock_smtp.return_value.send_message.assert_called_once() + @pytest.mark.asyncio async def test_discord_notification(notifier): """Test Discord notification sending""" @@ -113,20 +107,21 @@ async def test_discord_notification(notifier): mock_cm = AsyncMock() mock_cm.__aenter__.return_value = mock_session mock_cm.__aexit__.return_value = None - + with patch("aiohttp.ClientSession", return_value=mock_cm): await notifier.notify( "Test Discord", "high", data={"finding": "XSS vulnerability"}, - channels=["discord"] + channels=["discord"], ) - + mock_session.post.assert_called_once() call_args = mock_session.post.call_args assert call_args[0][0] == "https://discord.com/api/webhooks/test" assert "embeds" in call_args[1]["json"] + @pytest.mark.asyncio async def test_webhook_notification(notifier): """Test generic webhook notification""" @@ -137,21 +132,18 @@ async def test_webhook_notification(notifier): mock_cm = AsyncMock() mock_cm.__aenter__.return_value = mock_session mock_cm.__aexit__.return_value = None - + with patch("aiohttp.ClientSession", return_value=mock_cm): test_data = { "scan_id": "123", "target": "example.com", - "findings": ["vuln1", "vuln2"] + "findings": ["vuln1", "vuln2"], } - + await notifier.notify( - "Test webhook", - "medium", - data=test_data, - channels=["webhook"] + "Test webhook", "medium", data=test_data, channels=["webhook"] ) - + mock_session.post.assert_called_once() call_args = mock_session.post.call_args assert call_args[0][0] == "https://webhook.test.com" @@ -159,6 +151,7 @@ async def test_webhook_notification(notifier): assert call_args[1]["json"]["severity"] == "medium" assert call_args[1]["json"]["data"] == test_data + @pytest.mark.asyncio async def test_severity_colors(notifier): """Test severity color mapping""" @@ -167,9 +160,9 @@ async def test_severity_colors(notifier): "high": 0xFFA500, "medium": 0xFFFF00, "low": 0x00FF00, - "info": 0x0000FF + "info": 0x0000FF, } - + for severity, expected_color in colors.items(): color = notifier._get_severity_color(severity) - assert color == expected_color, f"Wrong color for {severity}" \ No newline at end of file + assert color == expected_color, f"Wrong color for {severity}" diff --git a/utils/__init__.py b/utils/__init__.py index e39b5e5..fd33727 100755 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -8,16 +8,16 @@ print_results_table, print_error, print_success, - print_warning + print_warning, ) from .logger import setup_logger __all__ = [ - 'print_banner', - 'create_progress', - 'print_results_table', - 'print_error', - 'print_success', - 'print_warning', - 'setup_logger' -] \ No newline at end of file + "print_banner", + "create_progress", + "print_results_table", + "print_error", + "print_success", + "print_warning", + "setup_logger", +] diff --git a/utils/cli_utils.py b/utils/cli_utils.py index 7cbce7b..9446cf7 100755 --- a/utils/cli_utils.py +++ b/utils/cli_utils.py @@ -13,6 +13,7 @@ console = Console() + def print_banner(version: str): """Print VulnForge banner""" banner = f""" @@ -22,69 +23,79 @@ def print_banner(version: str): ╚══════════════════════════════════════════════════════════════╝ """ console.print(Panel(banner, style="bold blue")) - + + def print_results_table(results: dict): """Print results in a formatted table""" for category, items in results.items(): if not items: continue - + table = Table(title=f"{category} Results") table.add_column("Item", style="cyan") table.add_column("Details", style="green") - + for item in items: if isinstance(item, dict): details = ", ".join(f"{k}: {v}" for k, v in item.items()) else: details = str(item) table.add_row(str(item), details) - + console.print(table) - + + def print_json_results(results: Dict[str, Any]): """Print results in JSON format""" console.print(json.dumps(results, indent=2)) - + + def print_error(message: str): """Print error message""" console.print(f"[bold red]Error:[/bold red] {message}") - + + def print_warning(message: str): """Print warning message""" console.print(f"[bold yellow]Warning:[/bold yellow] {message}") - + + def print_success(message: str): """Print success message""" console.print(f"[bold green]Success:[/bold green] {message}") - + + def print_info(message: str): """Print info message""" console.print(f"[bold blue]Info:[/bold blue] {message}") - + + def print_progress(message: str): """Print progress message""" console.print(f"[bold cyan]Progress:[/bold cyan] {message}") - + + def print_debug(message: str): """Print debug message""" console.print(f"[bold magenta]Debug:[/bold magenta] {message}") - + + def print_stealth_stats(stats: Dict[str, Any]): """Print stealth statistics""" table = Table(title="\nStealth Statistics") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") - + for key, value in stats.items(): - table.add_row(key.replace('_', ' ').title(), str(value)) - + table.add_row(key.replace("_", " ").title(), str(value)) + console.print(table) + def create_progress(description: str): """Create a progress bar""" return Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=console - ) \ No newline at end of file + console=console, + ) diff --git a/utils/config_manager.py b/utils/config_manager.py index bf8cadf..b6b9670 100755 --- a/utils/config_manager.py +++ b/utils/config_manager.py @@ -10,6 +10,7 @@ from pathlib import Path from typing import Dict, Any, Optional + class ConfigManager: def __init__(self, config_path: str): self.logger = logging.getLogger(__name__) @@ -17,81 +18,81 @@ def __init__(self, config_path: str): self.config_dir = self.config_path.parent / "configs" self.config_dir.mkdir(parents=True, exist_ok=True) self.config = self._load_config() - + def _load_config(self) -> Dict[str, Any]: """Load configuration from file""" if not self.config_path.exists(): self.logger.warning("Config file not found: %s", self.config_path) self._save_config() try: - with open(self.config_path, 'r') as f: + with open(self.config_path, "r") as f: self.config = json.load(f) except Exception as e: self.logger.error("Error loading config: %s", e) self.config = self._get_default_config() - + def get(self, key: str, default: Any = None) -> Any: """Get configuration value""" return self.config.get(key, default) - + def set(self, key: str, value: Any) -> None: """Set configuration value""" self.config[key] = value self._save_config() - + def _save_config(self) -> None: """Save configuration to file""" try: # SECURITY FIX: Set secure file permissions (0o600) for sensitive config files # This ensures only the owner can read/write the file - with open(self.config_path, 'w') as f: + with open(self.config_path, "w") as f: json.dump(self.config, f, indent=2) os.chmod(self.config_path, 0o600) except Exception as e: self.logger.error("Error saving config: %s", e) - + def get_ai_config(self) -> Dict[str, Any]: """Get AI configuration""" return self.config.get("ai", {}) - + def get_notification_config(self) -> Dict[str, Any]: """Get notification configuration""" return self.config.get("notifications", {}) - + def get_tool_config(self, tool_name: str) -> Dict[str, Any]: """Get tool-specific configuration""" return self.config.get("tools", {}).get(tool_name, {}) - + def save_config(self): """Save current configuration to file""" try: # SECURITY FIX: Set secure file permissions for sensitive config files - with open(self.config_path, 'w') as f: + with open(self.config_path, "w") as f: json.dump(self.config, f, indent=2) os.chmod(self.config_path, 0o600) except Exception as e: self.logger.error("Error saving config: %s", e) - + def update(self, updates: Dict[str, Any]): """Update multiple configuration values""" self.config = {**self.config, **updates} self._save_config() - + def reset(self): """Reset configuration to defaults""" self.config = self._load_config() self._save_config() - + def export(self, filepath: Path): """Export configuration to file""" try: # SECURITY FIX: Set secure file permissions for exported config files - with open(filepath, 'w') as f: + with open(filepath, "w") as f: json.dump(self.config, f, indent=2) os.chmod(filepath, 0o600) except Exception as e: self.logger.error("Error exporting config: %s", e) - + def import_config(self, filepath: Path): """Import configuration from file""" try: @@ -101,23 +102,23 @@ def import_config(self, filepath: Path): self._save_config() except Exception as e: self.logger.error("Error importing config: %s", e) - + def get_all(self) -> Dict[str, Any]: """Get complete configuration""" return self.config.copy() - + def validate(self) -> bool: """Validate configuration""" required_keys = [ "ai.preferred_model", "scanning.max_threads", "stealth.enabled", - "reporting.default_format" + "reporting.default_format", ] - + for key in required_keys: if self.get(key) is None: self.logger.error("Missing required config key: %s", key) return False - - return True \ No newline at end of file + + return True diff --git a/utils/context_builder.py b/utils/context_builder.py index f4dcf4b..94610db 100755 --- a/utils/context_builder.py +++ b/utils/context_builder.py @@ -11,6 +11,7 @@ from datetime import datetime import re + class ContextBuilder: def __init__(self, base_dir: Path): self.base_dir = base_dir @@ -19,61 +20,65 @@ def __init__(self, base_dir: Path): self.max_history = 10 self.prompt_templates_dir = self.base_dir / "prompts" self.prompt_templates_dir.mkdir(parents=True, exist_ok=True) - - def add_tool_output(self, tool_name: str, output: str, metadata: Optional[Dict] = None): + + def add_tool_output( + self, tool_name: str, output: str, metadata: Optional[Dict] = None + ): """Add tool output to context history""" context = { "type": "tool_output", "tool": tool_name, "output": output, "timestamp": datetime.now().isoformat(), - "metadata": metadata or {} + "metadata": metadata or {}, } self._add_to_history(context) - - def add_ai_response(self, prompt: str, response: str, metadata: Optional[Dict] = None): + + def add_ai_response( + self, prompt: str, response: str, metadata: Optional[Dict] = None + ): """Add AI response to context history""" context = { "type": "ai_response", "prompt": prompt, "response": response, "timestamp": datetime.now().isoformat(), - "metadata": metadata or {} + "metadata": metadata or {}, } self._add_to_history(context) - + def add_scan_result(self, result_type: str, data: Dict): """Add scan result to context history""" context = { "type": "scan_result", "result_type": result_type, "data": data, - "timestamp": datetime.now().isoformat() + "timestamp": datetime.now().isoformat(), } self._add_to_history(context) - + def _add_to_history(self, context: Dict): """Add context to history with size limit""" self.context_history.append(context) if len(self.context_history) > self.max_history: self.context_history.pop(0) - + def build_prompt(self, template_name: str, variables: Dict) -> str: """Build prompt from template with context""" template = self._load_template(template_name) if not template: return "" - + # Add context to variables variables["context"] = self._format_context() - + # Replace variables in template prompt = template for key, value in variables.items(): prompt = prompt.replace(f"{{{{ {key} }}}}", str(value)) - + return prompt - + def _load_template(self, template_name: str) -> Optional[str]: """Load prompt template""" template_path = self.prompt_templates_dir / f"{template_name}.jinja" @@ -85,46 +90,48 @@ def _load_template(self, template_name: str) -> Optional[str]: except Exception as e: self.logger.error(f"Error loading template {template_name}: {e}") return None - + def _format_context(self) -> str: """Format context history for prompt""" context_str = [] - + for ctx in self.context_history: if ctx["type"] == "tool_output": context_str.append(f"Tool: {ctx['tool']}\nOutput: {ctx['output']}") elif ctx["type"] == "ai_response": context_str.append(f"Previous AI Response: {ctx['response']}") elif ctx["type"] == "scan_result": - context_str.append(f"Scan Result ({ctx['result_type']}): {json.dumps(ctx['data'])}") - + context_str.append( + f"Scan Result ({ctx['result_type']}): {json.dumps(ctx['data'])}" + ) + return "\n\n".join(context_str) - + def save_template(self, template_name: str, content: str): """Save new prompt template""" template_path = self.prompt_templates_dir / f"{template_name}.jinja" try: - with open(template_path, 'w') as f: + with open(template_path, "w") as f: f.write(content) except Exception as e: self.logger.error(f"Error saving template {template_name}: {e}") - + def get_available_templates(self) -> List[str]: """Get list of available templates""" return [f.stem for f in self.prompt_templates_dir.glob("*.jinja")] - + def clear_history(self): """Clear context history""" self.context_history.clear() - + def export_context(self, filepath: Path): """Export context history to file""" try: - with open(filepath, 'w') as f: + with open(filepath, "w") as f: json.dump(self.context_history, f, indent=2) except Exception as e: self.logger.error(f"Error exporting context: {e}") - + def import_context(self, filepath: Path): """Import context history from file""" try: @@ -132,29 +139,29 @@ def import_context(self, filepath: Path): self.context_history = json.load(f) except Exception as e: self.logger.error(f"Error importing context: {e}") - + def get_relevant_context(self, query: str, max_items: int = 5) -> List[Dict]: """Get most relevant context items for a query""" # Simple relevance scoring based on keyword matching scored_items = [] - keywords = set(re.findall(r'\w+', query.lower())) - + keywords = set(re.findall(r"\w+", query.lower())) + for item in self.context_history: score = 0 content = "" - + if item["type"] == "tool_output": content = f"{item['tool']} {item['output']}" elif item["type"] == "ai_response": content = f"{item['prompt']} {item['response']}" elif item["type"] == "scan_result": - content = json.dumps(item['data']) - - content_words = set(re.findall(r'\w+', content.lower())) + content = json.dumps(item["data"]) + + content_words = set(re.findall(r"\w+", content.lower())) score = len(keywords.intersection(content_words)) - + scored_items.append((score, item)) - + # Sort by score and return top items scored_items.sort(reverse=True) - return [item for _, item in scored_items[:max_items]] \ No newline at end of file + return [item for _, item in scored_items[:max_items]] diff --git a/utils/logger.py b/utils/logger.py index 79e4d4a..2394ec8 100755 --- a/utils/logger.py +++ b/utils/logger.py @@ -15,28 +15,29 @@ import os from datetime import datetime + def setup_logger(name: str, log_dir: str = None) -> logging.Logger: """Set up a logger with the given name.""" if log_dir is None: log_dir = os.path.join(os.path.expanduser("~"), ".vulnforge", "logs") - + os.makedirs(log_dir, exist_ok=True) - + logger = logging.getLogger(name) logger.setLevel(logging.INFO) - + # Create handlers log_file = os.path.join(log_dir, f"{name}.log") file_handler = logging.FileHandler(log_file) console_handler = logging.StreamHandler() - + # Create formatters and add it to handlers - log_format = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') + log_format = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") file_handler.setFormatter(log_format) console_handler.setFormatter(log_format) - + # Add handlers to the logger logger.addHandler(file_handler) logger.addHandler(console_handler) - - return logger \ No newline at end of file + + return logger diff --git a/utils/notifier.py b/utils/notifier.py index aa33944..70d83c9 100755 --- a/utils/notifier.py +++ b/utils/notifier.py @@ -14,16 +14,17 @@ from typing import Dict, List, Optional, Any from datetime import datetime, timezone + class Notifier: """Handles notifications for various channels""" - + def __init__(self, base_dir: str, config_path: str): """Initialize notifier with configuration""" self.base_dir = Path(base_dir) self.config_path = Path(config_path) self.logger = logging.getLogger(__name__) self.config = self._load_config() - + def _load_config(self) -> Dict: """Load notification configuration""" try: @@ -32,46 +33,44 @@ def _load_config(self) -> Dict: except Exception as e: self.logger.error(f"Failed to load config: {e}") return {"notifications": {"enabled": False}} - + async def notify( self, message: str, severity: str = "info", data: Optional[Dict] = None, - channels: Optional[List[str]] = None + channels: Optional[List[str]] = None, ) -> None: """Send notification to specified channels""" if not self.config["notifications"]["enabled"]: return - + if channels is None: channels = [] - + await self._process_notification(message, severity, data, channels) - + async def _process_notification( - self, - message: str, - severity: str, - data: Optional[Dict], - channels: List[str] + self, message: str, severity: str, data: Optional[Dict], channels: List[str] ) -> None: """Process notification for each channel""" tasks = [] - + if "email" in channels and self.config["notifications"]["email"]["enabled"]: tasks.append(self._send_email(message, severity, data)) - + if "discord" in channels and self.config["notifications"]["discord"]["enabled"]: tasks.append(self._send_discord(message, severity, data)) - + if "webhook" in channels and self.config["notifications"]["webhook"]["enabled"]: tasks.append(self._send_webhook(message, severity, data)) - + if tasks: await asyncio.gather(*tasks, return_exceptions=True) - - async def _send_email(self, message: str, severity: str, data: Optional[Dict]) -> None: + + async def _send_email( + self, message: str, severity: str, data: Optional[Dict] + ) -> None: """Send email notification""" try: email_config = self.config["notifications"]["email"] @@ -79,77 +78,84 @@ async def _send_email(self, message: str, severity: str, data: Optional[Dict]) - msg["From"] = email_config["username"] msg["To"] = email_config["username"] msg["Subject"] = f"VulnForge Alert: {severity.upper()}" - + body = f"Message: {message}\nSeverity: {severity}\n" if data: body += f"\nAdditional Data:\n{json.dumps(data, indent=2)}" - + msg.attach(MIMEText(body, "plain")) - - with smtplib.SMTP(email_config["smtp_server"], email_config["smtp_port"]) as server: + + with smtplib.SMTP( + email_config["smtp_server"], email_config["smtp_port"] + ) as server: server.starttls() server.login(email_config["username"], email_config["password"]) server.send_message(msg) - + except Exception as e: self.logger.error(f"Failed to send email: {e}") - - async def _send_discord(self, message: str, severity: str, data: Optional[Dict]) -> None: + + async def _send_discord( + self, message: str, severity: str, data: Optional[Dict] + ) -> None: """Send Discord notification""" try: webhook_url = self.config["notifications"]["discord"]["webhook_url"] color = self._get_severity_color(severity) - + embed = { "title": f"VulnForge Alert: {severity.upper()}", "description": message, "color": color, - "timestamp": datetime.now(timezone.utc).isoformat() # Changed datetime.utcnow() to datetime.now(UTC) + "timestamp": datetime.now( + timezone.utc + ).isoformat(), # Changed datetime.utcnow() to datetime.now(UTC) } - + if data: embed["fields"] = [ {"name": k, "value": str(v), "inline": True} for k, v in data.items() ] - + async with aiohttp.ClientSession() as session: - response = await session.post( - webhook_url, - json={"embeds": [embed]} - ) + response = await session.post(webhook_url, json={"embeds": [embed]}) await response.text() - + except Exception as e: self.logger.error(f"Failed to send Discord notification: {e}") - - async def _send_webhook(self, message: str, severity: str, data: Optional[Dict]) -> None: + + async def _send_webhook( + self, message: str, severity: str, data: Optional[Dict] + ) -> None: """Send generic webhook notification""" try: webhook_url = self.config["notifications"]["webhook"]["url"] payload = { "message": message, "severity": severity, - "timestamp": datetime.now(timezone.utc).isoformat() # Changed datetime.utcnow() to datetime.now(UTC) + "timestamp": datetime.now( + timezone.utc + ).isoformat(), # Changed datetime.utcnow() to datetime.now(UTC) } - + if data: payload["data"] = data - + async with aiohttp.ClientSession() as session: response = await session.post(webhook_url, json=payload) await response.text() - + except Exception as e: self.logger.error(f"Failed to send webhook notification: {e}") - + def _get_severity_color(self, severity: str) -> int: """Get color code for severity level""" colors = { "critical": 0xFF0000, # Red - "high": 0xFFA500, # Orange - "medium": 0xFFFF00, # Yellow - "low": 0x00FF00, # Green - "info": 0x0000FF # Blue + "high": 0xFFA500, # Orange + "medium": 0xFFFF00, # Yellow + "low": 0x00FF00, # Green + "info": 0x0000FF, # Blue } - return colors.get(severity.lower(), 0x808080) # Default to gray \ No newline at end of file + return colors.get(severity.lower(), 0x808080) # Default to gray diff --git a/utils/report_generator.py b/utils/report_generator.py index 7c9982c..3bbbbb7 100644 --- a/utils/report_generator.py +++ b/utils/report_generator.py @@ -12,96 +12,101 @@ from typing import Dict, Any, Optional from jinja2 import Environment, FileSystemLoader + class ReportGenerator: """Handles generation of reports in multiple formats (HTML, Markdown, JSON).""" def __init__(self, session_dir: str): """Initialize the report generator. - + Args: session_dir: Directory to store reports """ self.session_dir = Path(session_dir) self.report_dir = self.session_dir / "reports" self.report_dir.mkdir(parents=True, exist_ok=True) - + # Setup Jinja2 environment template_dir = Path(__file__).parent / "templates" self.env = Environment(loader=FileSystemLoader(str(template_dir))) - + # Setup logging self.logger = logging.getLogger(__name__) self.logger.setLevel(logging.INFO) - + # Add file handler if not already added if not self.logger.handlers: log_file = self.session_dir / "logs" / "report_generator.log" log_file.parent.mkdir(parents=True, exist_ok=True) handler = logging.FileHandler(log_file) - formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) handler.setFormatter(formatter) self.logger.addHandler(handler) - def generate_reports(self, context: Dict[str, Any], output_format: str = "all") -> Dict[str, str]: + def generate_reports( + self, context: Dict[str, Any], output_format: str = "all" + ) -> Dict[str, str]: """Generate reports in specified format(s). - + Args: context: Dictionary containing report data output_format: Desired output format (markdown, json, html, or all) - + Returns: Dictionary mapping report types to their file paths """ self.logger.info(f"Generating reports in {output_format} format...") - + # Add timestamp if not present if "timestamp" not in context: context["timestamp"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S") - + # Add watermark metadata context["watermark"] = { "generated_by": "VulnForge", "author": "DemonKing369.0", - "github": "https://github.com/Arunking9" + "github": "https://github.com/Arunking9", } - + report_paths = {} - + try: if output_format in ["html", "all"]: html_path = self._generate_html_report(context) report_paths["html"] = str(html_path) - + if output_format in ["markdown", "all"]: md_path = self._generate_markdown_report(context) report_paths["markdown"] = str(md_path) - + if output_format in ["json", "all"]: json_path = self._generate_json_report(context) report_paths["json"] = str(json_path) - + self.logger.info("Reports generated successfully") return report_paths - + except Exception as e: self.logger.error(f"Error generating reports: {str(e)}") raise def _generate_html_report(self, context: Dict[str, Any]) -> Path: """Generate HTML report. - + Args: context: Dictionary containing report data - + Returns: Path to the generated HTML report """ self.logger.info("Generating HTML report...") - + try: template = self.env.get_template("report.html") html_content = template.render(**context) - + # Add watermark footer watermark_footer = f""" """ html_content = html_content.replace("", f"{watermark_footer}") - + output_path = self.report_dir / "report.html" with open(output_path, "w", encoding="utf-8") as f: f.write(html_content) - + self.logger.info(f"HTML report generated: {output_path}") return output_path - + except Exception as e: self.logger.error(f"Error generating HTML report: {str(e)}") raise def _generate_markdown_report(self, context: Dict[str, Any]) -> Path: """Generate Markdown report. - + Args: context: Dictionary containing report data - + Returns: Path to the generated Markdown report """ self.logger.info("Generating Markdown report...") - + try: template = self.env.get_template("report.md") md_content = template.render(**context) - + # Add watermark footer watermark_footer = f""" --- @@ -144,67 +149,67 @@ def _generate_markdown_report(self, context: Dict[str, Any]) -> Path: GitHub: https://github.com/Arunking9 """ md_content += watermark_footer - + output_path = self.report_dir / "report.md" with open(output_path, "w", encoding="utf-8") as f: f.write(md_content) - + self.logger.info(f"Markdown report generated: {output_path}") return output_path - + except Exception as e: self.logger.error(f"Error generating Markdown report: {str(e)}") raise def _generate_json_report(self, context: Dict[str, Any]) -> Path: """Generate JSON report. - + Args: context: Dictionary containing report data - + Returns: Path to the generated JSON report """ self.logger.info("Generating JSON report...") - + try: # Add watermark metadata context["watermark"] = { "generated_by": "VulnForge", "author": "DemonKing369.0", "github": "https://github.com/Arunking9", - "timestamp": datetime.now().isoformat() + "timestamp": datetime.now().isoformat(), } - + output_path = self.report_dir / "report.json" with open(output_path, "w", encoding="utf-8") as f: json.dump(context, f, indent=2, default=str) - + self.logger.info(f"JSON report generated: {output_path}") return output_path - + except Exception as e: self.logger.error(f"Error generating JSON report: {str(e)}") raise def get_report_paths(self) -> Dict[str, str]: """Get paths to all generated reports. - + Returns: Dictionary mapping report types to their file paths """ report_paths = {} - + html_path = self.report_dir / "report.html" if html_path.exists(): report_paths["html"] = str(html_path) - + md_path = self.report_dir / "report.md" if md_path.exists(): report_paths["markdown"] = str(md_path) - + json_path = self.report_dir / "report.json" if json_path.exists(): report_paths["json"] = str(json_path) - - return report_paths \ No newline at end of file + + return report_paths diff --git a/utils/session_manager.py b/utils/session_manager.py index 0c1e353..fae7437 100644 --- a/utils/session_manager.py +++ b/utils/session_manager.py @@ -10,12 +10,13 @@ from pathlib import Path from typing import Dict, Any, Optional + class SessionManager: """Manages session data and metadata.""" def __init__(self, base_dir: str): """Initialize session manager. - + Args: base_dir: Base directory for session data """ @@ -25,27 +26,27 @@ def __init__(self, base_dir: str): def create_session(self, target: str) -> str: """Create a new session directory. - + Args: target: Target identifier - + Returns: Path to session directory """ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") self.session_dir = self.base_dir / target / timestamp self.session_dir.mkdir(parents=True, exist_ok=True) - + # Create metadata file self.metadata_file = self.session_dir / "meta.txt" self._write_metadata() - + # Create subdirectories (self.session_dir / "logs").mkdir(exist_ok=True) (self.session_dir / "data").mkdir(exist_ok=True) (self.session_dir / "tools").mkdir(exist_ok=True) (self.session_dir / "exploits").mkdir(exist_ok=True) - + return str(self.session_dir) def _write_metadata(self): @@ -62,38 +63,38 @@ def _write_metadata(self): def save_data(self, data: Dict[str, Any], filename: str): """Save data to session directory. - + Args: data: Data to save filename: Output filename """ if not self.session_dir: raise RuntimeError("No active session") - + output_path = self.session_dir / "data" / filename with open(output_path, "w") as f: json.dump(data, f, indent=2, default=str) def load_data(self, filename: str) -> Dict[str, Any]: """Load data from session directory. - + Args: filename: Input filename - + Returns: Loaded data """ if not self.session_dir: raise RuntimeError("No active session") - + input_path = self.session_dir / "data" / filename with open(input_path, "r") as f: return json.load(f) def get_session_path(self) -> Optional[str]: """Get current session directory path. - + Returns: Path to current session directory or None if no active session """ - return str(self.session_dir) if self.session_dir else None \ No newline at end of file + return str(self.session_dir) if self.session_dir else None diff --git a/utils/stealth_utils.py b/utils/stealth_utils.py index 287a041..38a25b9 100755 --- a/utils/stealth_utils.py +++ b/utils/stealth_utils.py @@ -8,6 +8,7 @@ from typing import List, Dict, Optional import logging + class StealthManager: def __init__(self): self.min_delay = 1.0 @@ -17,12 +18,12 @@ def __init__(self): self.request_count = 0 self.last_request_time = 0 self.logger = logging.getLogger(__name__) - + def set_delay_range(self, min_delay: float, max_delay: float): """Set delay range for requests""" self.min_delay = min_delay self.max_delay = max_delay - + def load_proxies_from_file(self, proxy_file: str): """Load proxies from file""" try: @@ -31,38 +32,38 @@ def load_proxies_from_file(self, proxy_file: str): self.logger.info(f"Loaded {len(self.proxies)} proxies") except Exception as e: self.logger.error(f"Error loading proxies: {e}") - + def get_random_proxy(self) -> Optional[str]: """Get random proxy from list""" if not self.proxies: return None return random.choice(self.proxies) - + def rotate_proxy(self): """Rotate to next proxy""" if not self.proxies: return - + if self.current_proxy in self.proxies: current_index = self.proxies.index(self.current_proxy) next_index = (current_index + 1) % len(self.proxies) else: next_index = 0 - + self.current_proxy = self.proxies[next_index] self.logger.debug(f"Rotated to proxy: {self.current_proxy}") - + def get_request_delay(self) -> float: """Get random delay for request""" return random.uniform(self.min_delay, self.max_delay) - + def wait_before_request(self): """Wait before making request""" delay = self.get_request_delay() time.sleep(delay) self.last_request_time = time.time() self.request_count += 1 - + def get_request_stats(self) -> Dict: """Get request statistics""" return { @@ -70,11 +71,15 @@ def get_request_stats(self) -> Dict: "total_proxies": len(self.proxies), "min_delay": self.min_delay, "max_delay": self.max_delay, - "last_request": time.strftime( - "%Y-%m-%d %H:%M:%S", - time.localtime(self.last_request_time) - ) if self.last_request_time else "Never" + "last_request": ( + time.strftime( + "%Y-%m-%d %H:%M:%S", time.localtime(self.last_request_time) + ) + if self.last_request_time + else "Never" + ), } + # Global stealth manager instance -stealth_manager = StealthManager() \ No newline at end of file +stealth_manager = StealthManager() diff --git a/vulnforge_main.py b/vulnforge_main.py index 785e2c6..fc19c2d 100755 --- a/vulnforge_main.py +++ b/vulnforge_main.py @@ -27,7 +27,7 @@ from recon_module import EnhancedReconModule from ai_integration import AIAnalyzer, OllamaClient -from ai_orchestrator import AIOrchestrator # New Import +from ai_orchestrator import AIOrchestrator # New Import from modules.darkweb import ( run_darkweb_osint, ROBIN_DEFAULT_MODEL, @@ -135,8 +135,12 @@ def install_missing_tools(self): self.logger.info("Installing %s...", tool) try: # SECURITY FIX: Use full executable path and handle subprocess failures - result = subprocess.run(["/usr/bin/go", "install", package], - capture_output=True, text=True, check=True) + result = subprocess.run( + ["/usr/bin/go", "install", package], + capture_output=True, + text=True, + check=True, + ) self.logger.info("Successfully installed %s", tool) except subprocess.CalledProcessError as e: self.logger.error("Failed to install %s: %s", tool, e) @@ -150,7 +154,9 @@ def install_missing_tools(self): async def run_recon(self, target: str, output_dir: Optional[Path] = None): """Run reconnaissance on target""" - recon = EnhancedReconModule(self.base_dir, self.ai_analyzer, config_path="configs/tools.json") + recon = EnhancedReconModule( + self.base_dir, self.ai_analyzer, config_path="configs/tools.json" + ) return await recon.run_recon(target, output_dir) def ask_ai(self, question: str): @@ -389,7 +395,7 @@ async def _async_main(): For detailed documentation, visit: https://github.com/Arunking9/VulnForge """, - formatter_class=argparse.RawDescriptionHelpFormatter + formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--target", "-t", help="Target domain or IP") parser.add_argument( @@ -425,57 +431,83 @@ async def _async_main(): "--stealth", "-s", action="store_true", help="Enable stealth mode" ) parser.add_argument( - "--ai-pipeline", action="store_true", help="Enable the advanced multi-prompt AI pipeline." + "--ai-pipeline", + action="store_true", + help="Enable the advanced multi-prompt AI pipeline.", ) parser.add_argument( - "--prompt-dir", help="Directory for the AI pipeline prompts.", default="AI_Propmt/system-prompts-and-models-of-ai-tools" + "--prompt-dir", + help="Directory for the AI pipeline prompts.", + default="AI_Propmt/system-prompts-and-models-of-ai-tools", ) parser.add_argument( - "--uninstall", action="store_true", help="Uninstall VulnForge and its components" + "--uninstall", + action="store_true", + help="Uninstall VulnForge and its components", ) # Add subparsers for commands - subparsers = parser.add_subparsers(dest='command', help='Available commands') - + subparsers = parser.add_subparsers(dest="command", help="Available commands") + # Ask AI command - ask_ai_parser = subparsers.add_parser('ask-ai', help='Ask the AI assistant a question') - ask_ai_parser.add_argument('question', help='The question to ask the AI') - ask_ai_parser.add_argument('--verbose', action='store_true', help='Show detailed model logs') - ask_ai_parser.add_argument('--dangerous', action='store_true', help='Enable dangerous mode') - ask_ai_parser.add_argument('--confirm-danger', action='store_true', help='Confirm dangerous mode') - + ask_ai_parser = subparsers.add_parser( + "ask-ai", help="Ask the AI assistant a question" + ) + ask_ai_parser.add_argument("question", help="The question to ask the AI") + ask_ai_parser.add_argument( + "--verbose", action="store_true", help="Show detailed model logs" + ) + ask_ai_parser.add_argument( + "--dangerous", action="store_true", help="Enable dangerous mode" + ) + ask_ai_parser.add_argument( + "--confirm-danger", action="store_true", help="Confirm dangerous mode" + ) + # Generate tool command - generate_tool_parser = subparsers.add_parser('generate-tool', help='Generate a custom tool') - generate_tool_parser.add_argument('description', help='Description of the tool to generate') - generate_tool_parser.add_argument('--verbose', action='store_true', help='Show detailed generation logs') - + generate_tool_parser = subparsers.add_parser( + "generate-tool", help="Generate a custom tool" + ) + generate_tool_parser.add_argument( + "description", help="Description of the tool to generate" + ) + generate_tool_parser.add_argument( + "--verbose", action="store_true", help="Show detailed generation logs" + ) + # List tools command - list_tools_parser = subparsers.add_parser('list-tools', help='List all custom tools') - list_tools_parser.add_argument('--verbose', action='store_true', help='Show detailed tool information') + list_tools_parser = subparsers.add_parser( + "list-tools", help="List all custom tools" + ) + list_tools_parser.add_argument( + "--verbose", action="store_true", help="Show detailed tool information" + ) # Dark web OSINT command (Robin integration) darkweb_parser = subparsers.add_parser( - 'darkweb', help='Run the Robin dark web OSINT workflow' + "darkweb", help="Run the Robin dark web OSINT workflow" + ) + darkweb_parser.add_argument( + "--query", "-q", required=True, help="Dark web search query" ) - darkweb_parser.add_argument('--query', '-q', required=True, help='Dark web search query') darkweb_parser.add_argument( - '--model', - '-m', + "--model", + "-m", choices=get_robin_model_choices(), default=ROBIN_DEFAULT_MODEL, - help='LLM model to use for refinement/filtering', + help="LLM model to use for refinement/filtering", ) darkweb_parser.add_argument( - '--threads', - '-t', + "--threads", + "-t", type=int, default=5, - help='Number of concurrent requests for search/scrape', + help="Number of concurrent requests for search/scrape", ) darkweb_parser.add_argument( - '--output', - '-o', - help='Optional output file or directory for the markdown report', + "--output", + "-o", + help="Optional output file or directory for the markdown report", ) args = parser.parse_args() @@ -488,11 +520,11 @@ async def _async_main(): if args.uninstall: script_dir = Path(__file__).parent uninstall_script = script_dir / "uninstall_script.sh" - + if not uninstall_script.exists(): print(f"Error: Uninstall script not found at {uninstall_script}") return - + try: subprocess.run([str(uninstall_script)], check=True) except subprocess.CalledProcessError as e: @@ -504,14 +536,16 @@ async def _async_main(): # Handle AI Pipeline Mode if args.ai_pipeline: if not args.target: - print("Error: A target is required for AI pipeline mode, e.g., --target 'scan example.com'") + print( + "Error: A target is required for AI pipeline mode, e.g., --target 'scan example.com'" + ) return - + prompt_path = Path(args.prompt_dir) if not prompt_path.exists(): print(f"Error: Prompt directory not found at '{prompt_path}'") return - + orchestrator = AIOrchestrator(prompt_path) orchestrator.execute_task(f"Perform a security scan on {args.target}") return @@ -645,5 +679,6 @@ def main(): """Synchronous entrypoint for console_scripts.""" asyncio.run(_async_main()) + if __name__ == "__main__": main() diff --git a/vulnforge_robin_integration/api/app.py b/vulnforge_robin_integration/api/app.py index 69e2ce5..3c6692b 100644 --- a/vulnforge_robin_integration/api/app.py +++ b/vulnforge_robin_integration/api/app.py @@ -134,11 +134,15 @@ async def get_item(item_id: str, db: Session = Depends(get_db)): @app.post("/items/{item_id}/actions") -async def create_action(item_id: str, action: ActionCreate, db: Session = Depends(get_db)): +async def create_action( + item_id: str, action: ActionCreate, db: Session = Depends(get_db) +): leak = db.get(LeakItem, item_id) if not leak: raise HTTPException(status_code=404, detail="Item not found") - record = ActionLog(item_id=item_id, action=action.action, actor=action.actor, notes=action.notes) + record = ActionLog( + item_id=item_id, action=action.action, actor=action.actor, notes=action.notes + ) db.add(record) db.flush() logger.info("action.recorded", extra={"item_id": item_id, "action": action.action}) @@ -146,7 +150,9 @@ async def create_action(item_id: str, action: ActionCreate, db: Session = Depend @app.post("/items/{item_id}/decrypt") -async def decrypt_snippet(item_id: str, body: RawSnippetRequest, db: Session = Depends(get_db)): +async def decrypt_snippet( + item_id: str, body: RawSnippetRequest, db: Session = Depends(get_db) +): if body.reviewer_password != os.getenv("REVIEWER_PASSWORD"): raise HTTPException(status_code=403, detail="Invalid reviewer password") leak = db.get(LeakItem, item_id) @@ -177,4 +183,3 @@ async def healthz(db: Session = Depends(get_db)): @app.get("/metrics") async def metrics(): return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST) - diff --git a/vulnforge_robin_integration/api/dashboard.py b/vulnforge_robin_integration/api/dashboard.py index eac29e8..35d3958 100644 --- a/vulnforge_robin_integration/api/dashboard.py +++ b/vulnforge_robin_integration/api/dashboard.py @@ -100,4 +100,3 @@ def render_dashboard() -> str: return DASHBOARD_HTML - diff --git a/vulnforge_robin_integration/crypto.py b/vulnforge_robin_integration/crypto.py index 87ddb24..b2d1086 100644 --- a/vulnforge_robin_integration/crypto.py +++ b/vulnforge_robin_integration/crypto.py @@ -18,7 +18,9 @@ def _load_key() -> bytes: return key -def encrypt_raw_snippet(plaintext: bytes, aad: bytes | None = None) -> Tuple[bytes, bytes, bytes]: +def encrypt_raw_snippet( + plaintext: bytes, aad: bytes | None = None +) -> Tuple[bytes, bytes, bytes]: key = _load_key() aes = AESGCM(key) nonce = os.urandom(12) @@ -28,8 +30,9 @@ def encrypt_raw_snippet(plaintext: bytes, aad: bytes | None = None) -> Tuple[byt return body, nonce, tag -def decrypt_raw_snippet(ciphertext: bytes, nonce: bytes, tag: bytes, aad: bytes | None = None) -> bytes: +def decrypt_raw_snippet( + ciphertext: bytes, nonce: bytes, tag: bytes, aad: bytes | None = None +) -> bytes: key = _load_key() aes = AESGCM(key) return aes.decrypt(nonce, ciphertext + tag, aad) - diff --git a/vulnforge_robin_integration/db.py b/vulnforge_robin_integration/db.py index e865a40..d21a45a 100644 --- a/vulnforge_robin_integration/db.py +++ b/vulnforge_robin_integration/db.py @@ -6,7 +6,20 @@ from datetime import datetime from typing import Dict, Generator, Optional -from sqlalchemy import JSON, Column, DateTime, Float, ForeignKey, Integer, LargeBinary, String, Text, create_engine, func, select +from sqlalchemy import ( + JSON, + Column, + DateTime, + Float, + ForeignKey, + Integer, + LargeBinary, + String, + Text, + create_engine, + func, + select, +) from sqlalchemy.dialects.postgresql import UUID as PGUUID from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import Session, relationship, scoped_session, sessionmaker @@ -18,14 +31,20 @@ DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./data/vulnforge_robin.db") engine = create_engine(DATABASE_URL, future=True) -SessionLocal = scoped_session(sessionmaker(bind=engine, autoflush=False, autocommit=False)) +SessionLocal = scoped_session( + sessionmaker(bind=engine, autoflush=False, autocommit=False) +) Base = declarative_base() class LeakItem(Base): __tablename__ = "leak_items" - id = Column(PGUUID(as_uuid=True) if "postgres" in DATABASE_URL else String, primary_key=True, default=uuid4) + id = Column( + PGUUID(as_uuid=True) if "postgres" in DATABASE_URL else String, + primary_key=True, + default=uuid4, + ) target_type = Column(String(64), nullable=False) target_value = Column(String(512), nullable=False, index=True) leak_type = Column(String(128), nullable=False) @@ -42,10 +61,14 @@ class LeakItem(Base): score = Column(Integer, default=0) notes = Column(Text) created_at = Column(DateTime(timezone=True), server_default=func.now()) - updated_at = Column(DateTime(timezone=True), onupdate=func.now(), server_default=func.now()) + updated_at = Column( + DateTime(timezone=True), onupdate=func.now(), server_default=func.now() + ) hash_key = Column(String(128), unique=True, index=True) - actions = relationship("ActionLog", back_populates="item", cascade="all,delete-orphan") + actions = relationship( + "ActionLog", back_populates="item", cascade="all,delete-orphan" + ) class ActionLog(Base): @@ -62,7 +85,11 @@ class ActionLog(Base): notes = Column(Text) timestamp = Column(DateTime(timezone=True), default=datetime.utcnow) - item = relationship("LeakItem", back_populates="actions", primaryjoin="ActionLog.item_id==LeakItem.id") + item = relationship( + "LeakItem", + back_populates="actions", + primaryjoin="ActionLog.item_id==LeakItem.id", + ) def init_db() -> None: @@ -81,4 +108,3 @@ def get_session() -> Generator[Session, None, None]: raise finally: session.close() - diff --git a/vulnforge_robin_integration/enricher.py b/vulnforge_robin_integration/enricher.py index 09cef61..9ecc4d3 100644 --- a/vulnforge_robin_integration/enricher.py +++ b/vulnforge_robin_integration/enricher.py @@ -35,7 +35,9 @@ async def call_with_retries(func, *args, **kwargs): delay *= 2 -async def hibp_lookup(client: httpx.AsyncClient, email: str, settings: EnrichmentSettings): +async def hibp_lookup( + client: httpx.AsyncClient, email: str, settings: EnrichmentSettings +): if not settings.hibp_api_key or "@" not in email: return {"enabled": False} headers = {"hibp-api-key": settings.hibp_api_key} @@ -50,7 +52,9 @@ async def hibp_lookup(client: httpx.AsyncClient, email: str, settings: Enrichmen return {"found": True, "breaches": resp.json()} -async def shodan_lookup(client: httpx.AsyncClient, target: str, settings: EnrichmentSettings): +async def shodan_lookup( + client: httpx.AsyncClient, target: str, settings: EnrichmentSettings +): if not settings.shodan_api_key: return {"enabled": False} resp = await client.get( @@ -63,7 +67,9 @@ async def shodan_lookup(client: httpx.AsyncClient, target: str, settings: Enrich return {"matches": data.get("matches", [])[:5]} -async def censys_lookup(client: httpx.AsyncClient, target: str, settings: EnrichmentSettings): +async def censys_lookup( + client: httpx.AsyncClient, target: str, settings: EnrichmentSettings +): if not settings.censys_id or not settings.censys_secret: return {"enabled": False} resp = await client.get( @@ -76,7 +82,9 @@ async def censys_lookup(client: httpx.AsyncClient, target: str, settings: Enrich return resp.json() -async def passive_dns_lookup(client: httpx.AsyncClient, target: str, settings: EnrichmentSettings): +async def passive_dns_lookup( + client: httpx.AsyncClient, target: str, settings: EnrichmentSettings +): if not settings.passive_dns_endpoint: return {"enabled": False} headers = {} @@ -92,7 +100,9 @@ async def passive_dns_lookup(client: httpx.AsyncClient, target: str, settings: E return resp.json() -async def enrich(structured_fields: Dict[str, Any], target_value: str) -> Dict[str, Any]: +async def enrich( + structured_fields: Dict[str, Any], target_value: str +) -> Dict[str, Any]: settings = EnrichmentSettings() async with httpx.AsyncClient() as client: tasks = [] @@ -100,7 +110,9 @@ async def enrich(structured_fields: Dict[str, Any], target_value: str) -> Dict[s tasks.append(call_with_retries(hibp_lookup, client, email or "", settings)) tasks.append(call_with_retries(shodan_lookup, client, target_value, settings)) tasks.append(call_with_retries(censys_lookup, client, target_value, settings)) - tasks.append(call_with_retries(passive_dns_lookup, client, target_value, settings)) + tasks.append( + call_with_retries(passive_dns_lookup, client, target_value, settings) + ) results = await asyncio.gather(*tasks) return { @@ -109,4 +121,3 @@ async def enrich(structured_fields: Dict[str, Any], target_value: str) -> Dict[s "censys": results[2], "passive_dns": results[3], } - diff --git a/vulnforge_robin_integration/ingest.py b/vulnforge_robin_integration/ingest.py index 1cff8f7..4a72dde 100644 --- a/vulnforge_robin_integration/ingest.py +++ b/vulnforge_robin_integration/ingest.py @@ -17,11 +17,17 @@ logger = logging.getLogger(__name__) -INGEST_COUNTER = Counter("robin_ingest_total", "Total items accepted for processing", ["source"]) -BACKLOG_GAUGE = Gauge("robin_ingest_backlog", "Approximate queued items awaiting worker") +INGEST_COUNTER = Counter( + "robin_ingest_total", "Total items accepted for processing", ["source"] +) +BACKLOG_GAUGE = Gauge( + "robin_ingest_backlog", "Approximate queued items awaiting worker" +) -def ingest_payload(payload: str | bytes | dict | list, source: str = "webhook") -> List[str]: +def ingest_payload( + payload: str | bytes | dict | list, source: str = "webhook" +) -> List[str]: """Normalize and persist payload, returning IDs for worker processing.""" if isinstance(payload, (bytes, bytearray)): payload = payload.decode("utf-8", errors="ignore") @@ -49,7 +55,9 @@ def ingest_payload(payload: str | bytes | dict | list, source: str = "webhook") ).hexdigest() with get_session() as session: - existing = session.query(LeakItem).filter_by(hash_key=hash_key).one_or_none() + existing = ( + session.query(LeakItem).filter_by(hash_key=hash_key).one_or_none() + ) if existing: existing.raw_ciphertext = ciphertext existing.raw_nonce = nonce @@ -89,7 +97,10 @@ def ingest_payload(payload: str | bytes | dict | list, source: str = "webhook") process_item.delay(leak_id) stored_ids.append(leak_id) - logger.info("ingest.enqueue", extra={"module": "ingest", "item_id": leak_id, "source": source}) + logger.info( + "ingest.enqueue", + extra={"module": "ingest", "item_id": leak_id, "source": source}, + ) BACKLOG_GAUGE.set(len(stored_ids)) return stored_ids @@ -111,6 +122,8 @@ def run_forever(self) -> None: archive = file.with_suffix(file.suffix + ".processed") file.rename(archive) except Exception as exc: - logger.exception("watcher.error", extra={"module": "ingest", "path": str(file)}) + logger.exception( + "watcher.error", + extra={"module": "ingest", "path": str(file)}, + ) time.sleep(self.interval) - diff --git a/vulnforge_robin_integration/models.py b/vulnforge_robin_integration/models.py index 34d473f..4d0b793 100644 --- a/vulnforge_robin_integration/models.py +++ b/vulnforge_robin_integration/models.py @@ -71,4 +71,3 @@ class HealthResponse(BaseModel): status: str broker_ok: bool db_ok: bool - diff --git a/vulnforge_robin_integration/normalizer.py b/vulnforge_robin_integration/normalizer.py index 7b156c3..8ef20a6 100644 --- a/vulnforge_robin_integration/normalizer.py +++ b/vulnforge_robin_integration/normalizer.py @@ -19,9 +19,13 @@ def _parse_markdown(content: str) -> Dict[str, Any]: target_value = target_match.group("value").strip() else: target_value = "unknown" - leak_type_match = re.search(r"Leak Type\s*[:\-]\s*(?P[^\n]+)", content, re.IGNORECASE) + leak_type_match = re.search( + r"Leak Type\s*[:\-]\s*(?P[^\n]+)", content, re.IGNORECASE + ) leak_type = leak_type_match.group("lt").strip() if leak_type_match else "unknown" - source_match = re.search(r"Source\s*[:\-]\s*(?P[^\n]+)", content, re.IGNORECASE) + source_match = re.search( + r"Source\s*[:\-]\s*(?P[^\n]+)", content, re.IGNORECASE + ) source = source_match.group("src").strip() if source_match else "robin" first_seen = _extract_date(content, "First Seen") @@ -57,7 +61,7 @@ def _extract_sections(content: str) -> Dict[str, str]: def _extract_date(content: str, label: str) -> Optional[datetime]: - pattern = re.compile(fr"{label}\s*[:\-]\s*(?P[^\n]+)", re.IGNORECASE) + pattern = re.compile(rf"{label}\s*[:\-]\s*(?P[^\n]+)", re.IGNORECASE) match = pattern.search(content) if not match: return None @@ -122,4 +126,3 @@ def _parse_dt(value: Any) -> Optional[datetime]: "raw": (obj.get("raw_snippet") or obj.get("raw") or ""), "notes": obj.get("notes"), } - diff --git a/vulnforge_robin_integration/scoring.py b/vulnforge_robin_integration/scoring.py index a847fc5..2675550 100644 --- a/vulnforge_robin_integration/scoring.py +++ b/vulnforge_robin_integration/scoring.py @@ -36,9 +36,14 @@ def timeliness_factor(first_seen, last_seen) -> float: return 0.5 -def compute_score(confidence: float, structured_fields: Dict[str, any], first_seen, last_seen) -> int: +def compute_score( + confidence: float, structured_fields: Dict[str, any], first_seen, last_seen +) -> int: impact = impact_factor(structured_fields) timeliness = timeliness_factor(first_seen, last_seen) - score = round((confidence * CONFIDENCE_WEIGHT) + (impact * IMPACT_WEIGHT) + (timeliness * TIMELINESS_WEIGHT)) + score = round( + (confidence * CONFIDENCE_WEIGHT) + + (impact * IMPACT_WEIGHT) + + (timeliness * TIMELINESS_WEIGHT) + ) return min(score, 100) - diff --git a/vulnforge_robin_integration/tests/conftest.py b/vulnforge_robin_integration/tests/conftest.py index a3a7931..ce0671b 100644 --- a/vulnforge_robin_integration/tests/conftest.py +++ b/vulnforge_robin_integration/tests/conftest.py @@ -19,4 +19,3 @@ def clean_db(): Base.metadata.drop_all(bind=engine) Base.metadata.create_all(bind=engine) yield - diff --git a/vulnforge_robin_integration/tests/test_api.py b/vulnforge_robin_integration/tests/test_api.py index 295bebe..a3ef239 100644 --- a/vulnforge_robin_integration/tests/test_api.py +++ b/vulnforge_robin_integration/tests/test_api.py @@ -11,7 +11,10 @@ def test_ingest_list_and_decrypt(): "target": {"type": "domain", "value": "example.com"}, "leak_type": "credentials", "source": "unit-test", - "structured_fields": {"email": "demo@example.com", "password_present": True}, + "structured_fields": { + "email": "demo@example.com", + "password_present": True, + }, "raw": "user:demo@example.com pass:Secret", }, } @@ -22,7 +25,8 @@ def test_ingest_list_and_decrypt(): assert data["total"] == 1 item_id = data["items"][0]["id"] - decrypt_resp = client.post(f"/items/{item_id}/decrypt", json={"reviewer_password": "testpass"}) + decrypt_resp = client.post( + f"/items/{item_id}/decrypt", json={"reviewer_password": "testpass"} + ) assert decrypt_resp.status_code == 200 assert "user:demo@example.com" in decrypt_resp.json()["raw_snippet"] - diff --git a/vulnforge_robin_integration/tests/test_crypto.py b/vulnforge_robin_integration/tests/test_crypto.py index 3090400..4c363c7 100644 --- a/vulnforge_robin_integration/tests/test_crypto.py +++ b/vulnforge_robin_integration/tests/test_crypto.py @@ -6,4 +6,3 @@ def test_encrypt_decrypt_roundtrip(): ciphertext, nonce, tag = encrypt_raw_snippet(plaintext) recovered = decrypt_raw_snippet(ciphertext, nonce, tag) assert recovered == plaintext - diff --git a/vulnforge_robin_integration/tests/test_normalizer.py b/vulnforge_robin_integration/tests/test_normalizer.py index 009e949..933881d 100644 --- a/vulnforge_robin_integration/tests/test_normalizer.py +++ b/vulnforge_robin_integration/tests/test_normalizer.py @@ -18,4 +18,3 @@ def test_normalize_markdown(): assert item["leak_type"] == "credentials" assert item["source"] == "forum" assert item["first_seen"].isoformat().startswith("2024-08-01") - diff --git a/vulnforge_robin_integration/tests/test_scoring.py b/vulnforge_robin_integration/tests/test_scoring.py index 43a5ec1..eaadf99 100644 --- a/vulnforge_robin_integration/tests/test_scoring.py +++ b/vulnforge_robin_integration/tests/test_scoring.py @@ -13,4 +13,3 @@ def test_compute_score_high_impact(): last_seen=last, ) assert score >= 80 - diff --git a/vulnforge_robin_integration/workers.py b/vulnforge_robin_integration/workers.py index 49feb6f..e3440ea 100644 --- a/vulnforge_robin_integration/workers.py +++ b/vulnforge_robin_integration/workers.py @@ -16,11 +16,18 @@ BROKER_URL = os.getenv("BROKER_URL", "redis://localhost:6379/0") RESULT_BACKEND = os.getenv("RESULT_BACKEND", BROKER_URL) -app = Celery("vulnforge_robin", broker=BROKER_URL, backend=RESULT_BACKEND, include=["vulnforge_robin_integration.workers"]) +app = Celery( + "vulnforge_robin", + broker=BROKER_URL, + backend=RESULT_BACKEND, + include=["vulnforge_robin_integration.workers"], +) if os.getenv("CELERY_EAGER") == "1": app.conf.task_always_eager = True -PROCESSED_COUNTER = Counter("robin_items_processed_total", "Total processed items via worker") +PROCESSED_COUNTER = Counter( + "robin_items_processed_total", "Total processed items via worker" +) @app.task(name="vulnforge_robin.process_item") @@ -35,12 +42,18 @@ def process_item(leak_id: str) -> None: try: enrichment = asyncio.run(enrich(structured, leak.target_value)) except Exception as exc: - logger.exception("worker.enrich.error", extra={"item_id": leak_id, "error": str(exc)}) + logger.exception( + "worker.enrich.error", extra={"item_id": leak_id, "error": str(exc)} + ) enrichment = {"error": str(exc)} leak.enrichment = enrichment - leak.score = compute_score(leak.confidence or 0.5, structured, leak.first_seen, leak.last_seen) + leak.score = compute_score( + leak.confidence or 0.5, structured, leak.first_seen, leak.last_seen + ) session.add(leak) - logger.info("worker.process_item.completed", extra={"item_id": leak_id, "score": leak.score}) + logger.info( + "worker.process_item.completed", + extra={"item_id": leak_id, "score": leak.score}, + ) PROCESSED_COUNTER.inc() -