compliance-scanner-agent/compliance-dast/src/tools/recon.rs

use std::collections::HashMap;

use compliance_core::error::CoreError;
use compliance_core::traits::pentest_tool::{PentestTool, PentestToolContext, PentestToolResult};
use serde_json::json;
use tracing::info;

use crate::recon::ReconAgent;

/// PentestTool wrapper around the existing ReconAgent.
///
/// Performs HTTP header fingerprinting and technology detection.
/// Returns structured recon data for the LLM to use when planning attacks.
pub struct ReconTool {
    http: reqwest::Client,
    agent: ReconAgent,
}

impl ReconTool {
    pub fn new(http: reqwest::Client) -> Self {
        let agent = ReconAgent::new(http.clone());
        Self { http, agent }
    }
}

impl PentestTool for ReconTool {
    fn name(&self) -> &str {
        "recon"
    }

    fn description(&self) -> &str {
        "Performs reconnaissance on a target URL. Fingerprints HTTP headers, detects server \
         technologies and frameworks. Returns structured data about the target's technology stack."
    }

    fn input_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "url": {
                    "type": "string",
                    "description": "Base URL to perform reconnaissance on"
                },
                "additional_paths": {
                    "type": "array",
                    "description": "Optional additional paths to probe for technology fingerprinting",
                    "items": { "type": "string" }
                }
            },
            "required": ["url"]
        })
    }

    fn execute<'a>(
        &'a self,
        input: serde_json::Value,
        _context: &'a PentestToolContext,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<PentestToolResult, CoreError>> + Send + 'a>,
    > {
        Box::pin(async move {
            let url = input
                .get("url")
                .and_then(|v| v.as_str())
                .ok_or_else(|| CoreError::Dast("Missing required 'url' parameter".to_string()))?;

            let additional_paths: Vec<String> = input
                .get("additional_paths")
                .and_then(|v| v.as_array())
                .map(|arr| {
                    arr.iter()
                        .filter_map(|v| v.as_str().map(String::from))
                        .collect()
                })
                .unwrap_or_default();

            let result = self.agent.scan(url).await?;

            // Scan additional paths for more technology signals
            let mut extra_technologies: Vec<String> = Vec::new();
            let mut extra_headers: HashMap<String, String> = HashMap::new();

            let base_url = url.trim_end_matches('/');
            for path in &additional_paths {
                let probe_url = format!("{base_url}/{}", path.trim_start_matches('/'));
                if let Ok(resp) = self.http.get(&probe_url).send().await {
                    for (key, value) in resp.headers() {
                        let k = key.to_string().to_lowercase();
                        let v = value.to_str().unwrap_or("").to_string();

                        // Look for technology indicators
                        if (k == "x-powered-by" || k == "server" || k == "x-generator")
                            && !result.technologies.contains(&v)
                            && !extra_technologies.contains(&v)
                        {
                            extra_technologies.push(v.clone());
                        }
                        extra_headers.insert(format!("{probe_url} -> {k}"), v);
                    }
                }
            }

            let mut all_technologies = result.technologies.clone();
            all_technologies.extend(extra_technologies);
            all_technologies.dedup();

            let tech_count = all_technologies.len();
            info!(url, technologies = tech_count, "Recon complete");

            Ok(PentestToolResult {
                summary: format!(
                    "Recon complete for {url}. Detected {} technologies. Server: {}.",
                    tech_count,
                    result.server.as_deref().unwrap_or("unknown")
                ),
                findings: Vec::new(), // Recon produces data, not findings
                data: json!({
                    "base_url": url,
                    "server": result.server,
                    "technologies": all_technologies,
                    "interesting_headers": result.interesting_headers,
                    "extra_headers": extra_headers,
                    "open_ports": result.open_ports,
                }),
            })
        })
    }
}