ploit-malper 0.2.6

use std::collections::{BTreeMap, BTreeSet};

use crate::db::models::{Asset, Finding, ScanRun, Service};
use crate::db::Storage;
use crate::engine::matcher;
use crate::error::{AppError, Result};
use crate::msf;

use super::models::*;

/// One actionable (asset, finding, service, cve) tuple gathered for a run.
#[derive(Debug, Clone)]
pub struct ExploitTarget {
    pub asset: Asset,
    pub finding: Finding,
    pub service: Option<Service>,
    pub cve: Option<String>,
}

/// Plan discovery: turns a run into a list of module steps that could be
/// executed against its hosts. Works entirely offline against the database
/// and the local module catalogs; live framework verification happens later.
pub struct Planner;

impl Planner {
    /// Resolve a run id that may be a full `stable_id` or a unique prefix.
    pub fn resolve_run(storage: &mut dyn Storage, run_ref: &str) -> Result<ScanRun> {
        if let Some(run) = storage.get_scan_run(run_ref)? {
            return Ok(run);
        }
        let runs = storage.list_scan_runs()?;
        let matches: Vec<&ScanRun> = runs
            .iter()
            .filter(|run| run.stable_id.starts_with(run_ref))
            .collect();
        match matches.len() {
            1 => Ok(matches[0].clone()),
            0 => Err(AppError::Message(format!(
                "no scan run matches '{}'",
                run_ref
            ))),
            _ => Err(AppError::Message(format!(
                "run id '{}' is ambiguous ({} matches); use a longer prefix",
                run_ref,
                matches.len()
            ))),
        }
    }

    /// Collect the assets that belong to a run. A run's primary target usually
    /// matches the asset name; additionally any asset referenced by an
    /// observation recorded under this run id is included.
    pub fn collect_assets(storage: &mut dyn Storage, run: &ScanRun) -> Result<Vec<Asset>> {
        let mut seen = BTreeSet::new();
        let mut assets = Vec::new();

        for asset in storage.list_assets()? {
            let matches_target = asset.name.eq_ignore_ascii_case(&run.target)
                || asset
                    .ip
                    .as_deref()
                    .is_some_and(|ip| ip.eq_ignore_ascii_case(&run.target))
                || asset
                    .fqdn
                    .as_deref()
                    .is_some_and(|fqdn| fqdn.eq_ignore_ascii_case(&run.target));
            if matches_target && seen.insert(asset.stable_id.clone()) {
                assets.push(asset);
            }
        }

        for observation in storage.list_all_observations()? {
            if observation.run_id == run.stable_id
                && observation.subject_type == "asset"
                && seen.insert(observation.subject_id.clone())
            {
                if let Some(asset) = storage.get_asset(&observation.subject_id)? {
                    assets.push(asset);
                }
            }
        }

        Ok(assets)
    }

    /// Gather every (asset, finding, service, cve) tuple for a run.
    pub fn collect_targets(storage: &mut dyn Storage, run: &ScanRun) -> Result<Vec<ExploitTarget>> {
        let mut targets = Vec::new();
        for asset in Self::collect_assets(storage, run)? {
            let services = storage.list_services_for_asset(&asset.stable_id)?;
            for finding in storage.list_findings_for_asset(&asset.stable_id)? {
                if finding.status == crate::db::models::ASSET_REMOVED {
                    continue;
                }
                let service = finding
                    .service_id
                    .as_deref()
                    .and_then(|id| services.iter().find(|s| s.stable_id == id))
                    .cloned();
                let cves = finding.cves.clone();
                if cves.is_empty() {
                    targets.push(ExploitTarget {
                        asset: asset.clone(),
                        finding: finding.clone(),
                        service: service.clone(),
                        cve: None,
                    });
                } else {
                    for cve in cves {
                        targets.push(ExploitTarget {
                            asset: asset.clone(),
                            finding: finding.clone(),
                            service: service.clone(),
                            cve: Some(cve),
                        });
                    }
                }
            }
        }
        Ok(targets)
    }

