ploit-malper 0.2.2

use std::io::{self, Write};
use std::time::Duration;

use crate::config::ConfigManager;
use crate::db::models::ExploitExecution;
use crate::error::{AppError, Result};
use crate::exploit::models::*;
use crate::exploit::msf::MsfFramework;
use crate::exploit::planner::Planner;
use crate::exploit::{execute_step, ExploitFramework};
use crate::frontend::{open_backend, Args};

const DEFAULT_JOB_TIMEOUT_S: u64 = 300;

pub fn usage() {
    println!(
        "Usage: ploit-malper exploit <run_id> [--dry-run] [--yes] [--module <name>] \
         [--target <host>] [--payload <name>] [--workspace <name>] [--job-timeout <secs>] \
         [--backend pocketbase|sqlite] [--pocketbase-url URL] [--sqlite-path PATH] [--verbose]"
    );
}

pub fn cmd_exploit(args: &[String], config_mgr: &mut ConfigManager) -> Result<()> {
    let parsed = Args::parse(args)?;
    if parsed.has("help") {
        usage();
        return Ok(());
    }

    let run_ref = parsed.positionals().first().cloned().ok_or_else(|| {
        AppError::Message(
            "exploit requires a run id (e.g. ploit-malper exploit <run_id>)".to_string(),
        )
    })?;

    let opts = ExploitOptions {
        dry_run: parsed.has("dry-run"),
        yes: parsed.has("yes"),
        module_filter: parsed.get("module").map(str::to_string),
        target_filter: parsed.get("target").map(str::to_string),
        payload_override: parsed.get("payload").map(str::to_string),
        workspace: parsed.get("workspace").map(str::to_string),
        backend: parsed.get("backend").map(str::to_string),
        pocketbase_url: parsed.get("pocketbase-url").map(str::to_string),
        sqlite_path: parsed.get("sqlite-path").map(str::to_string),
        verbose: parsed.has("verbose"),
    };
    let job_timeout = parsed
        .get("job-timeout")
        .map(|v| v.parse::<u64>())
        .transpose()
        .map_err(|_| AppError::Message("--job-timeout must be a number of seconds".to_string()))?
        .unwrap_or(DEFAULT_JOB_TIMEOUT_S);

    let mut storage = open_backend(config_mgr, &parsed)?;
    let run = Planner::resolve_run(storage.as_mut(), &run_ref)?;

    println!(
        "[+] Run: {} (target={})",
        short(&run.stable_id, 16),
        run.target
    );

    let plan = Planner::build_plan(
        storage.as_mut(),
        &run,
        opts.module_filter.as_deref(),
        opts.target_filter.as_deref(),
    )?;

    if opts.dry_run {
        return dry_run(&mut storage, &run, &plan, &opts);
    }

    let mut framework = build_framework(config_mgr, &opts)?;
    if let Err(e) = framework.connect() {
        // Record the whole plan as skipped so the attempt is auditable.
        println!("[!] Could not connect to the exploitation framework: {e}");
        for step in &plan {
            persist(&mut storage, &run.stable_id, step, false, None, None)?;
        }
        println!("[!] Plan recorded as skipped; nothing was executed.");
        return Ok(());
    }
    println!(
        "[+] Connected to {} (workspace '{}')",
        framework.name(),
        framework.workspace()
    );

    // Verify every planned module exists on the connected instance.
    let mut verified: Vec<(ExploitStep, Option<ModuleInfo>)> = Vec::new();
    for step in &plan {
        match framework.module_exists(&step.module_type, &step.module_name) {
            Ok(true) => {
                let info = framework
                    .module_info(&step.module_type, &step.module_name)
                    .ok();
                verified.push((step.clone(), info));
            }
            Ok(false) => {
                println!(
                    "[!] Module {} not present on the connected instance; marking skipped.",
                    step.module_name
                );
                persist(
                    &mut storage,
                    &run.stable_id,
                    step,
                    false,
                    Some(ExploitStatus::Skipped),
                    None,
                )?;
            }
            Err(e) => {
                println!("[!] Could not query module {}: {e}", step.module_name);
                persist(
                    &mut storage,
                    &run.stable_id,
                    step,
                    false,
                    Some(ExploitStatus::Skipped),
                    None,
                )?;
            }
        }
    }

    if verified.is_empty() {
        println!("[!] No modules are available on the connected instance.");
        return Ok(());
    }

    render_plan(&verified);

    // Interactive checklist: toggle modules, Enter to run the selected set.
    let selected: Vec<bool> = if opts.yes {
        vec![true; verified.len()]
    } else {
        select_steps(&verified)?
    };

    let chosen: Vec<(ExploitStep, Option<ModuleInfo>)> = verified
        .into_iter()
        .zip(selected.iter())
        .filter(|(_, sel)| **sel)
        .map(|(step, _)| step)
        .collect();

    if chosen.is_empty() {
        println!("[!] No modules selected; recording plan as skipped.");
        for step in &plan {
            persist(&mut storage, &run.stable_id, step, false, None, None)?;
        }
        return Ok(());
    }

    if !opts.yes
        && !prompt_bool_with_default(
            &format!("Run {} selected module(s) now?", chosen.len()),
            false,
        )?
    {
        println!("[!] Execution declined; recording selection as skipped.");
        for (step, _) in &chosen {
            persist(
                &mut storage,
                &run.stable_id,
                step,
                true,
                Some(ExploitStatus::Skipped),
                None,
            )?;
        }
        return Ok(());
    }

    // Execute each selected step sequentially.
    println!();
    for (step, info) in &chosen {
        println!(
            "[+] Executing {} against {} ...",
            step.module_name, step.host
        );

        let options = resolve_options(info.as_ref(), step, opts.yes)?;
        let missing: Vec<String> = options
            .iter()
            .filter(|(_, value)| value.is_empty())
            .map(|(key, _)| key.clone())
            .collect();
        if !missing.is_empty() {
            println!(
                "[!] Required options missing ({}); skipping {}.",
                missing.join(", "),
                step.module_name
            );
            persist(
                &mut storage,
                &run.stable_id,
                step,
                true,
                Some(ExploitStatus::Skipped),
                None,
            )?;
            continue;
        }

        let mut step_clone = step.clone();
        step_clone.inferred_options = options;
        let payload = opts.payload_override.as_deref();
        match execute_step(
            framework.as_mut(),
            &step_clone,
            payload,
            Duration::from_secs(job_timeout),
        ) {
            Ok(outcome) => {
                render_outcome(&outcome);
                persist(
                    &mut storage,
                    &run.stable_id,
                    &outcome.step,
                    true,
                    Some(outcome.status),
                    Some(&outcome),
                )?;
            }
            Err(e) => {
                println!("[!] Execution error for {}: {e}", step.module_name);
                persist(
                    &mut storage,
                    &run.stable_id,
                    step,
                    true,
                    Some(ExploitStatus::Failed),
                    None,
                )?;
            }
        }
    }

    framework.disconnect();
    println!();
    println!("[+] Exploit session complete. Records written to the database.");
    Ok(())
}

