ploit-malper 0.2.6

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. A module
    // that is missing is skipped (auditable); a lookup that fails for any other
    // RPC reason is also skipped but reported distinctly so the operator can
    // tell "module absent" apart from "RPC down".
    let mut verified: Vec<(ExploitStep, Option<ModuleInfo>)> = Vec::new();
    for step in &plan {
        match framework.module_info(&step.module_type, &step.module_name) {
            Ok(info) => {
                verified.push((step.clone(), Some(info)));
            }
            Err(crate::error::AppError::RpcModuleNotFound(_)) => {
                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 {} (RPC failure, not a missing 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 resolved = resolve_options(info.as_ref(), step, opts.yes)?;
        if !resolved.missing.is_empty() {
            println!(
                "[!] Required options missing ({}); skipping {}.",
                resolved.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 = resolved.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>> {
    // Interactive TUI: a circular picker where the cursor wraps around the
    // ring (↑/k and ↓/j), space toggles, enter advances to the next step.
    if crate::frontend::tui::interactive() {
        let rows: Vec<crate::frontend::tui::PickerRow> = verified
            .iter()
            .map(|(step, _)| crate::frontend::tui::PickerRow {
                text: step.module_name.clone(),
                detail: Some(format!(
                    "rank={} conf={} host={} port={}",
                    step.rank,
                    step.confidence,
                    step.host,
                    step.service_port
                        .map(|p| p.to_string())
                        .unwrap_or_else(|| "-".to_string())
                )),
                selected: step.default_selected,
            })
            .collect();

        let first = verified
            .iter()
            .position(|(s, _)| s.default_selected)
            .unwrap_or(0);
        return match crate::frontend::tui::pick("Select modules to exploit", rows, first)? {
            Some(chosen) => Ok(verified
                .iter()
                .enumerate()
                .map(|(i, _)| chosen.contains(&i))
                .collect()),
            None => Ok(verified.iter().map(|_| false).collect()),
        };
    }

    // Non-interactive fallback: the classic checklist prompt.
    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)
}

/// Resolved option map for one step plus the options that remain truly
/// missing (required, no default, no value).
struct ResolvedOptions {
    /// Options to pass to `run_module`.
    options: std::collections::BTreeMap<String, String>,
    /// Required options with no default and no value after the pipeline.
    missing: Vec<String>,
}

/// Option resolution policy. This is the single place that decides how a
/// module's options become concrete values, and it applies identically to
/// every framework behind `ExploitFramework`.
///
/// Canonical required/optional × default matrix:
///
/// |                  | no default                | has default                    |
/// |------------------|---------------------------|--------------------------------|
/// | required         | prompt (interactive) else missing -> skip | auto-apply default; never prompt |
/// | optional         | omit (do not insert "")   | auto-apply default or omit     |
///
/// Additional rules:
/// - `inferred_options` (planner: RHOSTS/RPORT/target) always override defaults.
/// - Explicit CLI/payload values override inferred and defaults.
/// - UI/feedback options (ShowProgress, ShowProgressPercent, Verbose) never
///   block execution: apply default or omit; never hard-required.
/// - `--yes` / non-interactive: never prompt; missing required -> skip with
///   a clear reason (no stdin hang).
/// - An empty prompt answer is never "missing" when a default existed — the
///   default is applied before prompting, so nothing defaulted reaches the
///   prompt at all.
///
/// Defaults are applied first (step 1); inferred/CLI values always win (they
/// already occupy the map, so the default never overwrites them).
fn resolve_options(
    info: Option<&ModuleInfo>,
    step: &ExploitStep,
    yes: bool,
) -> Result<ResolvedOptions> {
    // Pure-framework UI toggles. They never gate execution: a default
    // auto-applies and no default means "omit", not "missing".
    const UI_OPTIONS: &[&str] = &["ShowProgress", "ShowProgressPercent", "Verbose", "VERBOSE"];

    let mut options = step.inferred_options.clone();
    let mut missing: Vec<String> = Vec::new();

    // 1. Apply defaults for every option that has one, but never over an
    //    inferred/explicit value already present.
    if let Some(info) = info {
        for opt in &info.options {
            if let Some(default) = &opt.default {
                options
                    .entry(opt.name.clone())
                    .or_insert_with(|| default.clone());
            }
        }
    }

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

    let interactive = crate::frontend::tui::interactive();

    // 2. Resolve the remaining required set. Defaulted/should-never-be here:
    //    required+default are auto-filled above, inferred are present above.
    for name in required {
        if options.contains_key(&name) {
            continue;
        }
        if UI_OPTIONS.contains(&name.as_str()) {
            continue;
        }
        if yes || !interactive {
            // Non-interactive or --yes: never prompt; mark missing.
            missing.push(name);
            continue;
        }
        let value = prompt_text(
            &format!("Required option {} for {}", name, step.module_name),
            "",
        )?;
        if value.trim().is_empty() {
            missing.push(name);
        } else {
            options.insert(name, value);
        }
    }

    Ok(ResolvedOptions { options, missing })
}

/// 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
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::BTreeMap;

    fn module_info_with(options: Vec<ModuleOption>, required: Vec<String>) -> ModuleInfo {
        ModuleInfo {
            name: "m".to_string(),
            module_type: MODULE_AUXILIARY.to_string(),
            rank: "normal".to_string(),
            description: String::new(),
            platforms: Vec::new(),
            targets: Vec::new(),
            options,
            required_options: required,
        }
    }

    fn option(name: &str, required: bool, default: Option<&str>) -> ModuleOption {
        ModuleOption {
            name: name.to_string(),
            required,
            default: default.map(str::to_string),
            description: String::new(),
        }
    }

    fn step(inferred: &[(&str, &str)]) -> ExploitStep {
        let inferred_options: BTreeMap<String, String> = inferred
            .iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect();
        ExploitStep {
            module_type: MODULE_AUXILIARY.to_string(),
            module_name: "auxiliary/scanner/http/robots_txt".to_string(),
            cve: None,
            vulnerability: "robots".to_string(),
            rank: "normal".to_string(),
            asset_id: "a1".to_string(),
            finding_id: "f1".to_string(),
            host: "10.0.0.1".to_string(),
            service_port: Some(80),
            service_protocol: Some("tcp".to_string()),
            inferred_options,
            missing_options: Vec::new(),
            payload_compatibility: "generic".to_string(),
            confidence: "high".to_string(),
            default_selected: true,
        }
    }

    #[test]
    fn required_with_default_auto_applies_without_prompting() {
        // In a test runner there is no tty, so interactive() == false. The
        // required+default options must be auto-filled (never "missing" and
        // never sent empty).
        let info = module_info_with(
            vec![
                option("ShowProgress", true, Some("false")),
                option("THREADS", true, Some("2")),
                option("RPORT", true, None),
            ],
            vec![
                "ShowProgress".to_string(),
                "THREADS".to_string(),
                "RPORT".to_string(),
            ],
        );
        let step = step(&[("RHOSTS", "10.0.0.1")]);
        let resolved = resolve_options(Some(&info), &step, true).unwrap();

        assert_eq!(
            resolved.options.get("ShowProgress").map(String::as_str),
            Some("false")
        );
        assert_eq!(
            resolved.options.get("THREADS").map(String::as_str),
            Some("2")
        );
        // RHOSTS inferred is preserved.
        assert_eq!(
            resolved.options.get("RHOSTS").map(String::as_str),
            Some("10.0.0.1")
        );
        // RPORT has no default and no inferred value -> missing.
        assert!(resolved.missing.contains(&"RPORT".to_string()));
    }

    #[test]
    fn inferred_options_always_win_over_defaults() {
        let info = module_info_with(
            vec![option("RHOSTS", true, Some("default.ignored"))],
            vec!["RHOSTS".to_string()],
        );
        let step = step(&[("RHOSTS", "10.0.0.9")]);
        let resolved = resolve_options(Some(&info), &step, true).unwrap();

        assert_eq!(
            resolved.options.get("RHOSTS").map(String::as_str),
            Some("10.0.0.9")
        );
        assert!(resolved.missing.is_empty());
    }

    #[test]
    fn non_interactive_yes_marks_missing_without_prompt() {
        // A required option with no default that cannot be filled will be
        // reported missing under --yes; it must not hang on stdin.
        let info = module_info_with(vec![option("LHOST", true, None)], vec!["LHOST".to_string()]);
        let step = step(&[]);
        let resolved = resolve_options(Some(&info), &step, true).unwrap();

        assert!(resolved.missing.contains(&"LHOST".to_string()));
        assert!(!resolved.options.contains_key("LHOST"));
    }

    #[test]
    fn optional_options_are_omitted_not_empty() {
        // An optional option with no default must simply be omitted from the
        // executed map (never inserted as "").
        let info = module_info_with(
            vec![
                option("OptionalOne", false, None),
                option("OptionalNow", false, None),
            ],
            Vec::new(),
        );
        let step = step(&[]);
        let resolved = resolve_options(Some(&info), &step, true).unwrap();

        assert!(resolved.missing.is_empty());
        assert!(!resolved.options.contains_key("OptionalOne"));
        assert!(!resolved.options.contains_key("OptionalNow"));
    }

    #[test]
    fn no_info_uses_step_missing_options() {
        let mut step = step(&[]);
        step.missing_options = vec!["RHOSTS".to_string()];
        let resolved = resolve_options(None, &step, true).unwrap();
        assert!(resolved.missing.contains(&"RHOSTS".to_string()));
    }
}