    /// Build the exploitation plan: every module mapped to any target, deduped
    /// by (module, host). `--module` / `--target` filters apply here.
    pub fn build_plan(
        storage: &mut dyn Storage,
        run: &ScanRun,
        module_filter: Option<&str>,
        target_filter: Option<&str>,
    ) -> Result<Vec<ExploitStep>> {
        let targets = Self::collect_targets(storage, run)?;
        if targets.is_empty() {
            return Err(AppError::Message(format!(
                "run '{}' (target={}) has no active findings to exploit",
                run.stable_id, run.target
            )));
        }

        let mut steps = Vec::new();
        let mut seen = BTreeSet::new();

        for target in &targets {
            if let Some(filter) = target_filter {
                let hay = format!(
                    "{} {} {}",
                    target.asset.name,
                    target.asset.ip.clone().unwrap_or_default(),
                    target.asset.fqdn.clone().unwrap_or_default()
                )
                .to_lowercase();
                if !hay.contains(&filter.to_lowercase()) {
                    continue;
                }
            }

            // Assemble the full suggestion set for this target. CVE maps,
            // banner matches, technology/CPE matches and title matches each
            // contribute independently; the union is deduped so a single CVE
            // module never silently crowds out title/technology candidates.
            let mut module_names: Vec<String> = Vec::new();
            if let Some(cve) = &target.cve {
                for module in msf::modules_for_cve(cve) {
                    module_names.push(module.name.clone());
                }
                if let Some(info) = crate::enrich::enrich_cve(cve) {
                    module_names.extend(info.metasploit_modules);
                }
            }
            if let Some(service) = &target.service {
                let banner = service
                    .banner
                    .clone()
                    .unwrap_or_else(|| service.product.clone().unwrap_or_default());
                if !banner.is_empty() {
                    for suggestion in matcher::suggest_modules(&banner) {
                        module_names.push(suggestion.suggested_module);
                    }
                }
                let techs: Vec<String> = service
                    .technologies
                    .iter()
                    .map(|t| t.to_lowercase())
                    .collect();
                if !techs.is_empty() {
                    for module in msf::modules_for_technology(&techs) {
                        module_names.push(module.name);
                    }
                }
            }
            module_names.extend(
                matcher::suggest_from_title(
                    &target.finding.title,
                    target.finding.detail.as_deref().unwrap_or(""),
                )
                .into_iter()
                .map(|s| s.suggested_module),
            );

            module_names.sort();
            module_names.dedup();
            if let Some(filter) = module_filter {
                module_names.retain(|name| name.contains(filter));
            }

            for name in module_names {
                let key = (name.clone(), target.asset.name.clone());
                if !seen.insert(key) {
                    continue;
                }
                let (module_type, module_leaf) = split_module_name(&name);
                let confidence = if target.cve.is_some() {
                    "high"
                } else {
                    "medium"
                };
                let rank = msf::modules_for_cve(target.cve.as_deref().unwrap_or(""))
                    .iter()
                    .find(|m| m.name == name)
                    .map(|m| m.rank.clone())
                    .unwrap_or_else(|| {
                        if module_type == MODULE_AUXILIARY {
                            "normal".to_string()
                        } else {
                            "good".to_string()
                        }
                    });
                let inferred = infer_options(target, module_type);
                let required = required_option_hint(&name, module_type);
                let missing = required
                    .iter()
                    .filter(|opt| !inferred.contains_key(*opt))
                    .cloned()
                    .collect::<Vec<_>>();

                steps.push(ExploitStep {
                    module_type: module_type.to_string(),
                    module_name: name.clone(),
                    cve: target.cve.clone(),
                    vulnerability: target.finding.title.clone(),
                    rank,
                    asset_id: target.asset.stable_id.clone(),
                    finding_id: target.finding.stable_id.clone(),
                    host: host_for(&target.asset),
                    service_port: target.service.as_ref().map(|s| s.port),
                    service_protocol: target.service.as_ref().map(|s| s.protocol.clone()),
                    inferred_options: inferred.clone(),
                    missing_options: missing.clone(),
                    payload_compatibility: payload_hint(&name, module_type),
                    confidence: confidence.to_string(),
                    default_selected: module_type == MODULE_EXPLOIT && missing.is_empty(),
                });
                let _ = module_leaf;
            }
        }

        if steps.is_empty() {
            return Err(AppError::Message(
                "no Metasploit modules could be mapped to this run's findings".to_string(),
            ));
        }

        steps.sort_by(|a, b| {
            a.module_name
                .cmp(&b.module_name)
                .then_with(|| a.host.cmp(&b.host))
        });
        Ok(steps)
    }
}

