Skip to main content

peat_protocol/cot/
peat_extension.rs

1//! Peat custom detail extension schema for CoT
2//!
3//! Implements the `<_peat_>` XML extension defined in ADR-028 for preserving
4//! Peat-specific semantics in CoT messages.
5
6use quick_xml::events::{BytesEnd, BytesStart, BytesText, Event};
7use quick_xml::Writer;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::io::Cursor;
11
12use super::event::CotError;
13use super::types::OperationalStatus;
14
15/// Peat version for the extension schema
16pub const PEAT_EXTENSION_VERSION: &str = "1.0";
17
18/// Peat custom detail extension
19///
20/// Contains Peat-specific metadata that doesn't map to standard CoT fields.
21#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
22pub struct PeatExtension {
23    /// Source attribution
24    pub source: Option<PeatSource>,
25    /// Confidence information
26    pub confidence: Option<PeatConfidence>,
27    /// Hierarchy membership
28    pub hierarchy: Option<PeatHierarchy>,
29    /// Custom attributes
30    pub attributes: HashMap<String, PeatAttribute>,
31    /// Operational status
32    pub status: Option<PeatStatus>,
33    /// Capability information (for capability advertisements)
34    pub capabilities: Vec<PeatCapability>,
35    /// Handoff information (for handoff messages)
36    pub handoff: Option<PeatHandoff>,
37    /// Classification information
38    pub classification: Option<PeatClassification>,
39}
40
41impl PeatExtension {
42    /// Create a new empty Peat extension
43    pub fn new() -> Self {
44        Self::default()
45    }
46
47    /// Set source attribution
48    pub fn with_source(mut self, source: PeatSource) -> Self {
49        self.source = Some(source);
50        self
51    }
52
53    /// Set confidence
54    pub fn with_confidence(mut self, value: f64, threshold: Option<f64>) -> Self {
55        self.confidence = Some(PeatConfidence { value, threshold });
56        self
57    }
58
59    /// Set hierarchy membership
60    pub fn with_hierarchy(mut self, hierarchy: PeatHierarchy) -> Self {
61        self.hierarchy = Some(hierarchy);
62        self
63    }
64
65    /// Add an attribute
66    pub fn with_attribute(mut self, key: &str, value: &str, attr_type: &str) -> Self {
67        self.attributes.insert(
68            key.to_string(),
69            PeatAttribute {
70                value: value.to_string(),
71                attr_type: attr_type.to_string(),
72            },
73        );
74        self
75    }
76
77    /// Set status
78    pub fn with_status(mut self, status: PeatStatus) -> Self {
79        self.status = Some(status);
80        self
81    }
82
83    /// Add a capability
84    pub fn with_capability(mut self, capability: PeatCapability) -> Self {
85        self.capabilities.push(capability);
86        self
87    }
88
89    /// Set handoff information
90    pub fn with_handoff(mut self, handoff: PeatHandoff) -> Self {
91        self.handoff = Some(handoff);
92        self
93    }
94
95    /// Set classification
96    pub fn with_classification(mut self, level: &str, caveat: Option<&str>) -> Self {
97        self.classification = Some(PeatClassification {
98            level: level.to_string(),
99            caveat: caveat.map(|s| s.to_string()),
100        });
101        self
102    }
103
104    /// Write the extension as XML
105    pub fn write_xml(&self, writer: &mut Writer<Cursor<Vec<u8>>>) -> Result<(), CotError> {
106        let mut peat_elem = BytesStart::new("_peat_");
107        peat_elem.push_attribute(("version", PEAT_EXTENSION_VERSION));
108
109        writer
110            .write_event(Event::Start(peat_elem))
111            .map_err(|e| CotError::XmlWrite(e.to_string()))?;
112
113        // Source
114        if let Some(ref source) = self.source {
115            let mut src_elem = BytesStart::new("source");
116            src_elem.push_attribute(("platform", source.platform.as_str()));
117            src_elem.push_attribute(("model", source.model.as_str()));
118            src_elem.push_attribute(("model_version", source.model_version.as_str()));
119
120            writer
121                .write_event(Event::Empty(src_elem))
122                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
123        }
124
125        // Confidence
126        if let Some(ref conf) = self.confidence {
127            let mut conf_elem = BytesStart::new("confidence");
128            conf_elem.push_attribute(("value", conf.value.to_string().as_str()));
129            if let Some(threshold) = conf.threshold {
130                conf_elem.push_attribute(("threshold", threshold.to_string().as_str()));
131            }
132
133            writer
134                .write_event(Event::Empty(conf_elem))
135                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
136        }
137
138        // Hierarchy
139        if let Some(ref hier) = self.hierarchy {
140            writer
141                .write_event(Event::Start(BytesStart::new("hierarchy")))
142                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
143
144            if let Some(ref cell) = hier.cell {
145                let mut cell_elem = BytesStart::new("cell");
146                cell_elem.push_attribute(("id", cell.id.as_str()));
147                if let Some(ref role) = cell.role {
148                    cell_elem.push_attribute(("role", role.as_str()));
149                }
150                writer
151                    .write_event(Event::Empty(cell_elem))
152                    .map_err(|e| CotError::XmlWrite(e.to_string()))?;
153            }
154
155            if let Some(ref formation) = hier.formation {
156                let mut form_elem = BytesStart::new("formation");
157                form_elem.push_attribute(("id", formation.as_str()));
158                writer
159                    .write_event(Event::Empty(form_elem))
160                    .map_err(|e| CotError::XmlWrite(e.to_string()))?;
161            }
162
163            if let Some(ref zone) = hier.zone {
164                let mut zone_elem = BytesStart::new("zone");
165                zone_elem.push_attribute(("id", zone.as_str()));
166                writer
167                    .write_event(Event::Empty(zone_elem))
168                    .map_err(|e| CotError::XmlWrite(e.to_string()))?;
169            }
170
171            writer
172                .write_event(Event::End(BytesEnd::new("hierarchy")))
173                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
174        }
175
176        // Attributes
177        if !self.attributes.is_empty() {
178            writer
179                .write_event(Event::Start(BytesStart::new("attributes")))
180                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
181
182            for (key, attr) in &self.attributes {
183                let mut attr_elem = BytesStart::new("attr");
184                attr_elem.push_attribute(("key", key.as_str()));
185                attr_elem.push_attribute(("type", attr.attr_type.as_str()));
186
187                writer
188                    .write_event(Event::Start(attr_elem))
189                    .map_err(|e| CotError::XmlWrite(e.to_string()))?;
190                writer
191                    .write_event(Event::Text(BytesText::new(&attr.value)))
192                    .map_err(|e| CotError::XmlWrite(e.to_string()))?;
193                writer
194                    .write_event(Event::End(BytesEnd::new("attr")))
195                    .map_err(|e| CotError::XmlWrite(e.to_string()))?;
196            }
197
198            writer
199                .write_event(Event::End(BytesEnd::new("attributes")))
200                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
201        }
202
203        // Status
204        if let Some(ref status) = self.status {
205            let mut status_elem = BytesStart::new("status");
206            status_elem.push_attribute(("operational", status.operational.as_str()));
207            status_elem.push_attribute(("readiness", status.readiness.to_string().as_str()));
208
209            writer
210                .write_event(Event::Empty(status_elem))
211                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
212        }
213
214        // Capabilities
215        for cap in &self.capabilities {
216            let mut cap_elem = BytesStart::new("capability");
217            cap_elem.push_attribute(("type", cap.capability_type.as_str()));
218            cap_elem.push_attribute(("model_version", cap.model_version.as_str()));
219            cap_elem.push_attribute(("precision", cap.precision.to_string().as_str()));
220            cap_elem.push_attribute(("status", cap.status.as_str()));
221
222            writer
223                .write_event(Event::Empty(cap_elem))
224                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
225        }
226
227        // Handoff
228        if let Some(ref handoff) = self.handoff {
229            let mut handoff_elem = BytesStart::new("handoff");
230            handoff_elem.push_attribute(("source_cell", handoff.source_cell.as_str()));
231            handoff_elem.push_attribute(("target_cell", handoff.target_cell.as_str()));
232            handoff_elem.push_attribute(("state", handoff.state.as_str()));
233            handoff_elem.push_attribute(("reason", handoff.reason.as_str()));
234
235            writer
236                .write_event(Event::Empty(handoff_elem))
237                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
238        }
239
240        // Classification
241        if let Some(ref class) = self.classification {
242            let mut class_elem = BytesStart::new("classification");
243            class_elem.push_attribute(("level", class.level.as_str()));
244            if let Some(ref caveat) = class.caveat {
245                class_elem.push_attribute(("caveat", caveat.as_str()));
246            }
247
248            writer
249                .write_event(Event::Empty(class_elem))
250                .map_err(|e| CotError::XmlWrite(e.to_string()))?;
251        }
252
253        writer
254            .write_event(Event::End(BytesEnd::new("_peat_")))
255            .map_err(|e| CotError::XmlWrite(e.to_string()))?;
256
257        Ok(())
258    }
259}
260
261/// Source attribution for track/capability origin
262#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
263pub struct PeatSource {
264    /// Platform that generated this data
265    pub platform: String,
266    /// Model/sensor name
267    pub model: String,
268    /// Model version
269    pub model_version: String,
270}
271
272impl PeatSource {
273    /// Create a new source attribution
274    pub fn new(platform: &str, model: &str, model_version: &str) -> Self {
275        Self {
276            platform: platform.to_string(),
277            model: model.to_string(),
278            model_version: model_version.to_string(),
279        }
280    }
281}
282
283/// Confidence information
284#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
285pub struct PeatConfidence {
286    /// Confidence value (0.0 - 1.0)
287    pub value: f64,
288    /// Optional threshold used
289    pub threshold: Option<f64>,
290}
291
292/// Hierarchy membership information
293#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
294pub struct PeatHierarchy {
295    /// Cell membership
296    pub cell: Option<PeatCellMembership>,
297    /// Formation membership
298    pub formation: Option<String>,
299    /// Zone membership
300    pub zone: Option<String>,
301}
302
303impl PeatHierarchy {
304    /// Create new hierarchy info
305    pub fn new() -> Self {
306        Self::default()
307    }
308
309    /// Set cell membership
310    pub fn with_cell(mut self, id: &str, role: Option<&str>) -> Self {
311        self.cell = Some(PeatCellMembership {
312            id: id.to_string(),
313            role: role.map(|s| s.to_string()),
314        });
315        self
316    }
317
318    /// Set formation
319    pub fn with_formation(mut self, id: &str) -> Self {
320        self.formation = Some(id.to_string());
321        self
322    }
323
324    /// Set zone
325    pub fn with_zone(mut self, id: &str) -> Self {
326        self.zone = Some(id.to_string());
327        self
328    }
329}
330
331/// Cell membership with role
332#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
333pub struct PeatCellMembership {
334    /// Cell ID
335    pub id: String,
336    /// Role within cell (leader, member, tracker, etc.)
337    pub role: Option<String>,
338}
339
340/// Custom attribute
341#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
342pub struct PeatAttribute {
343    /// Attribute value
344    pub value: String,
345    /// Attribute type (string, boolean, number, etc.)
346    pub attr_type: String,
347}
348
349/// Operational status
350#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
351pub struct PeatStatus {
352    /// Operational state
353    pub operational: OperationalStatus,
354    /// Readiness level (0.0 - 1.0)
355    pub readiness: f64,
356}
357
358impl PeatStatus {
359    /// Create a new status
360    pub fn new(operational: OperationalStatus, readiness: f64) -> Self {
361        Self {
362            operational,
363            readiness: readiness.clamp(0.0, 1.0),
364        }
365    }
366}
367
368/// Capability information
369#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
370pub struct PeatCapability {
371    /// Capability type
372    pub capability_type: String,
373    /// Model/version
374    pub model_version: String,
375    /// Precision (0.0 - 1.0)
376    pub precision: f64,
377    /// Current status
378    pub status: OperationalStatus,
379}
380
381impl PeatCapability {
382    /// Create a new capability
383    pub fn new(
384        capability_type: &str,
385        model_version: &str,
386        precision: f64,
387        status: OperationalStatus,
388    ) -> Self {
389        Self {
390            capability_type: capability_type.to_string(),
391            model_version: model_version.to_string(),
392            precision: precision.clamp(0.0, 1.0),
393            status,
394        }
395    }
396}
397
398/// Handoff information
399#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
400pub struct PeatHandoff {
401    /// Source cell releasing track
402    pub source_cell: String,
403    /// Target cell receiving track
404    pub target_cell: String,
405    /// Current handoff state
406    pub state: String,
407    /// Reason for handoff
408    pub reason: String,
409}
410
411impl PeatHandoff {
412    /// Create new handoff info
413    pub fn new(source_cell: &str, target_cell: &str, state: &str, reason: &str) -> Self {
414        Self {
415            source_cell: source_cell.to_string(),
416            target_cell: target_cell.to_string(),
417            state: state.to_string(),
418            reason: reason.to_string(),
419        }
420    }
421}
422
423/// Classification information
424#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
425pub struct PeatClassification {
426    /// Classification level
427    pub level: String,
428    /// Optional caveat
429    pub caveat: Option<String>,
430}
431
432#[cfg(test)]
433mod tests {
434    use super::*;
435
436    #[test]
437    fn test_peat_extension_creation() {
438        let ext = PeatExtension::new()
439            .with_source(PeatSource::new("Alpha-2", "object_tracker", "1.3.0"))
440            .with_confidence(0.89, Some(0.70))
441            .with_hierarchy(
442                PeatHierarchy::new()
443                    .with_cell("Alpha-Team", Some("tracker"))
444                    .with_formation("Formation-1"),
445            )
446            .with_attribute("jacket_color", "blue", "string")
447            .with_status(PeatStatus::new(OperationalStatus::Active, 0.91));
448
449        assert!(ext.source.is_some());
450        assert_eq!(ext.source.as_ref().unwrap().platform, "Alpha-2");
451        assert!(ext.confidence.is_some());
452        assert_eq!(ext.confidence.as_ref().unwrap().value, 0.89);
453    }
454
455    #[test]
456    fn test_peat_extension_to_xml() {
457        let ext = PeatExtension::new()
458            .with_source(PeatSource::new("Platform-1", "sensor", "1.0.0"))
459            .with_confidence(0.85, None);
460
461        let mut writer = Writer::new(Cursor::new(Vec::new()));
462        ext.write_xml(&mut writer).unwrap();
463
464        let xml = String::from_utf8(writer.into_inner().into_inner()).unwrap();
465        assert!(xml.contains("<_peat_"));
466        assert!(xml.contains("version=\"1.0\""));
467        assert!(xml.contains("platform=\"Platform-1\""));
468        assert!(xml.contains("value=\"0.85\""));
469    }
470
471    #[test]
472    fn test_peat_hierarchy() {
473        let hier = PeatHierarchy::new()
474            .with_cell("Alpha-Team", Some("leader"))
475            .with_formation("Formation-1")
476            .with_zone("Zone-A");
477
478        assert_eq!(hier.cell.as_ref().unwrap().id, "Alpha-Team");
479        assert_eq!(hier.cell.as_ref().unwrap().role, Some("leader".to_string()));
480        assert_eq!(hier.formation, Some("Formation-1".to_string()));
481        assert_eq!(hier.zone, Some("Zone-A".to_string()));
482    }
483
484    #[test]
485    fn test_peat_capability() {
486        let cap = PeatCapability::new("OBJECT_TRACKING", "1.3.0", 0.94, OperationalStatus::Active);
487
488        assert_eq!(cap.capability_type, "OBJECT_TRACKING");
489        assert_eq!(cap.precision, 0.94);
490    }
491
492    #[test]
493    fn test_peat_status_readiness_clamped() {
494        let status = PeatStatus::new(OperationalStatus::Active, 1.5);
495        assert_eq!(status.readiness, 1.0);
496
497        let status2 = PeatStatus::new(OperationalStatus::Degraded, -0.5);
498        assert_eq!(status2.readiness, 0.0);
499    }
500
501    #[test]
502    fn test_peat_extension_with_attributes() {
503        let ext = PeatExtension::new()
504            .with_attribute("color", "red", "string")
505            .with_attribute("count", "5", "number")
506            .with_attribute("active", "true", "boolean");
507
508        assert_eq!(ext.attributes.len(), 3);
509        assert_eq!(ext.attributes["color"].value, "red");
510        assert_eq!(ext.attributes["color"].attr_type, "string");
511    }
512
513    #[test]
514    fn test_peat_handoff() {
515        let handoff =
516            PeatHandoff::new("Alpha-Team", "Bravo-Team", "INITIATED", "boundary_crossing");
517
518        assert_eq!(handoff.source_cell, "Alpha-Team");
519        assert_eq!(handoff.target_cell, "Bravo-Team");
520        assert_eq!(handoff.state, "INITIATED");
521    }
522
523    #[test]
524    fn test_peat_classification() {
525        let ext = PeatExtension::new().with_classification("UNCLASSIFIED", Some("FOUO"));
526
527        assert!(ext.classification.is_some());
528        let class = ext.classification.unwrap();
529        assert_eq!(class.level, "UNCLASSIFIED");
530        assert_eq!(class.caveat, Some("FOUO".to_string()));
531    }
532}