#!/usr/bin/env python3
"""
Web Pentest Agent - Specialized agent for web application penetration testing.
"""

import json
import logging
from typing import Dict, List
from core.llm_manager import LLMManager
from tools.web_pentest import WebRecon # Import the moved WebRecon tool

logger = logging.getLogger(__name__)

class WebPentestAgent:
    """Agent responsible for comprehensive web application penetration testing."""
    
    def __init__(self, config: Dict):
        """Initializes the WebPentestAgent."""
        self.config = config
        self.llm = LLMManager(config)
        self.web_recon = WebRecon(config)
        # Placeholder for web exploitation tools if they become separate classes
        # self.web_exploiter = WebExploiter(config)
        logger.info("WebPentestAgent initialized")
    
    def execute(self, target: str, context: Dict) -> Dict:
        """Executes the web application penetration testing phase."""
        logger.info(f"Starting web pentest on {target}")
        
        results = {
            "target": target,
            "status": "running",
            "web_recon_results": {},
            "vulnerability_analysis": [],
            "exploitation_attempts": [],
            "ai_analysis": {}
        }
        
        try:
            # Phase 1: Web Reconnaissance
            logger.info("Phase 1: Web Reconnaissance (WebPentestAgent)")
            web_recon_output = self.web_recon.analyze(target)
            results["web_recon_results"] = web_recon_output
            
            # Phase 2: Vulnerability Analysis (AI-powered)
            logger.info("Phase 2: AI-powered Vulnerability Analysis")
            # This part will be improved later with more detailed vulnerability detection in WebRecon
            # For now, it will look for findings reported by WebRecon
            
            potential_vulnerabilities = self._identify_potential_web_vulnerabilities(web_recon_output)
            
            if potential_vulnerabilities:
                results["vulnerability_analysis"] = potential_vulnerabilities
                ai_vulnerability_analysis = self._ai_analyze_web_vulnerabilities(potential_vulnerabilities, target)
                results["ai_analysis"]["vulnerability_insights"] = ai_vulnerability_analysis
            else:
                logger.info("No immediate web vulnerabilities identified by WebRecon.")

            # Phase 3: Web Exploitation (Placeholder for now)
            # This will integrate with exploitation tools later.
            
            results["status"] = "completed"
            logger.info("Web pentest phase completed")
            
        except Exception as e:
            logger.error(f"Error during web pentest: {e}")
            results["status"] = "error"
            results["error"] = str(e)
        
        return results

    def _identify_potential_web_vulnerabilities(self, web_recon_output: Dict) -> List[Dict]:
        """
        Identifies potential web vulnerabilities based on WebRecon output.
        This is a placeholder and will be enhanced as WebRecon improves.
        """
        vulnerabilities = []
        if "vulnerabilities" in web_recon_output:
            vulnerabilities.extend(web_recon_output["vulnerabilities"])
        return vulnerabilities

    def _ai_analyze_web_vulnerabilities(self, vulnerabilities: List[Dict], target: str) -> Dict:
        """Uses AI to analyze identified web vulnerabilities."""
        prompt = self.llm.get_prompt(
            "web_recon",
            "ai_analysis_user",
            default=f"""
Analyze the following potential web vulnerabilities identified on {target} and provide insights:

Vulnerabilities: {json.dumps(vulnerabilities, indent=2)}

Provide:
1. Prioritized list of vulnerabilities
2. Recommended exploitation steps for each (if applicable)
3. Potential impact
4. Remediation suggestions

Response in JSON format with actionable recommendations.
"""
        )
        
        system_prompt = self.llm.get_prompt(
            "web_recon",
            "ai_analysis_system",
            default="""You are an expert web penetration tester and security analyst.
Provide precise analysis of web vulnerabilities and practical advice for exploitation and remediation."""
        )
        
        try:
            # Format the user prompt with recon_data
            formatted_prompt = prompt.format(
                target=target,
                vulnerabilities_json=json.dumps(vulnerabilities, indent=2)
            )
            response = self.llm.generate(formatted_prompt, system_prompt)
            return json.loads(response)
        except Exception as e:
            logger.error(f"AI web vulnerability analysis error: {e}")
            return {"error": str(e), "raw_response": response if 'response' in locals() else None}