fn build_framework(
    config_mgr: &ConfigManager,
    opts: &ExploitOptions,
) -> Result<Box<dyn ExploitFramework>> {
    if !config_mgr.is_configured() {
        return Err(AppError::Message(
            "no MSF-RPC configuration found; run 'ploit-malper setup' first".to_string(),
        ));
    }
    let cfg = config_mgr.get_msfrpc_config().clone();
    let mut framework =
        MsfFramework::new(&cfg.host, cfg.port, &cfg.username, &cfg.password, cfg.ssl);
    framework.set_workspace(opts.workspace.as_deref().unwrap_or(&cfg.workspace));
    Ok(Box::new(framework))
}

/// Dry-run: show the full plan with module metadata and record it as previews.
/// Never contacts the framework.
fn dry_run(
    storage: &mut Box<dyn crate::db::Storage>,
    run: &crate::db::models::ScanRun,
    plan: &[ExploitStep],
    opts: &ExploitOptions,
) -> Result<()> {
    println!();
    println!("[+] Dry run — no modules will be executed.");
    render_plan_dry(plan);
    println!();
    if let Some(payload) = &opts.payload_override {
        println!("[+] Payload override: {payload}");
    }
    println!("[+] Recording {} step(s) as dry_run previews.", plan.len());
    for step in plan {
        persist(
            storage,
            &run.stable_id,
            step,
            false,
            Some(ExploitStatus::DryRun),
            None,
        )?;
    }
    println!("[+] Preview recorded; nothing was executed.");
    Ok(())
}

