1pub mod cli;
2pub mod models;
3pub mod msf;
4pub mod planner;
5
6use std::collections::BTreeMap;
7use std::time::Duration;
8
9pub use models::*;
10
11use crate::db::models::now_utc;
12use crate::error::Result;
13
14pub fn execute_step(
18 framework: &mut dyn ExploitFramework,
19 step: &ExploitStep,
20 payload: Option<&str>,
21 run_timeout: Duration,
22) -> Result<ExploitOutcome> {
23 let start_time = now_utc();
24
25 let baseline_sessions: Vec<String> = framework
28 .list_sessions()
29 .unwrap_or_default()
30 .into_iter()
31 .map(|s| s.id)
32 .collect();
33 let baseline_loot: Vec<String> = framework
34 .list_loot()
35 .unwrap_or_default()
36 .into_iter()
37 .map(|l| l.id)
38 .collect();
39
40 let options = step.resolved_options();
41 let result = framework.run_module(&step.module_type, &step.module_name, &options, payload)?;
42
43 let mut session_id = None;
44 let mut status = ExploitStatus::Failed;
45 let mut error = None;
46 let job_id = result.job_id.clone();
47
48 if !result.success {
49 error = Some(result.message.clone());
50 } else if let Some(job) = &job_id {
51 let job_status = framework.monitor_job(job, run_timeout)?;
52 if job_status.error.is_some() {
53 error = job_status.error.clone();
54 } else if !job_status.running {
55 status = ExploitStatus::Completed;
56 } else {
57 status = ExploitStatus::Running;
58 }
59 } else {
60 status = ExploitStatus::Completed;
61 }
62
63 let after_sessions = framework.list_sessions().unwrap_or_default();
65 let after_loot = framework.list_loot().unwrap_or_default();
66 for session in &after_sessions {
67 if !baseline_sessions.contains(&session.id) {
68 session_id = Some(session.id.clone());
69 break;
70 }
71 }
72 let loot: Vec<String> = after_loot
73 .iter()
74 .filter(|l| !baseline_loot.contains(&l.id))
75 .map(|l| l.id.clone())
76 .collect();
77
78 Ok(ExploitOutcome {
79 step: step.clone(),
80 status,
81 job_id,
82 session_id,
83 error,
84 loot,
85 start_time,
86 finish_time: Some(now_utc()),
87 })
88}
89
90pub trait ExploitFramework: Send {
96 fn name(&self) -> &'static str;
97 fn workspace(&self) -> &str;
98 fn set_workspace(&mut self, workspace: &str);
99 fn connect(&mut self) -> Result<()>;
102 fn disconnect(&mut self);
103 fn is_connected(&self) -> bool;
104
105 fn module_info(&mut self, module_type: &str, name: &str) -> Result<ModuleInfo>;
109
110 fn run_module(
112 &mut self,
113 module_type: &str,
114 name: &str,
115 options: &BTreeMap<String, String>,
116 payload: Option<&str>,
117 ) -> Result<ModuleRunResult>;
118
119 fn monitor_job(&mut self, job_id: &str, timeout: Duration) -> Result<JobStatus>;
122
123 fn list_sessions(&mut self) -> Result<Vec<SessionInfo>>;
125
126 fn list_loot(&mut self) -> Result<Vec<LootInfo>>;
128}
129
130#[cfg(test)]
131mod tests {
132 use super::*;
133 use crate::db::models::ExploitExecution;
134 use crate::db::Storage;
135
136 struct MockFramework {
138 workspace: String,
139 run_succeeds: bool,
140 job_id: String,
141 job_finishes: bool,
142 preexisting_sessions: Vec<SessionInfo>,
144 preexisting_loot: Vec<LootInfo>,
145 new_sessions: Vec<SessionInfo>,
147 new_loot: Vec<LootInfo>,
148 ran: bool,
149 }
150
151 impl ExploitFramework for MockFramework {
152 fn name(&self) -> &'static str {
153 "mock"
154 }
155 fn workspace(&self) -> &str {
156 &self.workspace
157 }
158 fn set_workspace(&mut self, workspace: &str) {
159 self.workspace = workspace.to_string();
160 }
161 fn connect(&mut self) -> Result<()> {
162 Ok(())
163 }
164 fn disconnect(&mut self) {}
165 fn is_connected(&self) -> bool {
166 true
167 }
168 fn module_info(&mut self, _t: &str, _n: &str) -> Result<ModuleInfo> {
169 Ok(ModuleInfo::default())
170 }
171 fn run_module(
172 &mut self,
173 _t: &str,
174 _n: &str,
175 _o: &BTreeMap<String, String>,
176 _p: Option<&str>,
177 ) -> Result<ModuleRunResult> {
178 self.ran = true;
179 if self.run_succeeds {
180 Ok(ModuleRunResult {
181 job_id: Some(self.job_id.clone()),
182 success: true,
183 message: "started".to_string(),
184 raw: String::new(),
185 })
186 } else {
187 Ok(ModuleRunResult {
188 job_id: None,
189 success: false,
190 message: "boom".to_string(),
191 raw: String::new(),
192 })
193 }
194 }
195 fn monitor_job(&mut self, job_id: &str, _timeout: Duration) -> Result<JobStatus> {
196 Ok(JobStatus {
197 job_id: job_id.to_string(),
198 running: !self.job_finishes,
199 status: if self.job_finishes {
200 "completed".to_string()
201 } else {
202 "still_running".to_string()
203 },
204 error: None,
205 })
206 }
207 fn list_sessions(&mut self) -> Result<Vec<SessionInfo>> {
208 let mut all = self.preexisting_sessions.clone();
209 if self.ran {
210 all.extend(self.new_sessions.clone());
211 }
212 Ok(all)
213 }
214 fn list_loot(&mut self) -> Result<Vec<LootInfo>> {
215 let mut all = self.preexisting_loot.clone();
216 if self.ran {
217 all.extend(self.new_loot.clone());
218 }
219 Ok(all)
220 }
221 }
222
223 fn sample_step() -> ExploitStep {
224 ExploitStep {
225 module_type: MODULE_EXPLOIT.to_string(),
226 module_name: "exploit/multi/http/apache_path_traversal".to_string(),
227 cve: Some("CVE-2021-41773".to_string()),
228 vulnerability: "Apache path traversal".to_string(),
229 rank: "excellent".to_string(),
230 asset_id: "asset-1".to_string(),
231 finding_id: "finding-1".to_string(),
232 host: "10.0.0.5".to_string(),
233 service_port: Some(80),
234 service_protocol: Some("tcp".to_string()),
235 inferred_options: BTreeMap::from([("RHOSTS".to_string(), "10.0.0.5".to_string())]),
236 missing_options: vec!["LHOST".to_string()],
237 payload_compatibility: "generic".to_string(),
238 confidence: "high".to_string(),
239 default_selected: true,
240 }
241 }
242
243 #[test]
244 fn execute_step_records_completed_outcome() {
245 let mut framework = MockFramework {
246 workspace: "default".to_string(),
247 run_succeeds: true,
248 job_id: "7".to_string(),
249 job_finishes: true,
250 preexisting_sessions: Vec::new(),
251 preexisting_loot: Vec::new(),
252 new_sessions: vec![SessionInfo {
253 id: "1".to_string(),
254 typ: "meterpreter".to_string(),
255 target_host: "10.0.0.5".to_string(),
256 platform: "linux".to_string(),
257 info: "test".to_string(),
258 }],
259 new_loot: vec![LootInfo {
260 id: "loot-1".to_string(),
261 name: "shell".to_string(),
262 content_type: "text/plain".to_string(),
263 host: "10.0.0.5".to_string(),
264 path: "/tmp/loot".to_string(),
265 }],
266 ran: false,
267 };
268 let outcome =
269 execute_step(&mut framework, &sample_step(), None, Duration::from_secs(5)).unwrap();
270 assert_eq!(outcome.status, ExploitStatus::Completed);
271 assert_eq!(outcome.job_id.as_deref(), Some("7"));
272 assert_eq!(outcome.session_id.as_deref(), Some("1"));
274 assert!(outcome.loot.iter().any(|l| l == "loot-1"));
275 assert!(outcome.error.is_none());
276 }
277
278 #[test]
279 fn execute_step_does_not_credit_preexisting_sessions() {
280 let mut framework = MockFramework {
283 workspace: "default".to_string(),
284 run_succeeds: true,
285 job_id: "7".to_string(),
286 job_finishes: true,
287 preexisting_sessions: vec![SessionInfo {
288 id: "old-1".to_string(),
289 typ: "shell".to_string(),
290 target_host: "10.0.0.9".to_string(),
291 platform: "linux".to_string(),
292 info: "pre-existing".to_string(),
293 }],
294 preexisting_loot: vec![LootInfo {
295 id: "loot-old".to_string(),
296 name: "shell".to_string(),
297 content_type: "text/plain".to_string(),
298 host: "10.0.0.9".to_string(),
299 path: "/tmp/loot-old".to_string(),
300 }],
301 new_sessions: vec![SessionInfo {
302 id: "new-2".to_string(),
303 typ: "meterpreter".to_string(),
304 target_host: "10.0.0.5".to_string(),
305 platform: "linux".to_string(),
306 info: "from this run".to_string(),
307 }],
308 new_loot: vec![LootInfo {
309 id: "loot-2".to_string(),
310 name: "shell".to_string(),
311 content_type: "text/plain".to_string(),
312 host: "10.0.0.5".to_string(),
313 path: "/tmp/loot-2".to_string(),
314 }],
315 ran: false,
316 };
317 let outcome =
318 execute_step(&mut framework, &sample_step(), None, Duration::from_secs(5)).unwrap();
319 assert_eq!(outcome.session_id.as_deref(), Some("new-2"));
320 assert!(!outcome.loot.iter().any(|l| l == "loot-old"));
321 assert_eq!(outcome.loot, vec!["loot-2".to_string()]);
322 }
323
324 #[test]
325 fn execute_step_marks_failure() {
326 let mut framework = MockFramework {
327 workspace: "default".to_string(),
328 run_succeeds: false,
329 job_id: String::new(),
330 job_finishes: true,
331 preexisting_sessions: Vec::new(),
332 preexisting_loot: Vec::new(),
333 new_sessions: Vec::new(),
334 new_loot: Vec::new(),
335 ran: false,
336 };
337 let outcome =
338 execute_step(&mut framework, &sample_step(), None, Duration::from_secs(5)).unwrap();
339 assert_eq!(outcome.status, ExploitStatus::Failed);
340 assert_eq!(outcome.error.as_deref(), Some("boom"));
341 }
342
343 #[test]
344 fn execution_status_roundtrips() {
345 for status in [
346 ExploitStatus::Planned,
347 ExploitStatus::Running,
348 ExploitStatus::Completed,
349 ExploitStatus::Failed,
350 ExploitStatus::Skipped,
351 ExploitStatus::DryRun,
352 ] {
353 assert_eq!(ExploitStatus::parse(status.as_str()), status);
354 }
355 }
356
357 #[test]
358 fn persist_execution_to_sqlite_roundtrips() {
359 let mut storage = crate::db::sqlite::SqliteStorage::open(":memory:").unwrap();
360 storage.ensure_schema().unwrap();
361
362 let mut exec =
363 ExploitExecution::new("run123", "asset-1", "finding-1", "exploit/multi/http/x");
364 exec.module_type = MODULE_EXPLOIT.to_string();
365 exec.host = "10.0.0.5".to_string();
366 exec.status = ExploitStatus::Completed.as_str().to_string();
367 exec.job_id = Some("9".to_string());
368 exec.session_id = Some("2".to_string());
369 exec.loot = vec!["loot-9".to_string()];
370 exec.selected = true;
371 storage.upsert_exploit_execution(&exec).unwrap();
372
373 let loaded = storage
374 .get_exploit_execution(&exec.execution_id)
375 .unwrap()
376 .unwrap();
377 assert_eq!(loaded.run_id, "run123");
378 assert_eq!(loaded.module, "exploit/multi/http/x");
379 assert_eq!(loaded.status, "completed");
380 assert_eq!(loaded.job_id.as_deref(), Some("9"));
381 assert_eq!(loaded.session_id.as_deref(), Some("2"));
382 assert_eq!(loaded.loot, vec!["loot-9".to_string()]);
383 assert!(loaded.selected);
384
385 let by_run = storage.list_exploit_executions_for_run("run123").unwrap();
386 assert_eq!(by_run.len(), 1);
387 let all = storage.list_exploit_executions().unwrap();
388 assert_eq!(all.len(), 1);
389 }
390}