km003c_lib/
packet.rs

1#![allow(unused_parens)]
2
3use crate::constants::*;
4use crate::error::KMError;
5use bytes::Bytes;
6use modular_bitfield::prelude::*;
7use num_enum::{FromPrimitive, IntoPrimitive};
8
9#[bitfield(bytes = 4)]
10#[derive(Debug, Clone, Copy, PartialEq)]
11pub struct CtrlHeader {
12    pub packet_type: B7,
13    /// Reserved flag bit in the first byte, purpose unknown.
14    /// IMPORTANT: This is NOT an indicator that an "extended header" is present.
15    pub reserved_flag: bool,
16    pub id: u8,
17    #[skip]
18    unused: bool,
19    pub attribute: B15,
20}
21
22#[bitfield(bytes = 4)]
23#[derive(Debug, Clone, Copy, PartialEq)]
24pub struct DataHeader {
25    pub packet_type: B7,
26    /// Reserved flag bit in the first byte, purpose unknown.
27    /// IMPORTANT: This is NOT an indicator that an "extended header" is present.
28    /// As per protocol research, PutData (0x41) packets always include a 4-byte
29    /// extended header regardless of this bit.
30    pub reserved_flag: bool,
31    pub id: u8,
32    #[skip]
33    unused: B6,
34    pub obj_count_words: B10,
35}
36
37/// Header used by StreamingAuth, whose final two bytes are a raw result value.
38#[bitfield(bytes = 4)]
39#[derive(Debug, Clone, Copy, PartialEq)]
40pub(crate) struct StreamingAuthHeader {
41    pub packet_type: B7,
42    pub reserved_flag: bool,
43    pub id: u8,
44    pub attribute: u16,
45}
46
47#[bitfield(bytes = 4)]
48#[derive(Debug, Clone, Copy, PartialEq)]
49pub struct ExtendedHeader {
50    pub attribute: B15,
51    pub next: bool,
52    pub chunk: B6,
53    pub size: B10,
54}
55
56/// KM003C protocol packet types.
57///
58/// Values < 0x40 are control packet types, >= 0x40 are data packet types.
59/// Unknown types have been discovered through protocol analysis but their
60/// purpose has not yet been reverse engineered.
61#[derive(Debug, Clone, Copy, PartialEq, Eq, IntoPrimitive, FromPrimitive)]
62#[repr(u8)]
63pub enum PacketType {
64    // 0 is reserved
65    // less than 0x40 is ctrl type
66    Sync = 0x01,
67    Connect = 0x02,
68    Disconnect = 0x03,
69    Reset = 0x04,
70    Accept = 0x05,
71    Rejected = 0x06,
72    Finished = 0x07,
73    JumpAprom = 0x08,
74    JumpDfu = 0x09,
75    GetStatus = 0x0A,
76    Error = 0x0B,
77    GetData = 0x0C,
78    GetFile = 0x0D,
79    StartGraph = 0x0E,       // Start AdcQueue streaming with rate selector
80    StopGraph = 0x0F,        // Stop AdcQueue streaming
81    EnablePdMonitor = 0x10,  // Enable PD sniffer/capture
82    DisablePdMonitor = 0x11, // Disable PD sniffer/capture
83
84    // Authentication commands
85    MemoryRead = 0x44,    // 68 - Read device memory (encrypted)
86    StreamingAuth = 0x4C, // 76 - Authenticate for AdcQueue streaming
87
88    // Error responses
89    NotReadable = 0x27, // Memory address not accessible
90
91    // Unknown control types discovered in protocol analysis
92    Unknown26 = 26,
93    Unknown58 = 58,
94
95    // >= 0x40 is data type
96    Head = 64,
97    PutData = 65,
98    #[num_enum(catch_all)]
99    Unknown(u8),
100}
101
102impl PacketType {
103    pub fn is_ctrl_type(&self) -> bool {
104        let value: u8 = (*self).into();
105        value < 0x40
106    }
107}
108
109// Python support for PacketType
110#[cfg(feature = "python")]
111impl<'py> pyo3::IntoPyObject<'py> for PacketType {
112    type Target = pyo3::PyAny;
113    type Output = pyo3::Bound<'py, Self::Target>;
114    type Error = pyo3::PyErr;
115
116    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, <Self as pyo3::IntoPyObject<'py>>::Error> {
117        // Convert to u8 value for Python
118        let value: u8 = self.into();
119        Ok(value.into_pyobject(py).unwrap().into_any())
120    }
121}
122
123/// Attribute values used in command headers and extended headers.
124///
125/// These values specify the type of data or command being sent.
126/// Unknown attributes have been discovered through protocol analysis
127/// but their purpose has not yet been reverse engineered.
128#[derive(Debug, Clone, Copy, PartialEq, Eq, IntoPrimitive, FromPrimitive)]
129#[repr(u16)]
130pub enum Attribute {
131    None = 0,
132    Adc = 0x1,
133    AdcQueue = 0x2,
134    AdcQueue10k = 0x4,
135    Settings = 0x8,
136    /// PD protocol data. Response is either 12-byte PdStatus (measurements only)
137    /// or longer PdEventStream (12-byte preamble + PD wire events).
138    PdPacket = 0x10,
139
140    // Unknown attributes discovered in protocol analysis
141    // TODO: Reverse engineer the purpose of these attributes
142    Unknown512 = 512,     // 0x200 - found with PutData packets
143    Unknown1609 = 1609,   // 0x649 - found with Unknown26 packets
144    Unknown11046 = 11046, // 0x2B26 - found with Unknown44 packets
145    Unknown26817 = 26817, // 0x68C1 - found with Unknown58 packets
146
147    #[num_enum(catch_all)]
148    Unknown(u16),
149}
150
151// Python support for Attribute
152#[cfg(feature = "python")]
153impl<'py> pyo3::IntoPyObject<'py> for Attribute {
154    type Target = pyo3::PyAny;
155    type Output = pyo3::Bound<'py, Self::Target>;
156    type Error = pyo3::PyErr;
157
158    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, Self::Error> {
159        // Convert to u16 value for Python
160        let value: u16 = self.into();
161        Ok(value.into_pyobject(py).unwrap().into_any())
162    }
163}
164
165/// Set of attributes for use in request masks.
166/// Can represent single or multiple attributes combined with bitwise OR.
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
168#[cfg_attr(feature = "python", derive(pyo3::IntoPyObject))]
169pub struct AttributeSet {
170    mask: u16,
171}
172
173impl AttributeSet {
174    pub const fn empty() -> Self {
175        Self { mask: 0 }
176    }
177
178    pub const fn from_raw(mask: u16) -> Self {
179        Self { mask }
180    }
181
182    pub fn single(attr: Attribute) -> Self {
183        Self { mask: attr.into() }
184    }
185
186    pub fn from_attributes<I>(attrs: I) -> Self
187    where
188        I: IntoIterator<Item = Attribute>,
189    {
190        Self {
191            mask: attrs
192                .into_iter()
193                .map(|a| {
194                    let val: u16 = a.into();
195                    val
196                })
197                .fold(0, |acc, x| acc | x),
198        }
199    }
200
201    pub fn with(mut self, attr: Attribute) -> Self {
202        let val: u16 = attr.into();
203        self.mask |= val;
204        self
205    }
206
207    pub fn contains(&self, attr: Attribute) -> bool {
208        let val: u16 = attr.into();
209        self.mask & val != 0
210    }
211
212    /// Check if any of the given attributes are present
213    pub fn contains_any<I>(&self, attrs: I) -> bool
214    where
215        I: IntoIterator<Item = Attribute>,
216    {
217        attrs.into_iter().any(|a| self.contains(a))
218    }
219
220    /// Check if all of the given attributes are present
221    pub fn contains_all<I>(&self, attrs: I) -> bool
222    where
223        I: IntoIterator<Item = Attribute>,
224    {
225        attrs.into_iter().all(|a| self.contains(a))
226    }
227
228    /// Remove an attribute from the set
229    pub fn without(mut self, attr: Attribute) -> Self {
230        let val: u16 = attr.into();
231        self.mask &= !val;
232        self
233    }
234
235    /// Iterate over all attributes in the set
236    pub fn iter(&self) -> impl Iterator<Item = Attribute> + '_ {
237        (0..16).filter_map(move |bit| {
238            let value = 1u16 << bit;
239            if self.mask & value != 0 {
240                Some(Attribute::from_primitive(value))
241            } else {
242                None
243            }
244        })
245    }
246
247    pub fn to_vec(&self) -> Vec<Attribute> {
248        self.iter().collect()
249    }
250
251    pub const fn raw(&self) -> u16 {
252        self.mask
253    }
254
255    pub fn is_empty(&self) -> bool {
256        self.mask == 0
257    }
258
259    pub fn len(&self) -> usize {
260        self.mask.count_ones() as usize
261    }
262}
263
264impl From<Attribute> for AttributeSet {
265    fn from(attr: Attribute) -> Self {
266        Self::single(attr)
267    }
268}
269
270impl FromIterator<Attribute> for AttributeSet {
271    fn from_iter<I: IntoIterator<Item = Attribute>>(iter: I) -> Self {
272        Self::from_attributes(iter)
273    }
274}
275
276/// Represents a single logical packet within a PutData response.
277/// PutData packets can contain multiple chained logical packets,
278/// each with its own extended header and payload.
279#[derive(Debug, Clone, PartialEq)]
280#[cfg_attr(
281    feature = "python",
282    pyo3::pyclass(get_all, skip_from_py_object, name = "LogicalPacket")
283)]
284pub struct LogicalPacket {
285    pub attribute: Attribute,
286    pub next: bool,
287    pub chunk: u8,
288    pub size: u16,
289    pub payload: Vec<u8>,
290}
291
292#[derive(Debug, Clone, PartialEq)]
293pub enum RawPacket {
294    Ctrl {
295        header: CtrlHeader,
296        payload: Vec<u8>,
297    },
298    SimpleData {
299        header: DataHeader,
300        payload: Vec<u8>,
301    },
302    Data {
303        header: DataHeader,
304        logical_packets: Vec<LogicalPacket>,
305    },
306}
307
308impl RawPacket {
309    pub fn id(&self) -> u8 {
310        match self {
311            RawPacket::Ctrl { header, .. } => header.id(),
312            RawPacket::SimpleData { header, .. } => header.id(),
313            RawPacket::Data { header, .. } => header.id(),
314        }
315    }
316
317    pub fn packet_type(&self) -> PacketType {
318        match self {
319            RawPacket::Ctrl { header, .. } => PacketType::from_primitive(header.packet_type()),
320            RawPacket::SimpleData { header, .. } => PacketType::from_primitive(header.packet_type()),
321            RawPacket::Data { header, .. } => PacketType::from_primitive(header.packet_type()),
322        }
323    }
324
325    /// Get the attribute for control packets
326    pub fn get_attribute(&self) -> Option<Attribute> {
327        match self {
328            RawPacket::Ctrl { header, .. } => Some(Attribute::from_primitive(header.attribute())),
329            _ => None,
330        }
331    }
332
333    /// Get the attribute set for control packets (for GetData requests)
334    pub fn get_attribute_set(&self) -> Option<AttributeSet> {
335        match self {
336            RawPacket::Ctrl { header, .. } => Some(AttributeSet::from_raw(header.attribute())),
337            _ => None,
338        }
339    }
340
341    /// Get logical packets for Data variant
342    pub fn logical_packets(&self) -> Option<&[LogicalPacket]> {
343        match self {
344            RawPacket::Data { logical_packets, .. } => Some(logical_packets),
345            _ => None,
346        }
347    }
348
349    /// Validate that response attributes match the request mask
350    ///
351    /// Returns Ok(()) if all response attributes were requested in the mask,
352    /// or Err if there's a mismatch.
353    pub fn validate_correlation(&self, request_mask: u16) -> Result<(), KMError> {
354        match self {
355            RawPacket::Data { logical_packets, .. } => {
356                let request_set = AttributeSet::from_raw(request_mask);
357
358                // Check each logical packet's attribute
359                for lp in logical_packets {
360                    if !request_set.contains(lp.attribute) {
361                        let expected: Vec<u16> = request_set.iter().map(|a| a.into()).collect();
362                        let actual: Vec<u16> = logical_packets.iter().map(|lp| lp.attribute.into()).collect();
363
364                        return Err(KMError::AttributeMismatch { expected, actual });
365                    }
366                }
367
368                Ok(())
369            }
370            _ => Ok(()), // Non-data packets don't need validation
371        }
372    }
373
374    /// Check if this is an empty PutData response (no logical packets)
375    pub fn is_empty_response(&self) -> bool {
376        matches!(self, RawPacket::Data { logical_packets, .. } if logical_packets.is_empty())
377    }
378}
379
380impl TryFrom<Bytes> for RawPacket {
381    type Error = KMError;
382
383    fn try_from(mut bytes: Bytes) -> Result<Self, Self::Error> {
384        // Check minimum length first to prevent panic in split_to
385        if bytes.len() < MAIN_HEADER_SIZE {
386            return Err(KMError::InvalidPacket(format!(
387                "Packet too short for header: expected {}, got {}",
388                MAIN_HEADER_SIZE,
389                bytes.len()
390            )));
391        }
392
393        // the first byte contains packet type (7 bits) + header flag bit
394        let first_byte = bytes[0]; // Safe now that we know len >= 4
395        // Extract only the packet type (lower 7 bits), ignoring the header flag bit
396        let package_type_byte = first_byte & 0x7F;
397        let is_ctrl_packet = PacketType::from_primitive(package_type_byte).is_ctrl_type();
398
399        let header_bytes: [u8; 4] = bytes
400            .split_to(4) // Safe now - we know there are at least 4 bytes
401            .as_ref()
402            .try_into()
403            .unwrap(); // Safe to unwrap since we know the slice is exactly 4 bytes
404        let mut payload = bytes;
405
406        if is_ctrl_packet {
407            let header = CtrlHeader::from_bytes(header_bytes);
408            Ok(RawPacket::Ctrl {
409                header,
410                payload: payload.to_vec(),
411            })
412        } else {
413            let header = DataHeader::from_bytes(header_bytes);
414            let packet_type = PacketType::from_primitive(header.packet_type());
415
416            // Only PutData packets have chained logical packets with extended headers
417            if packet_type == PacketType::PutData {
418                // Check for empty PutData (obj_count_words == 0)
419                if header.obj_count_words() == 0 || payload.is_empty() {
420                    // Valid empty response - device has no data
421                    return Ok(RawPacket::Data {
422                        header,
423                        logical_packets: vec![],
424                    });
425                }
426
427                if payload.len() < EXTENDED_HEADER_SIZE {
428                    // Spec indicates all PutData (0x41) packets  must carry a 4-byte extended header.
429                    // We currently fall back to SimpleData when payload < 4 for
430                    // robustness against malformed frames.
431                    return Ok(RawPacket::SimpleData {
432                        header,
433                        payload: payload.to_vec(),
434                    });
435                }
436
437                // Parse chained logical packets
438                let mut logical_packets = Vec::new();
439
440                loop {
441                    if payload.len() < EXTENDED_HEADER_SIZE {
442                        return Err(KMError::InvalidPacket(format!(
443                            "Insufficient bytes for extended header: need {}, got {}",
444                            EXTENDED_HEADER_SIZE,
445                            payload.len()
446                        )));
447                    }
448
449                    // Parse extended header
450                    let ext_header_bytes = payload.split_to(4);
451                    let ext_header_array: [u8; 4] = ext_header_bytes
452                        .as_ref()
453                        .try_into()
454                        .map_err(|_| KMError::InvalidPacket("Failed to extract extended header bytes".to_string()))?;
455                    let ext = ExtendedHeader::from_bytes(ext_header_array);
456
457                    let has_next = ext.next();
458                    let attribute = Attribute::from_primitive(ext.attribute());
459                    // AdcQueue uses chunk as its sample count and size as bytes per sample.
460                    let payload_size = ext.size() as usize
461                        * if attribute == Attribute::AdcQueue {
462                            ext.chunk() as usize
463                        } else {
464                            1
465                        };
466
467                    // For AdcQueue, the size field indicates sample size (20 bytes),
468                    // but the actual payload contains multiple samples.
469                    // Take all remaining payload if this is the last logical packet.
470                    let logical_payload = if !has_next && attribute == Attribute::AdcQueue {
471                        // Last packet and AdcQueue: take all remaining bytes
472                        let all = payload.clone();
473                        payload = Bytes::new();
474                        all
475                    } else {
476                        // Normal case: take exactly size bytes
477                        if payload.len() < payload_size {
478                            return Err(KMError::InvalidPacket(format!(
479                                "Insufficient payload bytes: expected {}, got {}",
480                                payload_size,
481                                payload.len()
482                            )));
483                        }
484                        let chunk = payload.slice(..payload_size);
485                        payload = payload.slice(payload_size..);
486                        chunk
487                    };
488
489                    logical_packets.push(LogicalPacket {
490                        attribute,
491                        next: has_next,
492                        chunk: ext.chunk(),
493                        size: ext.size(),
494                        payload: logical_payload.to_vec(),
495                    });
496
497                    // Check if there are more logical packets
498                    if !has_next {
499                        break;
500                    }
501                }
502
503                if logical_packets.is_empty() {
504                    return Err(KMError::InvalidPacket(
505                        "PutData packet must have at least one logical packet".to_string(),
506                    ));
507                }
508
509                Ok(RawPacket::Data {
510                    header,
511                    logical_packets,
512                })
513            } else {
514                Ok(RawPacket::SimpleData {
515                    header,
516                    payload: payload.to_vec(),
517                })
518            }
519        }
520    }
521}
522
523// Python support
524#[cfg(feature = "python")]
525impl<'py> pyo3::IntoPyObject<'py> for RawPacket {
526    type Target = pyo3::PyAny;
527    type Output = pyo3::Bound<'py, Self::Target>;
528    type Error = pyo3::PyErr;
529
530    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, Self::Error> {
531        use pyo3::types::{PyDict, PyDictMethods, PyListMethods};
532
533        let dict = PyDict::new(py);
534        match self {
535            RawPacket::Ctrl { header, payload } => {
536                let inner = PyDict::new(py);
537                let header_dict = PyDict::new(py);
538                header_dict.set_item("packet_type", header.packet_type())?;
539                header_dict.set_item("reserved_flag", header.reserved_flag())?;
540                header_dict.set_item("id", header.id())?;
541                header_dict.set_item("attribute", header.attribute())?;
542                inner.set_item("header", header_dict)?;
543                inner.set_item("payload", payload)?;
544                dict.set_item("Ctrl", inner)?;
545            }
546            RawPacket::SimpleData { header, payload } => {
547                let inner = PyDict::new(py);
548                let header_dict = PyDict::new(py);
549                header_dict.set_item("packet_type", header.packet_type())?;
550                header_dict.set_item("reserved_flag", header.reserved_flag())?;
551                header_dict.set_item("id", header.id())?;
552                header_dict.set_item("obj_count_words", header.obj_count_words())?;
553                inner.set_item("header", header_dict)?;
554                inner.set_item("payload", payload)?;
555                dict.set_item("SimpleData", inner)?;
556            }
557            RawPacket::Data {
558                header,
559                logical_packets,
560            } => {
561                let inner = PyDict::new(py);
562                let header_dict = PyDict::new(py);
563                header_dict.set_item("packet_type", header.packet_type())?;
564                header_dict.set_item("reserved_flag", header.reserved_flag())?;
565                header_dict.set_item("id", header.id())?;
566                header_dict.set_item("obj_count_words", header.obj_count_words())?;
567                inner.set_item("header", header_dict)?;
568                // Convert LogicalPacket vec -> list of dicts for Python
569                let lp_list = pyo3::types::PyList::empty(py);
570                for lp in logical_packets {
571                    let lp_dict = PyDict::new(py);
572                    lp_dict.set_item("attribute", u16::from(lp.attribute))?;
573                    lp_dict.set_item("next", lp.next)?;
574                    lp_dict.set_item("chunk", lp.chunk)?;
575                    lp_dict.set_item("size", lp.size)?;
576                    lp_dict.set_item("payload", lp.payload)?;
577                    lp_list.append(lp_dict)?;
578                }
579                inner.set_item("logical_packets", lp_list)?;
580                dict.set_item("Data", inner)?;
581            }
582        }
583        Ok(dict.into_any())
584    }
585}
586
587impl From<RawPacket> for Bytes {
588    fn from(packet: RawPacket) -> Self {
589        let (header_bytes, payload) = match packet {
590            RawPacket::Ctrl { header, payload } => (header.into_bytes(), payload),
591            RawPacket::SimpleData { header, payload } => (header.into_bytes(), payload),
592            RawPacket::Data {
593                header,
594                logical_packets,
595            } => {
596                // Reconstruct chained logical packets
597                let mut full_payload = Vec::new();
598
599                for logical_packet in logical_packets {
600                    // Build extended header
601                    let ext = ExtendedHeader::new()
602                        .with_attribute(logical_packet.attribute.into())
603                        .with_next(logical_packet.next)
604                        .with_chunk(logical_packet.chunk)
605                        .with_size(logical_packet.size);
606
607                    full_payload.extend_from_slice(&ext.into_bytes());
608                    full_payload.extend_from_slice(&logical_packet.payload);
609                }
610
611                (header.into_bytes(), full_payload)
612            }
613        };
614
615        // Create the full message by combining header and payload
616        let mut message = Vec::with_capacity(4 + payload.len());
617        message.extend_from_slice(&header_bytes);
618        message.extend_from_slice(payload.as_ref());
619
620        Bytes::from(message)
621    }
622}
km003c_lib/packet.rs

km003c_lib/
packet.rs