fn split_module_name(name: &str) -> (&str, &str) {
    match name.split_once('/') {
        Some((first, rest)) if first == "exploit" || first == "auxiliary" => (first, rest),
        _ => (MODULE_EXPLOIT, name),
    }
}

fn host_for(asset: &Asset) -> String {
    if let Some(ip) = asset.ip.as_deref().filter(|ip| !ip.is_empty()) {
        return ip.to_string();
    }
    if let Some(fqdn) = asset.fqdn.as_deref().filter(|f| !f.is_empty()) {
        return fqdn.to_string();
    }
    asset.name.clone()
}

/// Infer options from database records for a target.
fn infer_options(target: &ExploitTarget, module_type: &str) -> BTreeMap<String, String> {
    let mut options = BTreeMap::new();
    let host = host_for(&target.asset);
    options.insert("RHOSTS".to_string(), host.clone());

    if let Some(service) = &target.service {
        options.insert("RPORT".to_string(), service.port.to_string());
        let protocol = service.protocol.to_lowercase();
        if protocol.contains("ssl") || protocol.contains("https") {
            options.insert("SSL".to_string(), "true".to_string());
        }
        if module_type == MODULE_EXPLOIT || protocol.contains("http") {
            if let Some(fqdn) = target.asset.fqdn.as_deref().filter(|f| !f.is_empty()) {
                options.insert("VHOST".to_string(), fqdn.to_string());
                options.insert("HTTPHOST".to_string(), fqdn.to_string());
            } else {
                options.insert("HTTPHOST".to_string(), host.clone());
            }
        }
    }

    if let Some(target_url) = &target.finding.target_url {
        let without_query = target_url.split('?').next().unwrap_or(target_url);
        let path = if let Some(scheme_end) = without_query.find("://") {
            let after_scheme = &without_query[scheme_end + 3..];
            match after_scheme.find('/') {
                Some(i) => after_scheme[i..].to_string(),
                None => String::new(),
            }
        } else {
            without_query.to_string()
        };
        let cleaned = path
            .trim_end_matches('/')
            .chars()
            .take(512)
            .collect::<String>();
        if !cleaned.is_empty() {
            options.insert("TARGETURI".to_string(), cleaned);
        }
    }

    options
}

/// Required options that cannot be inferred from the database for this module.
fn required_option_hint(name: &str, module_type: &str) -> Vec<String> {
    let mut required = vec!["RHOSTS".to_string()];
    if name.contains("http") || name.contains("webapp") || name.contains("struts") {
        required.push("TARGETURI".to_string());
    }
    if name.contains("cve_2019_0708")
        || name.contains("zerologon")
        || name.contains("smb")
        || name.contains("rdp")
        || name.contains("ssh_version")
        || name.contains("mysql_version")
        || name.contains("postgres_version")
        || name.contains("ftp_version")
    {
        required.push("RPORT".to_string());
    }
    if module_type == MODULE_EXPLOIT {
        required.push("LHOST".to_string());
    }
    required.dedup();
    required
}