fn render_plan(verified: &[(ExploitStep, Option<ModuleInfo>)]) {
    println!();
    println!("Exploitation plan (verified against the connected instance):");
    println!(
        "{:<4} {:<8} {:<6} {:<5} {:<16} {:<6} {:<10}",
        "#", "TYPE", "RANK", "CONF", "MODULE", "MISSING", "PAYLOAD"
    );
    println!("{}", "-".repeat(90));
    for (i, (step, info)) in verified.iter().enumerate() {
        let missing = if step.missing_options.is_empty() {
            "none".to_string()
        } else {
            step.missing_options.join(",")
        };
        println!(
            "{:<4} {:<8} {:<6} {:<5} {:<16} {:<6} {:<10}",
            i + 1,
            step.module_type,
            step.rank,
            step.confidence,
            short(&step.module_name, 16),
            short(&missing, 6),
            short(&step.payload_compatibility, 10),
        );
        println!(
            "      host={} port={} cve={} vuln={}",
            step.host,
            step.service_port
                .map(|p| p.to_string())
                .unwrap_or_else(|| "-".to_string()),
            step.cve.as_deref().unwrap_or("-"),
            short(&step.vulnerability, 60),
        );
        if !step.inferred_options.is_empty() {
            let opts = step
                .inferred_options
                .iter()
                .map(|(k, v)| format!("{}={}", k, v))
                .collect::<Vec<_>>()
                .join(" ");
            println!("      options: {opts}");
        }
        if let Some(info) = info {
            if !info.description.is_empty() {
                println!("      desc: {}", short(&info.description, 80));
            }
            let defaults = info
                .options
                .iter()
                .filter(|o| o.default.is_some())
                .map(|o| format!("{}={}", o.name, o.default.as_deref().unwrap_or("")))
                .collect::<Vec<_>>()
                .join(" ");
            if !defaults.is_empty() {
                println!("      defaults: {defaults}");
            }
        }
    }
}

fn render_plan_dry(plan: &[ExploitStep]) {
    println!();
    println!("Planned modules (offline metadata):");
    println!(
        "{:<4} {:<8} {:<6} {:<16} {:<16} {:<6} {:<10}",
        "#", "TYPE", "RANK", "MODULE", "CVE", "CONF", "PAYLOAD"
    );
    println!("{}", "-".repeat(90));
    for (i, step) in plan.iter().enumerate() {
        println!(
            "{:<4} {:<8} {:<6} {:<16} {:<16} {:<6} {:<10}",
            i + 1,
            step.module_type,
            step.rank,
            short(&step.module_name, 16),
            short(step.cve.as_deref().unwrap_or("-"), 16),
            step.confidence,
            short(&step.payload_compatibility, 10),
        );
        println!(
            "      target={} port={} vuln={}",
            step.host,
            step.service_port
                .map(|p| p.to_string())
                .unwrap_or_else(|| "-".to_string()),
            short(&step.vulnerability, 60),
        );
        if !step.missing_options.is_empty() {
            println!(
                "      required options to fill: {}",
                step.missing_options.join(", ")
            );
        } else {
            println!("      required options: fully inferred");
        }
    }
}

fn render_outcome(outcome: &ExploitOutcome) {
    let status = outcome.status.as_str();
    println!("      -> {} status={}", outcome.step.module_name, status);
    if let Some(job) = &outcome.job_id {
        println!("         job_id={}", job);
    }
    if let Some(session) = &outcome.session_id {
        println!("         session_id={}", session);
    }
    if !outcome.loot.is_empty() {
        println!("         loot={}", outcome.loot.join(","));
    }
    if let Some(error) = &outcome.error {
        println!("         error={}", short(error, 120));
    }
}

fn select_steps(verified: &[(ExploitStep, Option<ModuleInfo>)]) -> Result<Vec<bool>> {
    let mut selected: Vec<bool> = verified.iter().map(|(s, _)| s.default_selected).collect();
    loop {
        println!();
        println!("Module checklist — [x] = recommended, [ ] = optional:");
        for (i, (step, _)) in verified.iter().enumerate() {
            let mark = if selected[i] { "x" } else { " " };
            println!(
                "  [{}] {:<3} {}  [{}] host={}",
                mark,
                i + 1,
                step.module_name,
                step.rank,
                step.host
            );
        }
        print!("  Toggle (e.g. 1,3), 'a' all, 'n' none, or press Enter to run the selected set: ");
        io::stdout().flush()?;
        let mut buffer = String::new();
        io::stdin().read_line(&mut buffer)?;
        let input = buffer.trim().to_lowercase();
        if input.is_empty() {
            break;
        }
        match input.as_str() {
            "a" => selected.iter_mut().for_each(|s| *s = true),
            "n" => selected.iter_mut().for_each(|s| *s = false),
            _ => {
                for token in input.split([',', ' ', '\t']) {
                    if token.is_empty() {
                        continue;
                    }
                    match token.parse::<usize>() {
                        Ok(idx) if idx >= 1 && idx <= selected.len() => {
                            let slot = &mut selected[idx - 1];
                            *slot = !*slot;
                        }
                        _ => println!("    ignoring '{}' (not a module number)", token),
                    }
                }
            }
        }
    }
    Ok(selected)
}

