ploit_malper/exploit/

msf.rs

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();
            let required: Vec<String> = options
                .iter()
                .filter(|opt| opt.required)
                .map(|opt| opt.name.clone())
                .collect();
            if !required.is_empty() {
                info.required_options = required;
            }
        }
        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);
    let mut required_options = parse_strings(raw, "required_options");
    if required_options.is_empty() {
        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(MsgValue::as_str)
                .map(str::to_string);
            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
}

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"]);
    }
}