/// A short, honest payload compatibility hint shown in the plan.
fn payload_hint(name: &str, module_type: &str) -> String {
    if module_type == MODULE_AUXILIARY {
        return "n/a (auxiliary — no payload)".to_string();
    }
    if name.contains("windows") {
        "windows/meterpreter/reverse_tcp (LHOST required)".to_string()
    } else if name.contains("linux") || name.contains("unix") {
        "linux/x64/meterpreter/reverse_tcp (LHOST required)".to_string()
    } else {
        "generic/shell_reverse_tcp (LHOST required)".to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db::models::{Asset, Finding, ScanRun, Service};

    fn seed_storage() -> crate::db::sqlite::SqliteStorage {
        let mut storage = crate::db::sqlite::SqliteStorage::open(":memory:").unwrap();
        let now = crate::db::models::now_utc();

        let mut run = ScanRun::new("run123", "10.0.0.5", "/tmp/scan", "abc", None);
        run.finished_at = Some(now.clone());
        storage.upsert_scan_run(&run).unwrap();

        let mut asset = Asset::new("10.0.0.5", "ip");
        asset.ip = Some("10.0.0.5".to_string());
        storage.upsert_asset(&asset).unwrap();

        let mut service = Service::new(&asset.stable_id, 80, "tcp", "http");
        service.banner = Some("Apache/2.4.49".to_string());
        storage.upsert_service(&service).unwrap();

        let mut finding = Finding::new(&asset.stable_id, "vulnmalper", "Apache path traversal");
        finding.service_id = Some(service.stable_id.clone());
        finding.cves = vec!["CVE-2021-41773".to_string()];
        finding.target_url = Some("https://10.0.0.5/cgi-bin/".to_string());
        finding.severity = "high".to_string();
        storage.upsert_finding(&finding).unwrap();

        storage
    }

    #[test]
    fn resolve_run_by_prefix() {
        let mut storage = seed_storage();
        let run = Planner::resolve_run(&mut storage, "run12").unwrap();
        assert_eq!(run.stable_id, "run123");
        assert_eq!(run.target, "10.0.0.5");
    }

    #[test]
    fn collect_assets_matches_run_target() {
        let mut storage = seed_storage();
        let run = Planner::resolve_run(&mut storage, "run123").unwrap();
        let assets = Planner::collect_assets(&mut storage, &run).unwrap();
        assert_eq!(assets.len(), 1);
        assert_eq!(assets[0].name, "10.0.0.5");
    }

    #[test]
    fn build_plan_maps_cve_to_module() {
        let mut storage = seed_storage();
        let run = Planner::resolve_run(&mut storage, "run123").unwrap();
        let plan = Planner::build_plan(&mut storage, &run, None, None).unwrap();

        let step = plan
            .iter()
            .find(|s| s.module_name == "exploit/multi/http/apache_path_traversal")
            .expect("apache_path_traversal should be planned");
        assert_eq!(step.cve.as_deref(), Some("CVE-2021-41773"));
        assert_eq!(step.host, "10.0.0.5");
        assert_eq!(step.service_port, Some(80));
        assert_eq!(step.confidence, "high");
        assert_eq!(
            step.inferred_options.get("RHOSTS").map(String::as_str),
            Some("10.0.0.5")
        );
        assert_eq!(
            step.inferred_options.get("RPORT").map(String::as_str),
            Some("80")
        );
        assert_eq!(
            step.inferred_options.get("TARGETURI").map(String::as_str),
            Some("/cgi-bin")
        );
    }

    #[test]
    fn build_plan_applies_module_filter() {
        let mut storage = seed_storage();
        let run = Planner::resolve_run(&mut storage, "run123").unwrap();
        let plan = Planner::build_plan(&mut storage, &run, Some("apache_path"), None).unwrap();
        assert!(!plan.is_empty());
        assert!(plan.iter().all(|s| s.module_name.contains("apache_path")));
    }

    #[test]
    fn build_plan_keeps_the_full_suggestion_set_per_target() {
        // A CVE-driven module must not crowd out banner/title candidates for
        // the same target: the plan surfaces every suggestion, deduped.
        let mut storage = seed_storage();
        let run = Planner::resolve_run(&mut storage, "run123").unwrap();
        let plan = Planner::build_plan(&mut storage, &run, None, None).unwrap();

        let by_name: std::collections::HashMap<String, &ExploitStep> =
            plan.iter().map(|s| (s.module_name.clone(), s)).collect();

        assert!(by_name.contains_key("exploit/multi/http/apache_path_traversal"));
        // The Apache banner also maps to a scanner; the single CVE module must
        // not cause it (or its sibling exploit) to be dropped.
        assert!(by_name.contains_key("auxiliary/scanner/http/apache_version"));
        assert!(by_name.contains_key("exploit/multi/http/apache_mod_cgi_bash_env_exec"));
        assert_eq!(by_name.len(), 3);
    }

    #[test]
    fn build_plan_cve_free_target_uses_banner_and_title() {
        let mut storage = seed_storage();
        let now = crate::db::models::now_utc();

        let mut asset = Asset::new("10.1.1.1", "ip");
        asset.ip = Some("10.1.1.1".to_string());
        storage.upsert_asset(&asset).unwrap();

        let mut service = Service::new(&asset.stable_id, 22, "tcp", "ssh");
        service.banner = Some("OpenSSH_8.2p1".to_string());
        storage.upsert_service(&service).unwrap();

        let mut finding = Finding::new(&asset.stable_id, "vulnmalper", "OpenSSH banner disclosure");
        finding.service_id = Some(service.stable_id.clone());
        finding.severity = "low".to_string();
        storage.upsert_finding(&finding).unwrap();

        let mut run =
            crate::db::models::ScanRun::new("runfree", "10.1.1.1", "/tmp/free", "hashfree", None);
        run.finished_at = Some(now);
        storage.upsert_scan_run(&run).unwrap();

        let plan = Planner::build_plan(&mut storage, &run, None, None).unwrap();
        let names: Vec<&str> = plan.iter().map(|s| s.module_name.as_str()).collect();
        // No CVE in play; only the openssh banner/title suggestions should
        // surface, and all of them, not a single one.
        assert!(names.contains(&"auxiliary/scanner/ssh/ssh_version"));
        assert_eq!(plan.len(), 1, "expected exactly the ssh_version suggestion");
    }

    #[test]
    fn delete_scan_run_removes_run_and_owned_records_only() {
        let mut storage = seed_storage();
        let run = Planner::resolve_run(&mut storage, "run123").unwrap();
        let asset = storage.list_assets().unwrap().remove(0);

        let obs = crate::db::models::Observation::new(
            &run.stable_id,
            "asset",
            &asset.stable_id,
            "asset_discovered",
            serde_json::json!({}),
            serde_json::json!({"ip": "10.0.0.5"}),
            "discovered",
        );
        storage.add_observation(&obs).unwrap();

        let exec = crate::db::models::ExploitExecution::new(
            &run.stable_id,
            &asset.stable_id,
            "vuln-1",
            "exploit/multi/http/apache_path_traversal",
        );
        storage.upsert_exploit_execution(&exec).unwrap();

        assert_eq!(storage.list_all_observations().unwrap().len(), 1);
        assert_eq!(storage.list_exploit_executions().unwrap().len(), 1);

        storage.delete_scan_run(&run.stable_id).unwrap();

        assert!(storage.get_scan_run(&run.stable_id).unwrap().is_none());
        assert_eq!(storage.list_all_observations().unwrap().len(), 0);
        assert_eq!(storage.list_exploit_executions().unwrap().len(), 0);
        assert_eq!(storage.list_assets().unwrap().len(), 1);
        assert_eq!(storage.list_all_services().unwrap().len(), 1);
        assert_eq!(storage.list_all_findings().unwrap().len(), 1);
    }

    #[test]
    fn host_prefers_ip() {
        let mut asset = Asset::new("host.example", "hostname");
        asset.ip = Some("192.168.1.10".to_string());
        assert_eq!(host_for(&asset), "192.168.1.10");
    }
}