spice_framework/
rbac.rs

1use crate::agent::AgentConfig;
2use crate::assertion::Assertion;
3use crate::test_case::TestCase;
4use crate::TestCaseBuilder;
5
6/// Role-based access control matrix for test generation.
7pub struct RbacMatrix {
8    roles: Vec<(String, Vec<String>)>,
9}
10
11impl RbacMatrix {
12    pub fn new() -> Self {
13        Self { roles: vec![] }
14    }
15
16    /// Add a role with its allowed tools.
17    pub fn role(mut self, name: &str, tools: &[&str]) -> Self {
18        self.roles.push((
19            name.to_string(),
20            tools.iter().map(|s| s.to_string()).collect(),
21        ));
22        self
23    }
24
25    /// Get allowed tools for a role.
26    pub fn tools_for(&self, role: &str) -> Vec<&str> {
27        self.roles
28            .iter()
29            .find(|(r, _)| r == role)
30            .map(|(_, tools)| tools.iter().map(|s| s.as_str()).collect())
31            .unwrap_or_default()
32    }
33
34    /// Get all unique tools across all roles.
35    pub fn all_tools(&self) -> Vec<&str> {
36        let mut all: Vec<&str> = self
37            .roles
38            .iter()
39            .flat_map(|(_, tools)| tools.iter().map(|s| s.as_str()))
40            .collect();
41        all.sort();
42        all.dedup();
43        all
44    }
45
46    /// Build an AgentConfig for a specific role.
47    pub fn config_for_role(&self, role: &str) -> AgentConfig {
48        AgentConfig::new(serde_json::json!({"role": role}))
49    }
50
51    /// Auto-generate allowlist tests: one test per role verifying tools within allowlist.
52    pub fn generate_allowlist_tests(&self, user_message: &str) -> Vec<TestCase> {
53        let mut tests = Vec::new();
54
55        for (role, role_tools) in &self.roles {
56            let all = self.all_tools();
57            let forbidden: Vec<&str> = all
58                .iter()
59                .filter(|t| !role_tools.iter().any(|rt| rt.as_str() == **t))
60                .copied()
61                .collect();
62
63            let mut builder = TestCaseBuilder::new(
64                format!("rbac-allowlist-{}", role.to_lowercase()),
65                user_message,
66            )
67            .name(format!("{} — tools within allowlist", role))
68            .tags(&["rbac", "security"])
69            .config(self.config_for_role(role))
70            .expect_tools_within_allowlist();
71
72            if !forbidden.is_empty() {
73                builder = builder.forbid_tools(&forbidden);
74            }
75
76            tests.push(builder.build());
77        }
78
79        tests
80    }
81
82    /// Generate scenario tests with per-role custom assertions.
83    /// Takes ownership of assertions since `Assertion` is not Clone.
84    pub fn generate_scenario_tests(
85        &self,
86        id: &str,
87        msg: &str,
88        role_assertions: Vec<(&str, Vec<Assertion>)>,
89    ) -> Vec<TestCase> {
90        let mut tests = Vec::new();
91
92        for (role, assertions) in role_assertions {
93            let mut tc = TestCase {
94                id: format!("rbac-{}-{}", id, role.to_lowercase()),
95                name: Some(format!("{} — {} scenario", role, id)),
96                user_message: msg.to_string(),
97                config: self.config_for_role(role),
98                assertions,
99                tags: vec!["rbac".to_string(), "security".to_string()],
100                retries: 0,
101                consensus_runs: None,
102                consensus_required: None,
103                timeout: None,
104            };
105            // Always add allowlist check
106            tc.assertions.push(Assertion::ExpectToolsWithinAllowlist);
107            tests.push(tc);
108        }
109
110        tests
111    }
112
113    /// Generate injection tests: for each role × payload, verify tools stay within allowlist.
114    pub fn generate_injection_tests(
115        &self,
116        payloads: &[(&str, &str)],
117    ) -> Vec<TestCase> {
118        let mut tests = Vec::new();
119
120        for (role, _) in &self.roles {
121            for (payload_id, payload_msg) in payloads {
122                tests.push(
123                    TestCaseBuilder::new(
124                        format!(
125                            "rbac-injection-{}-{}",
126                            role.to_lowercase(),
127                            payload_id
128                        ),
129                        *payload_msg,
130                    )
131                    .name(format!("{} — injection: {}", role, payload_id))
132                    .tags(&["rbac", "security", "injection"])
133                    .config(self.config_for_role(role))
134                    .expect_tools_within_allowlist()
135                    .build(),
136                );
137            }
138        }
139
140        tests
141    }
142}
143
144impl Default for RbacMatrix {
145    fn default() -> Self {
146        Self::new()
147    }
148}
149
150// --- Convenience constructors for assertions ---
151
152/// Create a ForbidTools assertion.
153pub fn forbid_tools(tools: &[&str]) -> Assertion {
154    Assertion::ForbidTools(tools.iter().map(|s| s.to_string()).collect())
155}
156
157/// Create an ExpectTools assertion.
158pub fn expect_tools(tools: &[&str]) -> Assertion {
159    Assertion::ExpectTools(tools.iter().map(|s| s.to_string()).collect())
160}
spice_framework/rbac.rs

spice_framework/
rbac.rs
