ploit-malper 0.2.6

use std::collections::BTreeMap;
use std::time::Duration;

use crate::engine::msfrpc::MsfRpcClient;
use crate::error::{AppError, Result};
use crate::msgpack::MsgValue;

use super::models::*;
use super::ExploitFramework;

/// Metasploit implementation of [`ExploitFramework`].
///
/// All framework communication is delegated to the existing `MsfRpcClient`,
/// keeping Metasploit specifics isolated behind the trait.
pub struct MsfFramework {
    client: Option<MsfRpcClient>,
    host: String,
    port: u16,
    ssl: bool,
    username: String,
    password: String,
    workspace: String,
}

impl MsfFramework {
    pub fn new(host: &str, port: u16, username: &str, password: &str, ssl: bool) -> Self {
        Self {
            client: None,
            host: host.to_string(),
            port,
            ssl,
            username: username.to_string(),
            password: password.to_string(),
            workspace: "default".to_string(),
        }
    }
}

impl ExploitFramework for MsfFramework {
    fn name(&self) -> &'static str {
        "metasploit"
    }

    fn workspace(&self) -> &str {
        &self.workspace
    }

    fn set_workspace(&mut self, workspace: &str) {
        self.workspace = workspace.to_string();
    }

    fn connect(&mut self) -> Result<()> {
        let mut client = MsfRpcClient::new(
            &self.host,
            self.port,
            &self.username,
            &self.password,
            self.ssl,
        );
        let login = client
            .login()
            .map_err(|e| AppError::Network(format!("MSF-RPC login transport failure: {e}")))?;
        if !login.success {
            return Err(AppError::Message(format!(
                "MSF-RPC authentication failed: {}",
                login.error
            )));
        }
        self.client = Some(client);
        Ok(())
    }

    fn disconnect(&mut self) {
        if let Some(client) = self.client.as_mut() {
            client.logout();
        }
        self.client = None;
    }

    fn is_connected(&self) -> bool {
        self.client
            .as_ref()
            .is_some_and(MsfRpcClient::is_authenticated)
    }

    fn module_info(&mut self, module_type: &str, name: &str) -> Result<ModuleInfo> {
        let client = self.client()?;
        let raw = client.module_info(module_type, name)?;
        let mut info = parse_module_info(module_type, name, &raw)?;
        // `module.info` returns metadata only; the real datastore options come
        // from `module.options`. Merge them so required options and defaults
        // reflect the connected instance rather than offline hints.
        if let Ok(options_raw) = client.module_options(module_type, name) {
            let options = parse_options(&options_raw);
            info.options = options.clone();
            // "Required" here means required AND without a default: an option
            // that is required:true but carries a framework default (e.g.
            // ShowProgress, THREADS) never needs operator input. This keeps
            // `required_options` aligned with the offline path.
            info.required_options = options
                .iter()
                .filter(|opt| opt.required && opt.default.is_none())
                .map(|opt| opt.name.clone())
                .collect();
        }
        Ok(info)
    }

    fn run_module(
        &mut self,
        module_type: &str,
        name: &str,
        options: &BTreeMap<String, String>,
        payload: Option<&str>,
    ) -> Result<ModuleRunResult> {
        let client = self.client()?;
        let raw = client.module_execute(module_type, name, options, payload)?;
        parse_run_result(&raw)
    }

    fn monitor_job(&mut self, job_id: &str, timeout: Duration) -> Result<JobStatus> {
        let client = self.client()?;
        let started = std::time::Instant::now();
        loop {
            match client.job_list() {
                Ok(list) => {
                    let present = job_present(&list, job_id);
                    if !present {
                        return Ok(JobStatus {
                            job_id: job_id.to_string(),
                            running: false,
                            status: "completed".to_string(),
                            error: None,
                        });
                    }
                }
                Err(e) => {
                    return Ok(JobStatus {
                        job_id: job_id.to_string(),
                        running: false,
                        status: "failed".to_string(),
                        error: Some(e.to_string()),
                    });
                }
            }
            if started.elapsed() >= timeout {
                return Ok(JobStatus {
                    job_id: job_id.to_string(),
                    running: true,
                    status: "still_running".to_string(),
                    error: None,
                });
            }
            std::thread::sleep(Duration::from_millis(1000));
        }
    }

    fn list_sessions(&mut self) -> Result<Vec<SessionInfo>> {
        let client = self.client()?;
        let raw = client
            .session_list()
            .map_err(|e| AppError::Network(e.to_string()))?;
        Ok(parse_sessions(&raw))
    }

    fn list_loot(&mut self) -> Result<Vec<LootInfo>> {
        let workspace = self.workspace.clone();
        let client = self.client()?;
        let raw = client
            .db_loots(&workspace)
            .map_err(|e| AppError::Network(e.to_string()))?;
        Ok(parse_loot(&raw))
    }
}

