ploit-malper 0.2.6

pub mod cli;
pub mod models;
pub mod msf;
pub mod planner;

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

pub use models::*;

use crate::db::models::now_utc;
use crate::error::Result;

/// Execute a single planned step against the framework and report the
/// outcome. Runs the module, monitors the resulting job, then correlates any
/// new session or loot against the outcome.
pub fn execute_step(
    framework: &mut dyn ExploitFramework,
    step: &ExploitStep,
    payload: Option<&str>,
    run_timeout: Duration,
) -> Result<ExploitOutcome> {
    let start_time = now_utc();

    // Snapshot the workspace before the run so sessions/loot opened by this
    // job can be attributed to it instead of pinning the oldest existing one.
    let baseline_sessions: Vec<String> = framework
        .list_sessions()
        .unwrap_or_default()
        .into_iter()
        .map(|s| s.id)
        .collect();
    let baseline_loot: Vec<String> = framework
        .list_loot()
        .unwrap_or_default()
        .into_iter()
        .map(|l| l.id)
        .collect();

    let options = step.resolved_options();
    let result = framework.run_module(&step.module_type, &step.module_name, &options, payload)?;

    let mut session_id = None;
    let mut status = ExploitStatus::Failed;
    let mut error = None;
    let job_id = result.job_id.clone();

    if !result.success {
        error = Some(result.message.clone());
    } else if let Some(job) = &job_id {
        let job_status = framework.monitor_job(job, run_timeout)?;
        if job_status.error.is_some() {
            error = job_status.error.clone();
        } else if !job_status.running {
            status = ExploitStatus::Completed;
        } else {
            status = ExploitStatus::Running;
        }
    } else {
        status = ExploitStatus::Completed;
    }

    // Only surface sessions and loot that did not exist before this run.
    let after_sessions = framework.list_sessions().unwrap_or_default();
    let after_loot = framework.list_loot().unwrap_or_default();
    for session in &after_sessions {
        if !baseline_sessions.contains(&session.id) {
            session_id = Some(session.id.clone());
            break;
        }
    }
    let loot: Vec<String> = after_loot
        .iter()
        .filter(|l| !baseline_loot.contains(&l.id))
        .map(|l| l.id.clone())
        .collect();

    Ok(ExploitOutcome {
        step: step.clone(),
        status,
        job_id,
        session_id,
        error,
        loot,
        start_time,
        finish_time: Some(now_utc()),
    })
}

/// A framework-agnostic exploitation backend.
///
/// The CLI and database layers only ever talk to this trait; Metasploit is one
/// implementation behind it. This keeps the command and its history records
/// reusable if a different framework is added later.
pub trait ExploitFramework: Send {
    fn name(&self) -> &'static str;
    fn workspace(&self) -> &str;
    fn set_workspace(&mut self, workspace: &str);
    /// Verify connectivity and authenticate. Should be cheap when the client
    /// already holds a valid token.
    fn connect(&mut self) -> Result<()>;
    fn disconnect(&mut self);
    fn is_connected(&self) -> bool;

    /// Fetch full module metadata. Backends return an error whose variant is
    /// `AppError::RpcModuleNotFound` when the module does not exist, so
    /// callers can tell "missing" apart from a genuine RPC failure.
    fn module_info(&mut self, module_type: &str, name: &str) -> Result<ModuleInfo>;

    /// Launch a module with fully resolved options.
    fn run_module(
        &mut self,
        module_type: &str,
        name: &str,
        options: &BTreeMap<String, String>,
        payload: Option<&str>,
    ) -> Result<ModuleRunResult>;

    /// Poll a background job until it settles, `timeout` expires, or the
    /// optional `deadline` is reached.
    fn monitor_job(&mut self, job_id: &str, timeout: Duration) -> Result<JobStatus>;

    /// List current sessions.
    fn list_sessions(&mut self) -> Result<Vec<SessionInfo>>;

    /// List loot for the active workspace.
    fn list_loot(&mut self) -> Result<Vec<LootInfo>>;
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db::models::ExploitExecution;
    use crate::db::Storage;

    /// Deterministic framework stub for unit tests.
    struct MockFramework {
        workspace: String,
        run_succeeds: bool,
        job_id: String,
        job_finishes: bool,
        /// Sessions/loot that already existed when the run started.
        preexisting_sessions: Vec<SessionInfo>,
        preexisting_loot: Vec<LootInfo>,
        /// Sessions/loot opened by this run (surface after launch).
        new_sessions: Vec<SessionInfo>,
        new_loot: Vec<LootInfo>,
        ran: bool,
    }