/// Fill any required option not already inferred. Without `--yes` the missing
/// options are prompted for; with `--yes` framework defaults are used and the
/// value is left empty when no default exists.
fn resolve_options(
    info: Option<&ModuleInfo>,
    step: &ExploitStep,
    yes: bool,
) -> Result<std::collections::BTreeMap<String, String>> {
    let mut options = step.inferred_options.clone();

    let required: Vec<String> = match info {
        Some(info) => info.required_options.to_vec(),
        None => step.missing_options.clone(),
    };

    let default_for = |name: &str| -> Option<String> {
        info.and_then(|info| {
            info.options
                .iter()
                .find(|o| o.name == name)
                .and_then(|o| o.default.clone())
        })
    };

    for name in required {
        if options.contains_key(&name) {
            continue;
        }
        let default = default_for(&name);
        if yes {
            if let Some(default) = default {
                options.insert(name.clone(), default);
            } else {
                options.insert(name.clone(), String::new());
            }
            continue;
        }
        let value = prompt_text(
            &format!("Required option {} for {}", name, step.module_name),
            &default.unwrap_or_default(),
        )?;
        options.insert(name.clone(), value);
    }

    Ok(options)
}

/// Persist one step as an `ExploitExecution` record. When `outcome` is given,
/// the record reflects the finished execution; otherwise it reflects the
/// planned/skipped state.
fn persist(
    storage: &mut Box<dyn crate::db::Storage>,
    run_id: &str,
    step: &ExploitStep,
    selected: bool,
    status: Option<ExploitStatus>,
    outcome: Option<&ExploitOutcome>,
) -> Result<()> {
    let mut exec =
        ExploitExecution::new(run_id, &step.asset_id, &step.finding_id, &step.module_name);
    exec.module_type = step.module_type.clone();
    exec.host = step.host.clone();
    exec.options = step.inferred_options.clone();
    exec.selected = selected;
    match outcome {
        Some(outcome) => {
            exec.status = outcome.status.as_str().to_string();
            exec.start_time = outcome.start_time.clone();
            exec.finish_time = outcome.finish_time.clone();
            exec.job_id = outcome.job_id.clone();
            exec.session_id = outcome.session_id.clone();
            exec.error = outcome.error.clone();
            exec.loot = outcome.loot.clone();
        }
        None => {
            exec.status = status
                .unwrap_or(if selected {
                    ExploitStatus::Planned
                } else {
                    ExploitStatus::Skipped
                })
                .as_str()
                .to_string();
        }
    }
    storage.upsert_exploit_execution(&exec)
}

fn prompt_text(prompt: &str, default: &str) -> Result<String> {
    print!("  {} [{}]: ", prompt, default);
    io::stdout().flush()?;
    let mut buffer = String::new();
    io::stdin().read_line(&mut buffer)?;
    let value = buffer.trim();
    Ok(if value.is_empty() {
        default.to_string()
    } else {
        value.to_string()
    })
}

fn prompt_bool_with_default(prompt: &str, default: bool) -> Result<bool> {
    let suffix = if default { "[Y/n]" } else { "[y/N]" };
    loop {
        print!("  {} {}: ", prompt, suffix);
        io::stdout().flush()?;
        let mut buffer = String::new();
        io::stdin().read_line(&mut buffer)?;
        let value = buffer.trim().to_lowercase();
        if value.is_empty() {
            return Ok(default);
        }
        match value.as_str() {
            "y" | "yes" | "true" => return Ok(true),
            "n" | "no" | "false" => return Ok(false),
            _ => println!("  Please answer yes or no."),
        }
    }
}

fn short(value: &str, width: usize) -> String {
    let count = value.chars().count();
    if count <= width {
        return value.to_string();
    }
    if width <= 1 {
        return "…".to_string();
    }
    let mut out: String = value.chars().take(width - 1).collect();
    out.push('…');
    out
}