Refactor AI analysis functionality to support multiple providers and prompt types.

atiilla · atiilla · commit e7273637c6c2 · 2025-08-27T23:05:33.000+02:00
diff --git a/config.yaml b/config.yaml
@@ -21,7 +21,7 @@ ai_providers:
     api_key_env: ''
     enabled: false
     max_tokens: 512
-    model: microsoft/DialoGPT-medium
+    model: deepseek-r1:1.5b
     name: local
     priority: 4
     rate_limit: 0.1
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sqlmap-ai"
-version = "2.0.3"
+version = "2.0.4"
 description = "AI-powered SQL injection testing tool with multiple AI providers"
 readme = "README.md"
 license = "MIT"
diff --git a/sqlmap_ai/ai_analyzer.py b/sqlmap_ai/ai_analyzer.py
@@ -1,37 +1,81 @@
-from utils.groq_utils import get_groq_response
+from utils.ai_providers import ai_manager, AIProvider
 from sqlmap_ai.ui import print_info, print_warning, print_success
 from sqlmap_ai.parser import extract_sqlmap_info
 import json
-def ai_suggest_next_steps(report, scan_history=None, extracted_data=None):
+import asyncio
+def ai_suggest_next_steps(report, scan_history=None, extracted_data=None, ai_provider=None, use_advanced=None):
     print_info("Analyzing SQLMap results with AI...")
     if not report:
         return ["--technique=BT", "--level=2", "--risk=1"]
     if report.startswith("TIMEOUT_WITH_PARTIAL_DATA:"):
         report = report[len("TIMEOUT_WITH_PARTIAL_DATA:"):]
     structured_info = extract_sqlmap_info(report)
-    prompt = create_advanced_prompt(report, structured_info, scan_history, extracted_data)
-    print_info("Sending detailed analysis request to Groq AI...")
-    response = get_groq_response(prompt=prompt)
-    if not response:
+    
+    # Determine which prompt to use based on provider and user preference
+    use_simple = False
+    
+    # Default behavior: simple for Ollama, advanced for others
+    if ai_provider == AIProvider.OLLAMA or ai_provider == "ollama":
+        use_simple = True
+    
+    # Override based on user preference
+    if use_advanced is not None:
+        use_simple = not use_advanced
+    
+    if use_simple:
+        prompt = create_simple_prompt(report, structured_info, scan_history, extracted_data)
+        print_info("Using simple prompt for AI analysis")
+    else:
+        prompt = create_advanced_prompt(report, structured_info, scan_history, extracted_data)
+        print_info("Using advanced prompt for AI analysis")
+    
+    # Determine which AI provider to use
+    provider_name = "AI"
+    if ai_provider:
+        provider_name = ai_provider.upper()
+    print_info(f"Sending detailed analysis request to {provider_name}...")
+    
+    # Use the AI provider system
+    try:
+        # Convert string provider to AIProvider enum if needed
+        provider_enum = None
+        if ai_provider:
+            try:
+                provider_enum = AIProvider(ai_provider)
+                print_info(f"Using AI provider: {provider_enum}")
+            except ValueError:
+                print_warning(f"Invalid AI provider: {ai_provider}")
+                return ["--technique=BEU", "--level=3"]
+        
+        response = asyncio.run(ai_manager.get_response(prompt, provider=provider_enum))
+        if response and response.success:
+            response_text = response.content
+        else:
+            print_warning(f"AI provider {ai_provider} failed: {response.error if response else 'Unknown error'}")
+            return ["--technique=BEU", "--level=3"]
+    except Exception as e:
+        print_warning(f"AI analysis failed: {e}")
+        return ["--technique=BEU", "--level=3"]
+    if not response_text:
         print_warning("AI couldn't suggest options, using fallback options")
         return ["--technique=BEU", "--level=3"]
     print_success("Received AI recommendations!")
     try:
         # Try parsing JSON responses
-        if "```json" in response:
-            json_start = response.find("```json") + 7
-            json_end = response.find("```", json_start)
-            json_str = response[json_start:json_end].strip()
+        if "```json" in response_text:
+            json_start = response_text.find("```json") + 7
+            json_end = response_text.find("```", json_start)
+            json_str = response_text[json_start:json_end].strip()
             recommendation = json.loads(json_str)
             if "sqlmap_options" in recommendation:
                 return recommendation["sqlmap_options"]
             elif "options" in recommendation:
                 return recommendation["options"]
         # Look for code blocks without json tag
