moloch_core/agent/
attestation.rs

1//! Agent attestation types for agent identity verification.
2//!
3//! Agent attestation binds a cryptographic identity to a verifiable agent configuration.
4//! This answers: "What exactly was running when this action was taken?"
5
6use serde::{Deserialize, Serialize};
7use std::time::Duration;
8
9use crate::crypto::{hash, Hash, PublicKey, Sig};
10use crate::error::{Error, Result};
11
12/// Cryptographic attestation of agent state.
13///
14/// An attestation binds an agent's identity (public key) to its verifiable
15/// configuration (code, config, prompt, tools). This enables verification that
16/// an agent was in a known-good state when it took an action.
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct AgentAttestation {
19    /// Agent's public key (identity).
20    agent_id: PublicKey,
21
22    /// Hash of agent's executable code or model weights.
23    code_hash: Hash,
24
25    /// Hash of agent's configuration.
26    config_hash: Hash,
27
28    /// Hash of system prompt / instructions.
29    prompt_hash: Hash,
30
31    /// Available tools at time of attestation.
32    tools: Vec<ToolAttestation>,
33
34    /// Runtime environment attestation.
35    runtime: RuntimeAttestation,
36
37    /// When this attestation was created (Unix timestamp ms).
38    attested_at: i64,
39
40    /// How long this attestation is valid (milliseconds).
41    validity_period_ms: u64,
42
43    /// Signature from attestation authority.
44    authority_signature: Sig,
45
46    /// The authority that signed this attestation.
47    authority: PublicKey,
48}
49
50impl AgentAttestation {
51    /// Create a new attestation builder.
52    pub fn builder() -> AgentAttestationBuilder {
53        AgentAttestationBuilder::new()
54    }
55
56    /// Get the agent's public key.
57    pub fn agent_id(&self) -> &PublicKey {
58        &self.agent_id
59    }
60
61    /// Get the code hash.
62    pub fn code_hash(&self) -> &Hash {
63        &self.code_hash
64    }
65
66    /// Get the config hash.
67    pub fn config_hash(&self) -> &Hash {
68        &self.config_hash
69    }
70
71    /// Get the prompt hash.
72    pub fn prompt_hash(&self) -> &Hash {
73        &self.prompt_hash
74    }
75
76    /// Get the attested tools.
77    pub fn tools(&self) -> &[ToolAttestation] {
78        &self.tools
79    }
80
81    /// Get the runtime attestation.
82    pub fn runtime(&self) -> &RuntimeAttestation {
83        &self.runtime
84    }
85
86    /// Get when the attestation was created.
87    pub fn attested_at(&self) -> i64 {
88        self.attested_at
89    }
90
91    /// Get the validity period.
92    pub fn validity_period(&self) -> Duration {
93        Duration::from_millis(self.validity_period_ms)
94    }
95
96    /// Get the authority signature.
97    pub fn authority_signature(&self) -> &Sig {
98        &self.authority_signature
99    }
100
101    /// Get the attestation authority.
102    pub fn authority(&self) -> &PublicKey {
103        &self.authority
104    }
105
106    /// Check if this attestation is valid at a given time.
107    ///
108    /// Returns true if:
109    /// - `attested_at + validity_period > check_time`
110    pub fn is_valid_at(&self, check_time: i64) -> bool {
111        let expires_at = self
112            .attested_at
113            .saturating_add(self.validity_period_ms as i64);
114        check_time < expires_at
115    }
116
117    /// Get when this attestation expires.
118    pub fn expires_at(&self) -> i64 {
119        self.attested_at
120            .saturating_add(self.validity_period_ms as i64)
121    }
122
123    /// Compute the canonical bytes for signing/verification.
124    pub fn canonical_bytes(&self) -> Vec<u8> {
125        // Include all fields except authority_signature
126        let mut data = Vec::new();
127        data.extend_from_slice(&self.agent_id.as_bytes());
128        data.extend_from_slice(self.code_hash.as_bytes());
129        data.extend_from_slice(self.config_hash.as_bytes());
130        data.extend_from_slice(self.prompt_hash.as_bytes());
131        for tool in &self.tools {
132            data.extend_from_slice(tool.tool_id.as_bytes());
133            data.extend_from_slice(tool.version.as_bytes());
134            data.extend_from_slice(tool.implementation_hash.as_bytes());
135        }
136        data.extend_from_slice(self.runtime.runtime_id.as_bytes());
137        data.extend_from_slice(self.runtime.runtime_hash.as_bytes());
138        data.extend_from_slice(&self.attested_at.to_le_bytes());
139        data.extend_from_slice(&self.validity_period_ms.to_le_bytes());
140        data.extend_from_slice(&self.authority.as_bytes());
141        data
142    }
143
144    /// Compute the hash of this attestation.
145    pub fn hash(&self) -> Hash {
146        hash(&self.canonical_bytes())
147    }
148
149    /// Verify the authority signature.
150    pub fn verify_signature(&self) -> Result<()> {
151        let bytes = self.canonical_bytes();
152        self.authority
153            .verify(&bytes, &self.authority_signature)
154            .map_err(|_| Error::invalid_input("Attestation signature verification failed"))
155    }
156
157    /// Check if a tool is included in this attestation.
158    pub fn has_tool(&self, tool_id: &str) -> bool {
159        self.tools.iter().any(|t| t.tool_id == tool_id)
160    }
161
162    /// Get a specific tool attestation.
163    pub fn get_tool(&self, tool_id: &str) -> Option<&ToolAttestation> {
164        self.tools.iter().find(|t| t.tool_id == tool_id)
165    }
166
167    /// Validate attestation binding: the agent_id in the attestation must
168    /// match the actor's public key (G-4.2, Rule 4.3.3).
169    pub fn validate_binding(&self, actor_key: &PublicKey) -> Result<()> {
170        if self.agent_id != *actor_key {
171            return Err(Error::invalid_input(
172                "attestation agent_id does not match event actor",
173            ));
174        }
175        Ok(())
176    }
177
178    /// Validate tool consistency: the invoked tool must be listed in the
179    /// agent's attestation (G-4.3, Rule 4.3.4, INV-ATTEST-3).
180    pub fn validate_tool(&self, tool_id: &str) -> Result<()> {
181        if !self.has_tool(tool_id) {
182            return Err(Error::invalid_input(format!(
183                "tool '{}' not found in agent attestation",
184                tool_id
185            )));
186        }
187        Ok(())
188    }
189
190    /// Full attestation validation at action time (Rule 4.3.2).
191    ///
192    /// Checks:
193    /// - Signature validity
194    /// - Temporal validity (not expired)
195    /// - Binding to actor
196    pub fn validate_for_action(&self, actor_key: &PublicKey, action_time: i64) -> Result<()> {
197        // Check validity window
198        if !self.is_valid_at(action_time) {
199            return Err(Error::invalid_input("attestation has expired"));
200        }
201
202        // Check signature
203        self.verify_signature()?;
204
205        // Check binding
206        self.validate_binding(actor_key)?;
207
208        Ok(())
209    }
210}
211
212/// Builder for AgentAttestation.
213#[derive(Debug, Default)]
214pub struct AgentAttestationBuilder {
215    agent_id: Option<PublicKey>,
216    code_hash: Option<Hash>,
217    config_hash: Option<Hash>,
218    prompt_hash: Option<Hash>,
219    tools: Vec<ToolAttestation>,
220    runtime: Option<RuntimeAttestation>,
221    attested_at: Option<i64>,
222    validity_period_ms: Option<u64>,
223    authority: Option<PublicKey>,
224}
225
226impl AgentAttestationBuilder {
227    /// Create a new builder.
228    pub fn new() -> Self {
229        Self::default()
230    }
231
232    /// Set the agent ID.
233    pub fn agent_id(mut self, agent_id: PublicKey) -> Self {
234        self.agent_id = Some(agent_id);
235        self
236    }
237
238    /// Set the code hash.
239    pub fn code_hash(mut self, hash: Hash) -> Self {
240        self.code_hash = Some(hash);
241        self
242    }
243
244    /// Set the config hash.
245    pub fn config_hash(mut self, hash: Hash) -> Self {
246        self.config_hash = Some(hash);
247        self
248    }
249
250    /// Set the prompt hash.
251    pub fn prompt_hash(mut self, hash: Hash) -> Self {
252        self.prompt_hash = Some(hash);
253        self
254    }
255
256    /// Add a tool attestation.
257    pub fn tool(mut self, tool: ToolAttestation) -> Self {
258        self.tools.push(tool);
259        self
260    }
261
262    /// Set all tools.
263    pub fn tools(mut self, tools: Vec<ToolAttestation>) -> Self {
264        self.tools = tools;
265        self
266    }
267
268    /// Set the runtime attestation.
269    pub fn runtime(mut self, runtime: RuntimeAttestation) -> Self {
270        self.runtime = Some(runtime);
271        self
272    }
273
274    /// Set the attestation time.
275    pub fn attested_at(mut self, timestamp: i64) -> Self {
276        self.attested_at = Some(timestamp);
277        self
278    }
279
280    /// Set the validity period.
281    pub fn validity_period(mut self, duration: Duration) -> Self {
282        self.validity_period_ms = Some(duration.as_millis() as u64);
283        self
284    }
285
286    /// Set the authority.
287    pub fn authority(mut self, authority: PublicKey) -> Self {
288        self.authority = Some(authority);
289        self
290    }
291
292    /// Sign and build the attestation.
293    ///
294    /// # Arguments
295    /// * `authority_key` - The authority's secret key for signing
296    ///
297    /// # Errors
298    /// Returns error if required fields are missing.
299    pub fn sign(self, authority_key: &crate::crypto::SecretKey) -> Result<AgentAttestation> {
300        let agent_id = self
301            .agent_id
302            .ok_or_else(|| Error::invalid_input("agent_id is required"))?;
303
304        let code_hash = self
305            .code_hash
306            .ok_or_else(|| Error::invalid_input("code_hash is required"))?;
307
308        // Validate code_hash is not empty (all zeros)
309        if code_hash.as_bytes().iter().all(|&b| b == 0) {
310            return Err(Error::invalid_input("code_hash cannot be empty"));
311        }
312
313        let config_hash = self
314            .config_hash
315            .ok_or_else(|| Error::invalid_input("config_hash is required"))?;
316
317        let prompt_hash = self
318            .prompt_hash
319            .ok_or_else(|| Error::invalid_input("prompt_hash is required"))?;
320
321        let runtime = self
322            .runtime
323            .ok_or_else(|| Error::invalid_input("runtime is required"))?;
324
325        let attested_at = self
326            .attested_at
327            .unwrap_or_else(|| chrono::Utc::now().timestamp_millis());
328
329        let validity_period_ms = self.validity_period_ms.unwrap_or(3600 * 1000); // 1 hour default
330
331        if validity_period_ms == 0 {
332            return Err(Error::invalid_input("validity_period must be positive"));
333        }
334
335        let authority = self.authority.unwrap_or_else(|| authority_key.public_key());
336
337        // Create attestation without signature first
338        let mut attestation = AgentAttestation {
339            agent_id,
340            code_hash,
341            config_hash,
342            prompt_hash,
343            tools: self.tools,
344            runtime,
345            attested_at,
346            validity_period_ms,
347            authority_signature: Sig::empty(), // Placeholder
348            authority,
349        };
350
351        // Sign the canonical bytes
352        let canonical = attestation.canonical_bytes();
353        attestation.authority_signature = authority_key.sign(&canonical);
354
355        Ok(attestation)
356    }
357}
358
359/// A required capability for a tool.
360///
361/// This is a simplified representation of capability requirements that a tool
362/// needs to function. It matches the base types from `CapabilityKind`.
363#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
364#[serde(tag = "type", rename_all = "snake_case")]
365pub enum RequiredCapability {
366    /// Read data from resources.
367    Read,
368    /// Write/modify data in resources.
369    Write,
370    /// Delete data from resources.
371    Delete,
372    /// Execute commands or code.
373    Execute,
374    /// Invoke a specific tool.
375    InvokeTool { tool_id: String },
376    /// Spawn child agents.
377    SpawnAgent,
378    /// Delegate capabilities to other agents.
379    DelegateCapability,
380    /// Send messages on a channel.
381    SendMessage { channel: String },
382    /// Receive messages from a channel.
383    ReceiveMessage { channel: String },
384    /// Spend currency (any amount).
385    Spend { currency: String },
386    /// Modify permissions.
387    ModifyPermissions,
388    /// View audit logs.
389    ViewAuditLog,
390    /// File system access.
391    FileSystem,
392    /// Network access.
393    Network,
394}
395
396impl RequiredCapability {
397    /// Create a read capability requirement.
398    pub fn read() -> Self {
399        Self::Read
400    }
401
402    /// Create a write capability requirement.
403    pub fn write() -> Self {
404        Self::Write
405    }
406
407    /// Create an execute capability requirement.
408    pub fn execute() -> Self {
409        Self::Execute
410    }
411
412    /// Create a tool invocation requirement.
413    pub fn invoke_tool(tool_id: impl Into<String>) -> Self {
414        Self::InvokeTool {
415            tool_id: tool_id.into(),
416        }
417    }
418
419    /// Create a file system access requirement.
420    pub fn file_system() -> Self {
421        Self::FileSystem
422    }
423
424    /// Create a network access requirement.
425    pub fn network() -> Self {
426        Self::Network
427    }
428}
429
430impl std::fmt::Display for RequiredCapability {
431    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
432        match self {
433            RequiredCapability::Read => write!(f, "read"),
434            RequiredCapability::Write => write!(f, "write"),
435            RequiredCapability::Delete => write!(f, "delete"),
436            RequiredCapability::Execute => write!(f, "execute"),
437            RequiredCapability::InvokeTool { tool_id } => write!(f, "invoke_tool:{}", tool_id),
438            RequiredCapability::SpawnAgent => write!(f, "spawn_agent"),
439            RequiredCapability::DelegateCapability => write!(f, "delegate_capability"),
440            RequiredCapability::SendMessage { channel } => write!(f, "send_message:{}", channel),
441            RequiredCapability::ReceiveMessage { channel } => {
442                write!(f, "receive_message:{}", channel)
443            }
444            RequiredCapability::Spend { currency } => write!(f, "spend:{}", currency),
445            RequiredCapability::ModifyPermissions => write!(f, "modify_permissions"),
446            RequiredCapability::ViewAuditLog => write!(f, "view_audit_log"),
447            RequiredCapability::FileSystem => write!(f, "file_system"),
448            RequiredCapability::Network => write!(f, "network"),
449        }
450    }
451}
452
453/// Attestation of a specific tool.
454#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
455pub struct ToolAttestation {
456    /// Tool identifier.
457    pub tool_id: String,
458
459    /// Tool version.
460    pub version: String,
461
462    /// Hash of tool implementation.
463    pub implementation_hash: Hash,
464
465    /// Tool's capability requirements.
466    pub required_capabilities: Vec<RequiredCapability>,
467}
468
469impl ToolAttestation {
470    /// Create a new tool attestation.
471    pub fn new(
472        tool_id: impl Into<String>,
473        version: impl Into<String>,
474        implementation_hash: Hash,
475    ) -> Self {
476        Self {
477            tool_id: tool_id.into(),
478            version: version.into(),
479            implementation_hash,
480            required_capabilities: Vec::new(),
481        }
482    }
483
484    /// Add a required capability.
485    pub fn with_capability(mut self, capability: RequiredCapability) -> Self {
486        self.required_capabilities.push(capability);
487        self
488    }
489
490    /// Add multiple required capabilities.
491    pub fn with_capabilities(mut self, capabilities: Vec<RequiredCapability>) -> Self {
492        self.required_capabilities = capabilities;
493        self
494    }
495
496    /// Check if this tool requires a specific capability.
497    pub fn requires(&self, capability: &RequiredCapability) -> bool {
498        self.required_capabilities.contains(capability)
499    }
500}
501
502/// Runtime environment attestation.
503#[derive(Debug, Clone, Serialize, Deserialize)]
504pub struct RuntimeAttestation {
505    /// Runtime identifier (e.g., "claude-code-v1.2.3").
506    pub runtime_id: String,
507
508    /// Hash of runtime binary.
509    pub runtime_hash: Hash,
510
511    /// TEE attestation if available.
512    pub tee_quote: Option<TeeQuote>,
513
514    /// Platform measurements.
515    pub platform_hash: Option<Hash>,
516}
517
518impl RuntimeAttestation {
519    /// Create a new runtime attestation.
520    pub fn new(runtime_id: impl Into<String>, runtime_hash: Hash) -> Self {
521        Self {
522            runtime_id: runtime_id.into(),
523            runtime_hash,
524            tee_quote: None,
525            platform_hash: None,
526        }
527    }
528
529    /// Add a TEE quote.
530    pub fn with_tee(mut self, quote: TeeQuote) -> Self {
531        self.tee_quote = Some(quote);
532        self
533    }
534
535    /// Add a platform hash.
536    pub fn with_platform_hash(mut self, hash: Hash) -> Self {
537        self.platform_hash = Some(hash);
538        self
539    }
540}
541
542/// Trusted Execution Environment quote.
543#[derive(Debug, Clone, Serialize, Deserialize)]
544pub struct TeeQuote {
545    /// TEE type.
546    pub tee_type: TeeType,
547
548    /// Raw attestation quote.
549    pub quote: Vec<u8>,
550
551    /// Measurement registers.
552    pub measurements: Vec<Hash>,
553}
554
555impl TeeQuote {
556    /// Create a new TEE quote.
557    pub fn new(tee_type: TeeType, quote: Vec<u8>) -> Self {
558        Self {
559            tee_type,
560            quote,
561            measurements: Vec::new(),
562        }
563    }
564
565    /// Add a measurement.
566    pub fn with_measurement(mut self, hash: Hash) -> Self {
567        self.measurements.push(hash);
568        self
569    }
570}
571
572/// Type of Trusted Execution Environment.
573#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
574#[serde(rename_all = "snake_case")]
575pub enum TeeType {
576    /// Intel SGX.
577    IntelSgx,
578    /// Intel TDX.
579    IntelTdx,
580    /// AMD SEV-SNP.
581    AmdSevSnp,
582    /// ARM CCA.
583    ArmCca,
584    /// Software-based (for testing only).
585    Software,
586}
587
588/// Error types specific to attestation operations.
589#[derive(Debug, Clone, PartialEq, Eq)]
590pub enum AttestationError {
591    /// Attestation not found for agent.
592    NotFound,
593    /// Attestation has expired.
594    Expired,
595    /// Attestation has been revoked.
596    Revoked,
597    /// Attestation authority is not trusted.
598    UntrustedAuthority,
599    /// Signature verification failed.
600    InvalidSignature,
601    /// Agent ID mismatch.
602    AgentMismatch,
603    /// Tool not in attestation.
604    ToolNotAttested(String),
605}
606
607impl std::fmt::Display for AttestationError {
608    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
609        match self {
610            AttestationError::NotFound => write!(f, "Attestation not found"),
611            AttestationError::Expired => write!(f, "Attestation has expired"),
612            AttestationError::Revoked => write!(f, "Attestation has been revoked"),
613            AttestationError::UntrustedAuthority => write!(f, "Attestation authority not trusted"),
614            AttestationError::InvalidSignature => write!(f, "Attestation signature invalid"),
615            AttestationError::AgentMismatch => write!(f, "Agent ID does not match attestation"),
616            AttestationError::ToolNotAttested(tool) => {
617                write!(f, "Tool '{}' not in agent's attestation", tool)
618            }
619        }
620    }
621}
622
623impl std::error::Error for AttestationError {}
624
625#[cfg(test)]
626mod tests {
627    use super::*;
628    use crate::crypto::{hash, SecretKey};
629
630    fn test_runtime() -> RuntimeAttestation {
631        RuntimeAttestation::new("test-runtime-v1.0.0", hash(b"runtime-binary"))
632    }
633
634    fn test_tool() -> ToolAttestation {
635        ToolAttestation::new("read_file", "1.0.0", hash(b"read-file-impl"))
636    }
637
638    // === Construction Tests ===
639
640    #[test]
641    fn attestation_requires_code_hash() {
642        let authority = SecretKey::generate();
643        let agent = SecretKey::generate();
644
645        let result = AgentAttestation::builder()
646            .agent_id(agent.public_key())
647            // Missing code_hash
648            .config_hash(hash(b"config"))
649            .prompt_hash(hash(b"prompt"))
650            .runtime(test_runtime())
651            .sign(&authority);
652
653        assert!(result.is_err());
654    }
655
656    #[test]
657    fn attestation_rejects_empty_code_hash() {
658        let authority = SecretKey::generate();
659        let agent = SecretKey::generate();
660
661        let result = AgentAttestation::builder()
662            .agent_id(agent.public_key())
663            .code_hash(Hash::from_bytes([0u8; 32])) // Empty hash
664            .config_hash(hash(b"config"))
665            .prompt_hash(hash(b"prompt"))
666            .runtime(test_runtime())
667            .sign(&authority);
668
669        assert!(result.is_err());
670        assert!(result.unwrap_err().to_string().contains("empty"));
671    }
672
673    #[test]
674    fn attestation_requires_valid_signature() {
675        let authority = SecretKey::generate();
676        let agent = SecretKey::generate();
677
678        let attestation = AgentAttestation::builder()
679            .agent_id(agent.public_key())
680            .code_hash(hash(b"code"))
681            .config_hash(hash(b"config"))
682            .prompt_hash(hash(b"prompt"))
683            .runtime(test_runtime())
684            .sign(&authority)
685            .unwrap();
686
687        // Verify with correct authority
688        assert!(attestation.verify_signature().is_ok());
689
690        // Tamper with attestation and verify fails
691        let mut tampered = attestation.clone();
692        tampered.attested_at += 1;
693        assert!(tampered.verify_signature().is_err());
694    }
695
696    #[test]
697    fn attestation_validity_period_must_be_positive() {
698        let authority = SecretKey::generate();
699        let agent = SecretKey::generate();
700
701        let result = AgentAttestation::builder()
702            .agent_id(agent.public_key())
703            .code_hash(hash(b"code"))
704            .config_hash(hash(b"config"))
705            .prompt_hash(hash(b"prompt"))
706            .runtime(test_runtime())
707            .validity_period(Duration::from_secs(0)) // Zero duration
708            .sign(&authority);
709
710        assert!(result.is_err());
711    }
712
713    // === Validity Tests ===
714
715    #[test]
716    fn attestation_valid_within_validity_period() {
717        let authority = SecretKey::generate();
718        let agent = SecretKey::generate();
719        let start = 1000000i64;
720
721        let attestation = AgentAttestation::builder()
722            .agent_id(agent.public_key())
723            .code_hash(hash(b"code"))
724            .config_hash(hash(b"config"))
725            .prompt_hash(hash(b"prompt"))
726            .runtime(test_runtime())
727            .attested_at(start)
728            .validity_period(Duration::from_secs(3600)) // 1 hour
729            .sign(&authority)
730            .unwrap();
731
732        // Valid at start
733        assert!(attestation.is_valid_at(start));
734
735        // Valid 30 minutes later
736        assert!(attestation.is_valid_at(start + 1800 * 1000));
737
738        // Valid 59 minutes later
739        assert!(attestation.is_valid_at(start + 3540 * 1000));
740    }
741
742    #[test]
743    fn attestation_invalid_after_expiry() {
744        let authority = SecretKey::generate();
745        let agent = SecretKey::generate();
746        let start = 1000000i64;
747
748        let attestation = AgentAttestation::builder()
749            .agent_id(agent.public_key())
750            .code_hash(hash(b"code"))
751            .config_hash(hash(b"config"))
752            .prompt_hash(hash(b"prompt"))
753            .runtime(test_runtime())
754            .attested_at(start)
755            .validity_period(Duration::from_secs(3600)) // 1 hour
756            .sign(&authority)
757            .unwrap();
758
759        // Invalid at exactly expiry
760        assert!(!attestation.is_valid_at(start + 3600 * 1000));
761
762        // Invalid after expiry
763        assert!(!attestation.is_valid_at(start + 3601 * 1000));
764    }
765
766    // === Tool Attestation Tests ===
767
768    #[test]
769    fn tool_attestation_includes_version() {
770        let tool = ToolAttestation::new("bash", "5.1.0", hash(b"bash-impl"));
771        assert_eq!(tool.version, "5.1.0");
772    }
773
774    #[test]
775    fn tool_attestation_includes_implementation_hash() {
776        let impl_hash = hash(b"tool-implementation");
777        let tool = ToolAttestation::new("read_file", "1.0.0", impl_hash);
778        assert_eq!(tool.implementation_hash, impl_hash);
779    }
780
781    #[test]
782    fn tool_attestation_with_capabilities() {
783        let tool = ToolAttestation::new("bash", "5.1.0", hash(b"bash-impl"))
784            .with_capability(RequiredCapability::Execute)
785            .with_capability(RequiredCapability::FileSystem);
786
787        assert_eq!(tool.required_capabilities.len(), 2);
788        assert!(tool.requires(&RequiredCapability::Execute));
789        assert!(tool.requires(&RequiredCapability::FileSystem));
790        assert!(!tool.requires(&RequiredCapability::Network));
791    }
792
793    #[test]
794    fn attestation_has_tool_check() {
795        let authority = SecretKey::generate();
796        let agent = SecretKey::generate();
797
798        let attestation = AgentAttestation::builder()
799            .agent_id(agent.public_key())
800            .code_hash(hash(b"code"))
801            .config_hash(hash(b"config"))
802            .prompt_hash(hash(b"prompt"))
803            .runtime(test_runtime())
804            .tool(test_tool())
805            .tool(ToolAttestation::new("bash", "5.1", hash(b"bash")))
806            .sign(&authority)
807            .unwrap();
808
809        assert!(attestation.has_tool("read_file"));
810        assert!(attestation.has_tool("bash"));
811        assert!(!attestation.has_tool("write_file"));
812    }
813
814    // === Runtime Attestation Tests ===
815
816    #[test]
817    fn runtime_attestation_with_tee() {
818        let runtime = RuntimeAttestation::new("claude-v1", hash(b"runtime"))
819            .with_tee(TeeQuote::new(TeeType::IntelSgx, vec![1, 2, 3, 4]));
820
821        assert!(runtime.tee_quote.is_some());
822        assert_eq!(runtime.tee_quote.unwrap().tee_type, TeeType::IntelSgx);
823    }
824
825    #[test]
826    fn runtime_attestation_with_platform_hash() {
827        let platform = hash(b"platform-measurements");
828        let runtime =
829            RuntimeAttestation::new("claude-v1", hash(b"runtime")).with_platform_hash(platform);
830
831        assert_eq!(runtime.platform_hash, Some(platform));
832    }
833
834    // === Hash Tests ===
835
836    #[test]
837    fn attestation_hash_deterministic() {
838        let authority = SecretKey::generate();
839        let agent = SecretKey::generate();
840
841        let attestation = AgentAttestation::builder()
842            .agent_id(agent.public_key())
843            .code_hash(hash(b"code"))
844            .config_hash(hash(b"config"))
845            .prompt_hash(hash(b"prompt"))
846            .runtime(test_runtime())
847            .attested_at(1000000)
848            .sign(&authority)
849            .unwrap();
850
851        let h1 = attestation.hash();
852        let h2 = attestation.hash();
853        assert_eq!(h1, h2);
854    }
855
856    #[test]
857    fn different_attestations_different_hash() {
858        let authority = SecretKey::generate();
859        let agent = SecretKey::generate();
860
861        let attestation1 = AgentAttestation::builder()
862            .agent_id(agent.public_key())
863            .code_hash(hash(b"code1"))
864            .config_hash(hash(b"config"))
865            .prompt_hash(hash(b"prompt"))
866            .runtime(test_runtime())
867            .attested_at(1000000)
868            .sign(&authority)
869            .unwrap();
870
871        let attestation2 = AgentAttestation::builder()
872            .agent_id(agent.public_key())
873            .code_hash(hash(b"code2")) // Different code
874            .config_hash(hash(b"config"))
875            .prompt_hash(hash(b"prompt"))
876            .runtime(test_runtime())
877            .attested_at(1000000)
878            .sign(&authority)
879            .unwrap();
880
881        assert_ne!(attestation1.hash(), attestation2.hash());
882    }
883
884    // === Getter Tests ===
885
886    #[test]
887    fn attestation_getters_work() {
888        let authority = SecretKey::generate();
889        let agent = SecretKey::generate();
890        let code = hash(b"code");
891        let config = hash(b"config");
892        let prompt = hash(b"prompt");
893        let runtime = test_runtime();
894
895        let attestation = AgentAttestation::builder()
896            .agent_id(agent.public_key())
897            .code_hash(code)
898            .config_hash(config)
899            .prompt_hash(prompt)
900            .runtime(runtime.clone())
901            .attested_at(1000000)
902            .validity_period(Duration::from_secs(7200))
903            .sign(&authority)
904            .unwrap();
905
906        assert_eq!(attestation.agent_id(), &agent.public_key());
907        assert_eq!(attestation.code_hash(), &code);
908        assert_eq!(attestation.config_hash(), &config);
909        assert_eq!(attestation.prompt_hash(), &prompt);
910        assert_eq!(attestation.runtime().runtime_id, runtime.runtime_id);
911        assert_eq!(attestation.attested_at(), 1000000);
912        assert_eq!(attestation.validity_period(), Duration::from_secs(7200));
913        assert_eq!(attestation.authority(), &authority.public_key());
914        assert_eq!(attestation.expires_at(), 1000000 + 7200 * 1000);
915    }
916
917    // === Attestation Binding Tests (G-4.2) ===
918
919    #[test]
920    fn validate_binding_passes_for_matching_agent() {
921        let authority = SecretKey::generate();
922        let agent = SecretKey::generate();
923
924        let attestation = AgentAttestation::builder()
925            .agent_id(agent.public_key())
926            .code_hash(hash(b"code"))
927            .config_hash(hash(b"config"))
928            .prompt_hash(hash(b"prompt"))
929            .runtime(test_runtime())
930            .sign(&authority)
931            .unwrap();
932
933        assert!(attestation.validate_binding(&agent.public_key()).is_ok());
934    }
935
936    #[test]
937    fn validate_binding_fails_for_different_agent() {
938        let authority = SecretKey::generate();
939        let agent = SecretKey::generate();
940        let other = SecretKey::generate();
941
942        let attestation = AgentAttestation::builder()
943            .agent_id(agent.public_key())
944            .code_hash(hash(b"code"))
945            .config_hash(hash(b"config"))
946            .prompt_hash(hash(b"prompt"))
947            .runtime(test_runtime())
948            .sign(&authority)
949            .unwrap();
950
951        let result = attestation.validate_binding(&other.public_key());
952        assert!(result.is_err());
953        assert!(result.unwrap_err().to_string().contains("does not match"));
954    }
955
956    // === Tool Consistency Tests (G-4.3) ===
957
958    #[test]
959    fn validate_tool_passes_for_attested_tool() {
960        let authority = SecretKey::generate();
961        let agent = SecretKey::generate();
962
963        let attestation = AgentAttestation::builder()
964            .agent_id(agent.public_key())
965            .code_hash(hash(b"code"))
966            .config_hash(hash(b"config"))
967            .prompt_hash(hash(b"prompt"))
968            .runtime(test_runtime())
969            .tool(test_tool())
970            .sign(&authority)
971            .unwrap();
972
973        assert!(attestation.validate_tool("read_file").is_ok());
974    }
975
976    #[test]
977    fn validate_tool_fails_for_unattested_tool() {
978        let authority = SecretKey::generate();
979        let agent = SecretKey::generate();
980
981        let attestation = AgentAttestation::builder()
982            .agent_id(agent.public_key())
983            .code_hash(hash(b"code"))
984            .config_hash(hash(b"config"))
985            .prompt_hash(hash(b"prompt"))
986            .runtime(test_runtime())
987            .tool(test_tool())
988            .sign(&authority)
989            .unwrap();
990
991        let result = attestation.validate_tool("execute_code");
992        assert!(result.is_err());
993        assert!(result.unwrap_err().to_string().contains("not found"));
994    }
995
996    // === Full Action Validation Tests ===
997
998    #[test]
999    fn validate_for_action_passes_when_valid() {
1000        let authority = SecretKey::generate();
1001        let agent = SecretKey::generate();
1002
1003        let attestation = AgentAttestation::builder()
1004            .agent_id(agent.public_key())
1005            .code_hash(hash(b"code"))
1006            .config_hash(hash(b"config"))
1007            .prompt_hash(hash(b"prompt"))
1008            .runtime(test_runtime())
1009            .attested_at(1000)
1010            .validity_period(Duration::from_secs(3600))
1011            .sign(&authority)
1012            .unwrap();
1013
1014        // Within validity window
1015        assert!(attestation
1016            .validate_for_action(&agent.public_key(), 2000)
1017            .is_ok());
1018    }
1019
1020    #[test]
1021    fn validate_for_action_fails_when_expired() {
1022        let authority = SecretKey::generate();
1023        let agent = SecretKey::generate();
1024
1025        let attestation = AgentAttestation::builder()
1026            .agent_id(agent.public_key())
1027            .code_hash(hash(b"code"))
1028            .config_hash(hash(b"config"))
1029            .prompt_hash(hash(b"prompt"))
1030            .runtime(test_runtime())
1031            .attested_at(1000)
1032            .validity_period(Duration::from_secs(1))
1033            .sign(&authority)
1034            .unwrap();
1035
1036        // Past validity window (1000ms + 1000ms = 2000, checking at 3000)
1037        let result = attestation.validate_for_action(&agent.public_key(), 3000);
1038        assert!(result.is_err());
1039    }
1040
1041    #[test]
1042    fn validate_for_action_fails_for_wrong_agent() {
1043        let authority = SecretKey::generate();
1044        let agent = SecretKey::generate();
1045        let other = SecretKey::generate();
1046
1047        let attestation = AgentAttestation::builder()
1048            .agent_id(agent.public_key())
1049            .code_hash(hash(b"code"))
1050            .config_hash(hash(b"config"))
1051            .prompt_hash(hash(b"prompt"))
1052            .runtime(test_runtime())
1053            .attested_at(1000)
1054            .validity_period(Duration::from_secs(3600))
1055            .sign(&authority)
1056            .unwrap();
1057
1058        let result = attestation.validate_for_action(&other.public_key(), 2000);
1059        assert!(result.is_err());
1060    }
1061}
moloch_core/agent/attestation.rs

moloch_core/agent/
attestation.rs