impl MsfFramework {
    fn client(&mut self) -> Result<&mut MsfRpcClient> {
        self.client.as_mut().ok_or_else(|| {
            AppError::Message("MSF-RPC not connected; call connect() first".to_string())
        })
    }
}

fn parse_module_info(module_type: &str, name: &str, raw: &MsgValue) -> Result<ModuleInfo> {
    let rank = raw
        .get("rank")
        .and_then(MsgValue::as_str)
        .map(str::to_string)
        .or_else(|| {
            // Modern servers report rank as a numeric value; translate it back
            // to the familiar label so the plan stays readable.
            raw.get("rank").and_then(MsgValue::as_u64).map(rank_label)
        })
        .unwrap_or_default();
    let description = raw
        .get("description")
        .and_then(MsgValue::as_str)
        .unwrap_or("")
        .to_string();

    let mut platforms = parse_strings(raw, "platforms");
    if platforms.is_empty() {
        if let Some(single) = raw.get("platform").and_then(MsgValue::as_str) {
            platforms.push(single.to_string());
        }
    }
    let targets = parse_strings(raw, "targets");

    let options = parse_options(raw);
    // Canonical definition: an option is "required" only when it is marked
    // required AND carries no default. Options that are required but defaulted
    // (ShowProgress, THREADS, ...) never need operator input, so they must not
    // surface in the required list.
    let required_options = options
        .iter()
        .filter(|opt| opt.required && opt.default.is_none())
        .map(|opt| opt.name.clone())
        .collect();

    Ok(ModuleInfo {
        name: name.to_string(),
        module_type: module_type.to_string(),
        rank,
        description,
        platforms,
        targets,
        options,
        required_options,
    })
}

/// Translate Metasploit's numeric rank value to the familiar label.
fn rank_label(rank: u64) -> String {
    match rank {
        0 => "manual".to_string(),
        100 => "low".to_string(),
        200 => "average".to_string(),
        300 => "normal".to_string(),
        400 => "good".to_string(),
        500 => "great".to_string(),
        600 => "excellent".to_string(),
        other => format!("rank-{other}"),
    }
}

fn parse_options(raw: &MsgValue) -> Vec<ModuleOption> {
    let mut options = Vec::new();
    if let Some(map) = raw.get("options").and_then(MsgValue::as_map) {
        for (name, opt_val) in map {
            let required = opt_val.get("required").is_some_and(|v| match v {
                MsgValue::Bool(b) => *b,
                MsgValue::String(s) => s == "true",
                MsgValue::UInt(n) => *n != 0,
                MsgValue::Int(n) => *n != 0,
                _ => false,
            });
            let default = opt_val.get("default").and_then(parse_default_value);
            let description = opt_val
                .get("desc")
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            options.push(ModuleOption {
                name: name.clone(),
                required,
                default,
                description,
            });
        }
    }
    options
}

/// Convert a framework option default into its canonical string form so every
/// scalar (bool/int/uint/float/string) is represented identically. A
/// non-scalar or absent default maps to `None` ("no default"); an empty string
/// default stays `Some("")` — it is distinct from "no default".
fn parse_default_value(value: &MsgValue) -> Option<String> {
    match value {
        MsgValue::Bool(b) => Some(b.to_string()),
        MsgValue::Int(n) => Some(n.to_string()),
        MsgValue::UInt(n) => Some(n.to_string()),
        MsgValue::Float(f) => Some(format_float(f)),
        MsgValue::String(s) => Some(s.clone()),
        _ => None,
    }
}

/// Stable string form for a float default. `to_string` on f64 yields the
/// shortest round-trippable representation (e.g. `2.0` -> `"2"`, `2.5` ->
/// `"2.5"`), which is what MSF expects for a datastore value.
fn format_float(value: &f64) -> String {
    value.to_string()
}