-        elif "```" in response:
-            code_start = response.find("```") + 3
-            code_end = response.find("```", code_start)
-            code_block = response[code_start:code_end].strip()
+        elif "```" in response_text:
+            code_start = response_text.find("```") + 3
+            code_end = response_text.find("```", code_start)
+            code_block = response_text[code_start:code_end].strip()
             # Check if content is JSON
             try:
                 recommendation = json.loads(code_block)
@@ -44,7 +88,7 @@ def ai_suggest_next_steps(report, scan_history=None, extracted_data=None):
                 
         # Extract options from the response text
         options = []
-        for line in response.split('\n'):
+        for line in response_text.split('\n'):
             line = line.strip()
             if line.startswith('--') or line.startswith('-p ') or line.startswith('-D ') or line.startswith('-T ') or \
                line.startswith('--data=') or line.startswith('--cookie=') or line.startswith('--headers=') or \
@@ -58,7 +102,7 @@ def ai_suggest_next_steps(report, scan_history=None, extracted_data=None):
         print_warning(f"Error parsing AI response: {str(e)}")
         # Fallback to simple extraction
         options = []
-        for line in response.strip().split('\n'):
+        for line in response_text.strip().split('\n'):
             for part in line.split():
                 if part.startswith('--') or part.startswith('-p ') or part.startswith('-D ') or part.startswith('-T ') or \
                    part.startswith('--data=') or part.startswith('--cookie=') or part.startswith('--headers=') or \
@@ -126,6 +170,35 @@ def ai_suggest_next_steps(report, scan_history=None, extracted_data=None):
             return ["--technique=BEU", "--level=3"]
         
     return valid_options
