1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3
4use crate::core::safety::RiskLevel;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
7pub enum ModuleKind {
8 Exploit,
9 Payload,
10 Listener,
11 Post,
12 Auxiliary,
13 Scanner,
14 Backdoor,
15 Encoder,
16 Transform,
17 Analysis,
18}
19
20impl ModuleKind {
21 pub fn as_str(&self) -> &'static str {
22 match self {
23 ModuleKind::Exploit => "exploit",
24 ModuleKind::Payload => "payload",
25 ModuleKind::Listener => "listener",
26 ModuleKind::Post => "post",
27 ModuleKind::Auxiliary => "auxiliary",
28 ModuleKind::Scanner => "scanner",
29 ModuleKind::Backdoor => "backdoor",
30 ModuleKind::Encoder => "encoder",
31 ModuleKind::Transform => "transform",
32 ModuleKind::Analysis => "analysis",
33 }
34 }
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
38#[serde(rename_all = "snake_case")]
39pub enum Capability {
40 NetworkScan,
41 CredentialAccess,
42 PrivilegeEscalation,
43 Persistence,
44 LateralMovement,
45 DataCollection,
46 FilesystemModification,
47 CloudEnumeration,
48}
49
50impl<'de> Deserialize<'de> for Capability {
51 fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
52 let s = String::deserialize(deserializer)?;
53 s.parse::<Capability>().map_err(serde::de::Error::custom)
54 }
55}
56
57impl Capability {
58 pub fn as_str(&self) -> &'static str {
59 match self {
60 Capability::NetworkScan => "network_scan",
61 Capability::CredentialAccess => "credential_access",
62 Capability::PrivilegeEscalation => "privilege_escalation",
63 Capability::Persistence => "persistence",
64 Capability::LateralMovement => "lateral_movement",
65 Capability::DataCollection => "data_collection",
66 Capability::FilesystemModification => "filesystem_modification",
67 Capability::CloudEnumeration => "cloud_enumeration",
68 }
69 }
70
71 pub fn intent(&self) -> Intent {
72 match self {
73 Capability::CredentialAccess => Intent::Dump,
74 Capability::FilesystemModification | Capability::Persistence => Intent::Modify,
75 Capability::LateralMovement | Capability::PrivilegeEscalation => Intent::Execute,
76 _ => Intent::Read,
77 }
78 }
79
80 pub fn impact(&self) -> RiskLevel {
81 match self {
82 Capability::NetworkScan | Capability::DataCollection | Capability::CloudEnumeration => {
83 RiskLevel::Low
84 }
85 Capability::FilesystemModification => RiskLevel::Medium,
86 Capability::PrivilegeEscalation
87 | Capability::Persistence
88 | Capability::LateralMovement => RiskLevel::High,
89 Capability::CredentialAccess => RiskLevel::Critical,
90 }
91 }
92
93 pub fn from_kind(kind: ModuleKind) -> Vec<Capability> {
94 match kind {
95 ModuleKind::Scanner | ModuleKind::Auxiliary | ModuleKind::Analysis => {
96 vec![Capability::NetworkScan]
97 }
98 ModuleKind::Exploit => vec![
99 Capability::PrivilegeEscalation,
100 Capability::CredentialAccess,
101 ],
102 ModuleKind::Post => vec![Capability::PrivilegeEscalation, Capability::DataCollection],
103 ModuleKind::Payload => vec![Capability::Persistence, Capability::LateralMovement],
104 ModuleKind::Listener => vec![Capability::Persistence],
105 ModuleKind::Backdoor => vec![Capability::Persistence, Capability::LateralMovement],
106 ModuleKind::Encoder | ModuleKind::Transform => vec![Capability::DataCollection],
107 }
108 }
109}
110
111impl std::str::FromStr for Capability {
112 type Err = String;
113 fn from_str(s: &str) -> Result<Self, Self::Err> {
114 Ok(match s.to_ascii_lowercase().as_str() {
115 "networkscan" | "network_scan" => Capability::NetworkScan,
116 "credentialaccess" | "credential_access" => Capability::CredentialAccess,
117 "privilegeescalation" | "privilege_escalation" => Capability::PrivilegeEscalation,
118 "persistence" => Capability::Persistence,
119 "lateralmovement" | "lateral_movement" => Capability::LateralMovement,
120 "datacollection" | "data_collection" | "data_exfiltration" | "exfiltration"
121 | "exfil" => Capability::DataCollection,
122 "filesystemmodification" | "filesystem_modification" => {
123 Capability::FilesystemModification
124 }
125 "cloudenumeration" | "cloud_enumeration" => Capability::CloudEnumeration,
126 other => return Err(format!("unknown capability: {other}")),
127 })
128 }
129}
130
131#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
132pub enum Intent {
133 Read,
134 Modify,
135 Execute,
136 Dump,
137}
138
139impl Intent {
140 pub fn as_str(&self) -> &'static str {
141 match self {
142 Intent::Read => "read",
143 Intent::Modify => "modify",
144 Intent::Execute => "execute",
145 Intent::Dump => "dump",
146 }
147 }
148}
149
150impl std::str::FromStr for Intent {
151 type Err = String;
152 fn from_str(s: &str) -> Result<Self, Self::Err> {
153 Ok(match s.to_ascii_lowercase().as_str() {
154 "read" => Intent::Read,
155 "modify" => Intent::Modify,
156 "execute" => Intent::Execute,
157 "dump" => Intent::Dump,
158 other => return Err(format!("unknown intent: {other}")),
159 })
160 }
161}
162
163#[derive(Debug, Clone, Serialize)]
164pub struct ModuleInfo {
165 pub name: String,
166 pub description: String,
167 pub author: String,
168 pub kind: ModuleKind,
169 pub capabilities: Vec<Capability>,
170 pub impact: Option<RiskLevel>,
171 pub intent: Option<Intent>,
172 pub sandbox_image: Option<String>,
173 pub cve: Option<String>,
174}
175
176impl ModuleInfo {
177 pub fn effective_impact(&self) -> RiskLevel {
178 if let Some(i) = self.impact {
179 return i;
180 }
181 let base = RiskLevel::from_kind(self.kind);
182 self.capabilities
183 .iter()
184 .map(|c| c.impact())
185 .max()
186 .unwrap_or(base)
187 .max(base)
188 }
189
190 pub fn effective_intents(&self) -> Vec<Intent> {
191 if let Some(i) = self.intent {
192 return vec![i];
193 }
194 self.capabilities.iter().map(|c| c.intent()).collect()
195 }
196}
197
198#[derive(Debug, Clone, Default, Serialize, Deserialize)]
199pub struct ModuleResult {
200 pub success: bool,
201 pub finding: Option<String>,
202 pub evidence: Vec<String>,
203 pub error: Option<String>,
204 pub session_id: Option<String>,
205 pub data: serde_json::Value,
206}
207
208#[derive(Debug, thiserror::Error)]
209pub enum ModuleError {
210 #[error("missing required option: {0}")]
211 MissingOption(String),
212 #[error("option parse error: {0}")]
213 Parse(String),
214 #[error("{0}")]
215 Other(String),
216}
217
218#[async_trait]
219pub trait Module: Send + Sync {
220 fn options_json(&self) -> serde_json::Value {
221 serde_json::Value::Null
222 }
223
224 async fn run(&self) -> Result<ModuleResult, ModuleError>;
225
226 async fn dry_run(&self) -> Result<ModuleResult, ModuleError> {
229 Err(ModuleError::Other("dry_run not supported".into()))
230 }
231
232 fn set_option(&mut self, _name: &str, _value: &str) -> Result<(), ModuleError> {
233 Err(ModuleError::Other("module does not accept options".into()))
234 }
235
236 fn validate(&self) -> Result<(), ModuleError> {
237 Ok(())
238 }
239}
240
241#[derive(Clone, Copy)]
242pub struct ModuleEntry {
243 pub info: fn() -> ModuleInfo,
244 pub make: fn() -> Box<dyn Module>,
245}
246
247pub struct LoadedModule {
248 pub info: ModuleInfo,
249 pub module: Box<dyn Module>,
250}
251
252#[cfg(test)]
253mod tests {
254 use super::*;
255 use serde_json::json;
256 use std::str::FromStr;
257
258 #[test]
259 fn capability_from_str_accepts_aliases() {
260 assert_eq!(
261 Capability::from_str("data_exfiltration").unwrap(),
262 Capability::DataCollection
263 );
264 assert_eq!(
265 "exfiltration".parse::<Capability>().unwrap(),
266 Capability::DataCollection
267 );
268 assert_eq!(
269 "exfil".parse::<Capability>().unwrap(),
270 Capability::DataCollection
271 );
272 assert_eq!(
273 "DataCollection".parse::<Capability>().unwrap(),
274 Capability::DataCollection
275 );
276 assert_eq!(
277 "data_collection".parse::<Capability>().unwrap(),
278 Capability::DataCollection
279 );
280 assert!("reverse_shell".parse::<Capability>().is_err());
281 }
282
283 #[test]
284 fn capability_deserialize_accepts_aliases() {
285 let c: Capability = serde_json::from_value(json!("data_exfiltration")).unwrap();
286 assert_eq!(c, Capability::DataCollection);
287 let c: Capability = serde_json::from_value(json!("exfil")).unwrap();
288 assert_eq!(c, Capability::DataCollection);
289 let c: Capability = serde_json::from_value(json!("data_collection")).unwrap();
290 assert_eq!(c, Capability::DataCollection);
291 }
292
293 #[test]
294 fn taskspec_deserialize_alias_and_bare_cvss() {
295 use crate::core::sdk::TaskSpec;
296 let t: TaskSpec = serde_json::from_value(json!({
297 "name": "x",
298 "target": "10.0.0.5",
299 "capabilities": ["data_exfiltration"],
300 "impact": "medium",
301 "destructive": false,
302 "cvss": 9.5
303 }))
304 .unwrap();
305 assert_eq!(t.capabilities[0], Capability::DataCollection);
306 assert_eq!(t.cvss.unwrap().effective_score(), 9.5);
307 }
308}