fn parse_run_result(raw: &MsgValue) -> Result<ModuleRunResult> {
    let raw_text = format!("{raw:?}");

    let error_flag = match raw.get("error") {
        Some(MsgValue::Bool(b)) => *b,
        Some(MsgValue::String(s)) => s == "true",
        Some(MsgValue::UInt(n)) => *n != 0,
        _ => false,
    };
    if error_flag {
        let message = raw
            .get("error_message")
            .or_else(|| raw.get("error_string"))
            .and_then(MsgValue::as_str)
            .unwrap_or("module execution failed")
            .to_string();
        return Ok(ModuleRunResult {
            job_id: None,
            success: false,
            message,
            raw: raw_text,
        });
    }

    let job_id = raw
        .get("job_id")
        .and_then(MsgValue::as_u64)
        .map(|id| id.to_string());
    let message = raw
        .get("message")
        .and_then(MsgValue::as_str)
        .unwrap_or("module started")
        .to_string();
    Ok(ModuleRunResult {
        job_id,
        success: true,
        message,
        raw: raw_text,
    })
}

fn job_present(list: &MsgValue, job_id: &str) -> bool {
    if let Some(map) = list.get("jobs").and_then(MsgValue::as_map) {
        map.keys().any(|id| id == job_id)
    } else if let Some(map) = list.as_map() {
        map.get("jobs").is_none() && map.keys().any(|id| id == job_id)
    } else {
        false
    }
}

fn parse_sessions(raw: &MsgValue) -> Vec<SessionInfo> {
    let mut sessions = Vec::new();
    if let Some(map) = raw.get("sessions").and_then(MsgValue::as_map) {
        for (id, val) in map {
            let typ = val
                .get("type")
                .or_else(|| val.get("session_type"))
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            let target_host = val
                .get("target_host")
                .or_else(|| val.get("host"))
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            let platform = val
                .get("platform")
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            let info = val
                .get("info")
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            sessions.push(SessionInfo {
                id: id.clone(),
                typ,
                target_host,
                platform,
                info,
            });
        }
    }
    sessions
}

fn parse_loot(raw: &MsgValue) -> Vec<LootInfo> {
    let mut loot = Vec::new();
    if let Some(array) = raw.get("loot").and_then(MsgValue::as_array) {
        for val in array {
            let id = val
                .get("loot_id")
                .or_else(|| val.get("id"))
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            let name = val
                .get("name")
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            let content_type = val
                .get("content_type")
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            let host = val
                .get("host")
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            let path = val
                .get("path")
                .and_then(MsgValue::as_str)
                .unwrap_or("")
                .to_string();
            loot.push(LootInfo {
                id,
                name,
                content_type,
                host,
                path,
            });
        }
    }
    loot
}