+
+def create_simple_prompt(report, structured_info, scan_history=None, extracted_data=None):
+    """Create a simpler prompt for Ollama to avoid timeouts"""
+    prompt = """
+    You are a SQLMap expert. Analyze this SQL injection scan result and suggest the next steps.
+
+    DBMS: {dbms}
+    Vulnerable Parameters: {vulnerable_params}
+    Databases Found: {databases}
+
+    Based on this information, suggest the next SQLMap options to use. Focus on:
+    1. Extracting more database information
+    2. Using conservative settings to prevent timeouts
+    3. Using specific techniques rather than broad scanning
+
+    Return your recommendation as a simple list of options, one per line:
+    --level=2
+    --risk=1
+    --dbs
+    """
+    
+    formatted_prompt = prompt.format(
+        dbms=structured_info.get("dbms", "Unknown"),
+        vulnerable_params=', '.join(structured_info.get("vulnerable_parameters", [])) or "None",
+        databases=', '.join(structured_info.get("databases", [])) or "None"
+    )
+    
+    return formatted_prompt
+
 def create_advanced_prompt(report, structured_info, scan_history=None, extracted_data=None):
     prompt = """
     You are a SQLMap expert. You are given a SQLMap scan report and a list of previous scan steps.
diff --git a/sqlmap_ai/enhanced_cli.py b/sqlmap_ai/enhanced_cli.py
@@ -95,14 +95,28 @@ def create_parser(self) -> argparse.ArgumentParser:
         ai_group = parser.add_argument_group('AI Configuration')
         ai_group.add_argument(
             '--ai-provider',
-            choices=['groq', 'openai', 'anthropic', 'local', 'auto'],
+            choices=['groq', 'openai', 'anthropic', 'local', 'ollama', 'auto'],
             default='auto',
             help='AI provider to use (default: auto)'
         )
         ai_group.add_argument(
             '--ai-model',
             help='Specific AI model to use'
         )
+        ai_group.add_argument(
+            '--ollama-model',
+            help='Specific Ollama model to use (overrides OLLAMA_MODEL env var)'
+        )
+        ai_group.add_argument(
+            '--advanced',
+            action='store_true',
+            help='Use advanced AI prompts (default: simple for Ollama, advanced for other providers)'
+        )
+        ai_group.add_argument(
+            '--simple',
+            action='store_true',
+            help='Use simple AI prompts (default: simple for Ollama, advanced for other providers)'
+        )
         ai_group.add_argument(
             '--disable-ai',
             action='store_true',
diff --git a/sqlmap_ai/main.py b/sqlmap_ai/main.py
@@ -23,7 +23,7 @@
 from sqlmap_ai.adaptive_testing import run_adaptive_test_sequence
 from sqlmap_ai.advanced_reporting import report_generator
 from sqlmap_ai.evasion_engine import evasion_engine
-from utils.ai_providers import ai_manager, get_available_ai_providers
+from utils.ai_providers import ai_manager, get_available_ai_providers, AIProvider
 from typing import Optional
 def main():
     """Enhanced main function with improved CLI and security"""
@@ -160,7 +160,7 @@ def run_enhanced_adaptive_mode(runner, target_url, user_timeout, interactive_mod
     # Enhance with AI analysis if available
     if result and not args.disable_ai:
         try:
-            asyncio.run(enhance_with_ai_analysis(result, target_url))
+            asyncio.run(enhance_with_ai_analysis(result, target_url, args))
         except Exception as e:
             print_warning(f"AI analysis failed: {e}")
     
@@ -171,17 +171,62 @@ def run_enhanced_standard_mode(runner, target_url, user_timeout, interactive_mod
     """Run enhanced standard mode"""
     print_info("Starting enhanced standard testing...")
     
+    # Determine AI provider from arguments
+    ai_provider = None
+    if hasattr(args, 'ai_provider') and args.ai_provider and args.ai_provider != 'auto':
+        ai_provider = args.ai_provider
+        
+        # Enable Ollama if selected
+        if ai_provider == 'ollama':
+            import os
+            os.environ['ENABLE_OLLAMA'] = 'true'
+            # Reinitialize AI manager to include Ollama provider
+            ai_manager.reinitialize_providers()
+    
+    # Set Ollama model if specified
+    if hasattr(args, 'ollama_model') and args.ollama_model:
+        import os
+        os.environ['OLLAMA_MODEL'] = args.ollama_model
+    
+    # Set Ollama model if specified
+    if hasattr(args, 'ollama_model') and args.ollama_model:
+        import os
+        os.environ['OLLAMA_MODEL'] = args.ollama_model
+        # Update the Ollama provider's model if it exists
+        if AIProvider.OLLAMA in ai_manager.providers:
+            ai_manager.providers[AIProvider.OLLAMA].update_model(args.ollama_model)
+    
     # Use the existing standard mode but with enhancements
-    return run_standard_mode(runner, target_url, user_timeout, interactive_mode)
+    return run_standard_mode(runner, target_url, user_timeout, interactive_mode, ai_provider, args)
 
 
-async def enhance_with_ai_analysis(result, target_url):
+async def enhance_with_ai_analysis(result, target_url, args):
     """Enhance scan results with AI analysis"""
     if not result:
         return
     
     print_info("Enhancing results with AI analysis...")
     
+    # Determine AI provider from arguments
+    ai_provider = None
+    if hasattr(args, 'ai_provider') and args.ai_provider and args.ai_provider != 'auto':
+        ai_provider = args.ai_provider
+        
+        # Enable Ollama if selected
+        if ai_provider == 'ollama':
+            import os
+            os.environ['ENABLE_OLLAMA'] = 'true'
+            # Reinitialize AI manager to include Ollama provider
+            ai_manager.reinitialize_providers()
+    
+    # Set Ollama model if specified
+    if hasattr(args, 'ollama_model') and args.ollama_model:
+        import os
+        os.environ['OLLAMA_MODEL'] = args.ollama_model
+        # Update the Ollama provider's model if it exists
+        if AIProvider.OLLAMA in ai_manager.providers:
+            ai_manager.providers[AIProvider.OLLAMA].update_model(args.ollama_model)
+    
     # Get AI response for advanced analysis
     scan_data = result.get('scan_history', [])
     if scan_data:
@@ -197,7 +242,17 @@ async def enhance_with_ai_analysis(result, target_url):
         """
         
         try:
