Skip to main content

icebox/core/
module.rs

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    /// Preview the module's output without side effects (no network I/O).
227    /// Used by the executor to enforce DenyPayload rules pre-execution.
228    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}