    impl ExploitFramework for MockFramework {
        fn name(&self) -> &'static str {
            "mock"
        }
        fn workspace(&self) -> &str {
            &self.workspace
        }
        fn set_workspace(&mut self, workspace: &str) {
            self.workspace = workspace.to_string();
        }
        fn connect(&mut self) -> Result<()> {
            Ok(())
        }
        fn disconnect(&mut self) {}
        fn is_connected(&self) -> bool {
            true
        }
        fn module_info(&mut self, _t: &str, _n: &str) -> Result<ModuleInfo> {
            Ok(ModuleInfo::default())
        }
        fn run_module(
            &mut self,
            _t: &str,
            _n: &str,
            _o: &BTreeMap<String, String>,
            _p: Option<&str>,
        ) -> Result<ModuleRunResult> {
            self.ran = true;
            if self.run_succeeds {
                Ok(ModuleRunResult {
                    job_id: Some(self.job_id.clone()),
                    success: true,
                    message: "started".to_string(),
                    raw: String::new(),
                })
            } else {
                Ok(ModuleRunResult {
                    job_id: None,
                    success: false,
                    message: "boom".to_string(),
                    raw: String::new(),
                })
            }
        }
        fn monitor_job(&mut self, job_id: &str, _timeout: Duration) -> Result<JobStatus> {
            Ok(JobStatus {
                job_id: job_id.to_string(),
                running: !self.job_finishes,
                status: if self.job_finishes {
                    "completed".to_string()
                } else {
                    "still_running".to_string()
                },
                error: None,
            })
        }
        fn list_sessions(&mut self) -> Result<Vec<SessionInfo>> {
            let mut all = self.preexisting_sessions.clone();
            if self.ran {
                all.extend(self.new_sessions.clone());
            }
            Ok(all)
        }
        fn list_loot(&mut self) -> Result<Vec<LootInfo>> {
            let mut all = self.preexisting_loot.clone();
            if self.ran {
                all.extend(self.new_loot.clone());
            }
            Ok(all)
        }
    }

    fn sample_step() -> ExploitStep {
        ExploitStep {
            module_type: MODULE_EXPLOIT.to_string(),
            module_name: "exploit/multi/http/apache_path_traversal".to_string(),
            cve: Some("CVE-2021-41773".to_string()),
            vulnerability: "Apache path traversal".to_string(),
            rank: "excellent".to_string(),
            asset_id: "asset-1".to_string(),
            finding_id: "finding-1".to_string(),
            host: "10.0.0.5".to_string(),
            service_port: Some(80),
            service_protocol: Some("tcp".to_string()),
            inferred_options: BTreeMap::from([("RHOSTS".to_string(), "10.0.0.5".to_string())]),
            missing_options: vec!["LHOST".to_string()],
            payload_compatibility: "generic".to_string(),
            confidence: "high".to_string(),
            default_selected: true,
        }
    }

    #[test]
    fn execute_step_records_completed_outcome() {
        let mut framework = MockFramework {
            workspace: "default".to_string(),
            run_succeeds: true,
            job_id: "7".to_string(),
            job_finishes: true,
            preexisting_sessions: Vec::new(),
            preexisting_loot: Vec::new(),
            new_sessions: vec![SessionInfo {
                id: "1".to_string(),
                typ: "meterpreter".to_string(),
                target_host: "10.0.0.5".to_string(),
                platform: "linux".to_string(),
                info: "test".to_string(),
            }],
            new_loot: vec![LootInfo {
                id: "loot-1".to_string(),
                name: "shell".to_string(),
                content_type: "text/plain".to_string(),
                host: "10.0.0.5".to_string(),
                path: "/tmp/loot".to_string(),
            }],
            ran: false,
        };
        let outcome =
            execute_step(&mut framework, &sample_step(), None, Duration::from_secs(5)).unwrap();
        assert_eq!(outcome.status, ExploitStatus::Completed);
        assert_eq!(outcome.job_id.as_deref(), Some("7"));
        // The session/loot were opened by this run, so they are attributed.
        assert_eq!(outcome.session_id.as_deref(), Some("1"));
        assert!(outcome.loot.iter().any(|l| l == "loot-1"));
        assert!(outcome.error.is_none());
    }

    #[test]
    fn execute_step_does_not_credit_preexisting_sessions() {
        // A session that was already in the workspace before the run must not
        // be attributed to this job; only newly created ones are.
        let mut framework = MockFramework {
            workspace: "default".to_string(),
            run_succeeds: true,
            job_id: "7".to_string(),
            job_finishes: true,
            preexisting_sessions: vec![SessionInfo {
                id: "old-1".to_string(),
                typ: "shell".to_string(),
                target_host: "10.0.0.9".to_string(),
                platform: "linux".to_string(),
                info: "pre-existing".to_string(),
            }],
            preexisting_loot: vec![LootInfo {
                id: "loot-old".to_string(),
                name: "shell".to_string(),
                content_type: "text/plain".to_string(),
                host: "10.0.0.9".to_string(),
                path: "/tmp/loot-old".to_string(),
            }],
            new_sessions: vec![SessionInfo {
                id: "new-2".to_string(),
                typ: "meterpreter".to_string(),
                target_host: "10.0.0.5".to_string(),
                platform: "linux".to_string(),
                info: "from this run".to_string(),
            }],
            new_loot: vec![LootInfo {
                id: "loot-2".to_string(),
                name: "shell".to_string(),
                content_type: "text/plain".to_string(),
                host: "10.0.0.5".to_string(),
                path: "/tmp/loot-2".to_string(),
            }],
            ran: false,
        };
        let outcome =
            execute_step(&mut framework, &sample_step(), None, Duration::from_secs(5)).unwrap();
        assert_eq!(outcome.session_id.as_deref(), Some("new-2"));
        assert!(!outcome.loot.iter().any(|l| l == "loot-old"));
        assert_eq!(outcome.loot, vec!["loot-2".to_string()]);
    }

    #[test]
    fn execute_step_marks_failure() {
        let mut framework = MockFramework {
            workspace: "default".to_string(),
            run_succeeds: false,
            job_id: String::new(),
            job_finishes: true,
            preexisting_sessions: Vec::new(),
            preexisting_loot: Vec::new(),
            new_sessions: Vec::new(),
            new_loot: Vec::new(),
            ran: false,
        };
        let outcome =
            execute_step(&mut framework, &sample_step(), None, Duration::from_secs(5)).unwrap();
        assert_eq!(outcome.status, ExploitStatus::Failed);
        assert_eq!(outcome.error.as_deref(), Some("boom"));
    }

    #[test]
    fn execution_status_roundtrips() {
        for status in [
            ExploitStatus::Planned,
            ExploitStatus::Running,
            ExploitStatus::Completed,
            ExploitStatus::Failed,
            ExploitStatus::Skipped,
            ExploitStatus::DryRun,
        ] {
            assert_eq!(ExploitStatus::parse(status.as_str()), status);
        }
    }

    #[test]
    fn persist_execution_to_sqlite_roundtrips() {
        let mut storage = crate::db::sqlite::SqliteStorage::open(":memory:").unwrap();
        storage.ensure_schema().unwrap();

        let mut exec =
            ExploitExecution::new("run123", "asset-1", "finding-1", "exploit/multi/http/x");
        exec.module_type = MODULE_EXPLOIT.to_string();
        exec.host = "10.0.0.5".to_string();
        exec.status = ExploitStatus::Completed.as_str().to_string();
        exec.job_id = Some("9".to_string());
        exec.session_id = Some("2".to_string());
        exec.loot = vec!["loot-9".to_string()];
        exec.selected = true;
        storage.upsert_exploit_execution(&exec).unwrap();

        let loaded = storage
            .get_exploit_execution(&exec.execution_id)
            .unwrap()
            .unwrap();
        assert_eq!(loaded.run_id, "run123");
        assert_eq!(loaded.module, "exploit/multi/http/x");
        assert_eq!(loaded.status, "completed");
        assert_eq!(loaded.job_id.as_deref(), Some("9"));
        assert_eq!(loaded.session_id.as_deref(), Some("2"));
        assert_eq!(loaded.loot, vec!["loot-9".to_string()]);
        assert!(loaded.selected);

        let by_run = storage.list_exploit_executions_for_run("run123").unwrap();
        assert_eq!(by_run.len(), 1);
        let all = storage.list_exploit_executions().unwrap();
        assert_eq!(all.len(), 1);
    }
}