-            ai_response = await ai_manager.get_response(prompt)
+            # Convert string provider to AIProvider enum if needed
+            provider_enum = None
+            if ai_provider:
+                try:
+                    from utils.ai_providers import AIProvider
+                    provider_enum = AIProvider(ai_provider)
+                except ValueError:
+                    print_warning(f"Invalid AI provider: {ai_provider}")
+                    return
+            
+            ai_response = await ai_manager.get_response(prompt, provider=provider_enum)
             if ai_response.success:
                 result['ai_analysis'] = ai_response.content
                 print_success("AI analysis completed")
@@ -317,7 +372,7 @@ def run_adaptive_mode(runner, target_url, user_timeout, interactive_mode):
         print_error("Adaptive testing failed. Check target URL and try again.")
         if result and "message" in result:
             print_info(f"Error: {result['message']}")
-def run_standard_mode(runner, target_url, user_timeout, interactive_mode):
+def run_standard_mode(runner, target_url, user_timeout, interactive_mode, ai_provider=None, args=None):
     print_info("Starting initial reconnaissance...")
     scan_history = []
     extracted_data = {}
@@ -346,11 +401,24 @@ def run_standard_mode(runner, target_url, user_timeout, interactive_mode):
             print_warning("Scan was interrupted by user. Stopping here.")
             return
         display_report(report)
-        print_info("Analyzing results with Groq AI and determining next steps...")
+        # Determine AI provider name for display
+        provider_name = "AI"
+        if ai_provider:
+            provider_name = ai_provider.upper()
+        print_info(f"Analyzing results with {provider_name} and determining next steps...")
+        # Determine if user wants advanced prompts
+        use_advanced = None
+        if hasattr(args, 'advanced') and args.advanced:
+            use_advanced = True
+        elif hasattr(args, 'simple') and args.simple:
+            use_advanced = False
+        
         next_options = ai_suggest_next_steps(
             report=report, 
             scan_history=scan_history,
-            extracted_data=extracted_data
+            extracted_data=extracted_data,
+            ai_provider=ai_provider,
+            use_advanced=use_advanced
         )
         if next_options:
             user_options = get_user_choice(next_options)
diff --git a/utils/ai_providers.py b/utils/ai_providers.py
@@ -123,7 +123,6 @@ async def get_response(
         """Get response from AI provider with fallback"""
         
         providers_to_try = [provider] if provider else self.active_providers
-        
         for attempt_provider in providers_to_try:
             if attempt_provider not in self.providers:
                 continue
@@ -149,6 +148,12 @@ async def get_response(
     def get_available_providers(self) -> List[AIProvider]:
         """Get list of available providers"""
         return self.active_providers.copy()
+    
+    def reinitialize_providers(self):
+        """Reinitialize providers (useful when environment variables change)"""
+        self.providers = {}
+        self.active_providers = []
+        self._setup_providers()
 
 
 class BaseAIProvider:
@@ -351,6 +356,10 @@ def __init__(self):
         self.default_model = os.getenv("OLLAMA_MODEL", "llama3.2")
         self.rate_limit_delay = 0.5
     
+    def update_model(self, model_name: str):
+        """Update the default model name"""
+        self.default_model = model_name
+    
     async def get_response(
         self, 
         prompt: str, 
@@ -363,6 +372,9 @@ async def get_response(
         model = model or self.default_model
         start_time = time.time()
         
+        logger.info(f"Ollama provider: Using model {model}")
+        logger.info(f"Ollama provider: Base URL {self.base_url}")
+        
         for attempt in range(max_retries):
             try:
                 await self._rate_limit()
@@ -388,7 +400,7 @@ async def get_response(
                             "num_predict": kwargs.get('max_tokens', 512)
                         }
                     },
-                    timeout=kwargs.get('timeout', 60)
+                    timeout=kwargs.get('timeout', 120)  # Increased timeout for complex prompts
                 )
                 
                 if response.status_code != 200:
@@ -407,6 +419,7 @@ async def get_response(
                 )
                 
             except Exception as e:
+                logger.warning(f"Ollama provider attempt {attempt + 1} failed: {e}")
                 if attempt == max_retries - 1:
                     return AIResponse(
                         content="",