fn parse_strings(raw: &MsgValue, key: &str) -> Vec<String> {
    match raw.get(key) {
        Some(MsgValue::Array(arr)) => arr
            .iter()
            .filter_map(MsgValue::as_str)
            .map(str::to_string)
            .collect(),
        // Modern servers report `targets`/`platforms` as a map keyed by the
        // numeric index, e.g. {0=>"Automatic Target", 1=>"Windows 7", ...}.
        // The decoder stringifies integer keys, so sort them back into order.
        Some(MsgValue::Map(map)) => {
            let mut items: Vec<&String> = map.keys().collect();
            items.sort();
            items
                .into_iter()
                .filter_map(|k| map.get(k))
                .filter_map(MsgValue::as_str)
                .map(str::to_string)
                .collect()
        }
        _ => Vec::new(),
    }
}

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

    fn raw_map(pairs: &[(&str, MsgValue)]) -> MsgValue {
        MsgValue::Map(
            pairs
                .iter()
                .map(|(k, v)| (k.to_string(), v.clone()))
                .collect(),
        )
    }

    #[test]
    fn parse_strings_reads_array_form() {
        let raw = raw_map(&[(
            "targets",
            MsgValue::Array(vec![
                MsgValue::String("Linux x86".to_string()),
                MsgValue::String("Linux x86_64".to_string()),
            ]),
        )]);
        assert_eq!(
            parse_strings(&raw, "targets"),
            vec!["Linux x86", "Linux x86_64"]
        );
    }

    #[test]
    fn parse_strings_reads_indexed_map_form() {
        // Modern servers encode targets as {0=>"A", 1=>"B", ...}; the decoder
        // stringifies the integer keys so they must be re-sorted here.
        let mut targets = BTreeMap::new();
        targets.insert("1".to_string(), MsgValue::String("Windows 7".to_string()));
        targets.insert(
            "0".to_string(),
            MsgValue::String("Automatic Target".to_string()),
        );
        targets.insert("2".to_string(), MsgValue::String("Windows 8".to_string()));
        let raw = raw_map(&[("targets", MsgValue::Map(targets))]);
        assert_eq!(
            parse_strings(&raw, "targets"),
            vec!["Automatic Target", "Windows 7", "Windows 8"]
        );
    }

    #[test]
    fn parse_module_info_extracts_indexed_targets() {
        // Full-shaped module.info with integer-keyed targets (as returned by
        // exploit modules like ms17_010_eternalblue).
        let mut targets = BTreeMap::new();
        targets.insert(
            "0".to_string(),
            MsgValue::String("Automatic Target".to_string()),
        );
        targets.insert("1".to_string(), MsgValue::String("Windows 7".to_string()));
        let raw = raw_map(&[
            ("name", MsgValue::String("eternalblue".to_string())),
            ("description", MsgValue::String("MS17-010".to_string())),
            ("targets", MsgValue::Map(targets)),
            ("options", MsgValue::Map(BTreeMap::new())),
        ]);
        let info = parse_module_info("exploit", "windows/smb/ms17_010_eternalblue", &raw).unwrap();
        assert_eq!(info.targets, vec!["Automatic Target", "Windows 7"]);
    }

    #[test]
    fn parse_default_value_stringifies_all_scalars() {
        assert_eq!(
            parse_default_value(&MsgValue::Bool(true)),
            Some("true".to_string())
        );
        assert_eq!(
            parse_default_value(&MsgValue::Bool(false)),
            Some("false".to_string())
        );
        assert_eq!(
            parse_default_value(&MsgValue::Int(-3)),
            Some("-3".to_string())
        );
        assert_eq!(
            parse_default_value(&MsgValue::UInt(2)),
            Some("2".to_string())
        );
        assert_eq!(
            parse_default_value(&MsgValue::Float(2.0)),
            Some("2".to_string())
        );
        assert_eq!(
            parse_default_value(&MsgValue::Float(2.5)),
            Some("2.5".to_string())
        );
        assert_eq!(
            parse_default_value(&MsgValue::String("eth0".to_string())),
            Some("eth0".to_string())
        );
        assert_eq!(parse_default_value(&MsgValue::Nil), None);
        assert_eq!(parse_default_value(&MsgValue::Map(BTreeMap::new())), None);
    }

    #[test]
    fn required_with_default_is_not_in_required_options() {
        // Simulates auxiliary/scanner/http/robots_txt: options that are
        // required:true but carry bool/int defaults (ShowProgress, THREADS)
        // must NOT surface in `required_options`.
        let mut options = BTreeMap::new();
        options.insert(
            "ShowProgress".to_string(),
            MsgValue::Map(BTreeMap::from([
                ("required".to_string(), MsgValue::Bool(true)),
                ("default".to_string(), MsgValue::Bool(false)),
                (
                    "desc".to_string(),
                    MsgValue::String("show progress".to_string()),
                ),
            ])),
        );
        options.insert(
            "THREADS".to_string(),
            MsgValue::Map(BTreeMap::from([
                ("required".to_string(), MsgValue::Bool(true)),
                ("default".to_string(), MsgValue::UInt(2)),
                ("desc".to_string(), MsgValue::String("threads".to_string())),
            ])),
        );
        options.insert(
            "RPORT".to_string(),
            MsgValue::Map(BTreeMap::from([
                ("required".to_string(), MsgValue::Bool(true)),
                ("desc".to_string(), MsgValue::String("port".to_string())),
            ])),
        );

        let raw = raw_map(&[
            ("description", MsgValue::String("robots".to_string())),
            ("options", MsgValue::Map(options)),
        ]);
        let info = parse_module_info("auxiliary", "scanner/http/robots_txt", &raw).unwrap();

        // ShowProgress/THREADS keep their defaults and are dropped from the
        // required list; RPORT has no default so it stays required.
        assert!(info.required_options.contains(&"RPORT".to_string()));
        assert!(!info.required_options.contains(&"ShowProgress".to_string()));
        assert!(!info.required_options.contains(&"THREADS".to_string()));

        // Defaults are stringified so downstream resolution can auto-apply.
        let progress = info
            .options
            .iter()
            .find(|o| o.name == "ShowProgress")
            .unwrap();
        assert_eq!(progress.default.as_deref(), Some("false"));
        let threads = info.options.iter().find(|o| o.name == "THREADS").unwrap();
        assert_eq!(threads.default.as_deref(), Some("2"));
        let rport = info.options.iter().find(|o| o.name == "RPORT").unwrap();
        assert_eq!(rport.default, None);
    }
}