zkteco 1.0.0

//! Pure, side-effect-free building blocks of the ZKTeco wire protocol.
//!
//! Everything here is testable without a device: packet framing, the checksum,
//! the comm-key scrambler, and the timestamp codec. The networking layer in
//! [`crate::device`] composes these functions.
//!
//! ## Byte order
//! The protocol is little-endian throughout. We avoid an external `byteorder`
//! dependency and instead use small explicit helpers ([`wr`]) plus
//! `u16::from_le_bytes` / `u32::from_le_bytes` at the call sites, so every layout
//! stays visible and commented.

use crate::constants;

/// A tiny little-endian byte writer.
///
/// This stands in for Python's `struct.pack`: instead of format strings like
/// `"<HHI"`, callers push fields in order. Keeping it explicit makes the exact
/// on-wire layout obvious at each call site.
#[derive(Default)]
pub(crate) struct Writer {
    buf: Vec<u8>,
}

impl Writer {
    /// Start an empty writer.
    pub(crate) fn new() -> Self {
        Writer { buf: Vec::new() }
    }

    /// Append a single byte (`struct` format `B`/`b`).
    pub(crate) fn u8(&mut self, v: u8) -> &mut Self {
        self.buf.push(v);
        self
    }

    /// Append a little-endian `u16` (`struct` format `H`).
    pub(crate) fn u16(&mut self, v: u16) -> &mut Self {
        self.buf.extend_from_slice(&v.to_le_bytes());
        self
    }

    /// Append a little-endian `u32` (`struct` format `I`).
    pub(crate) fn u32(&mut self, v: u32) -> &mut Self {
        self.buf.extend_from_slice(&v.to_le_bytes());
        self
    }

    /// Append `n` zero bytes (`struct` padding `x` / fixed-width fields).
    pub(crate) fn pad(&mut self, n: usize) -> &mut Self {
        self.buf.resize(self.buf.len() + n, 0);
        self
    }

    /// Append raw bytes verbatim.
    pub(crate) fn bytes(&mut self, b: &[u8]) -> &mut Self {
        self.buf.extend_from_slice(b);
        self
    }

    /// Append `src` truncated or zero-padded to exactly `width` bytes
    /// (`struct` fixed-width string field `"{width}s"`).
    pub(crate) fn fixed(&mut self, src: &[u8], width: usize) -> &mut Self {
        let take = src.len().min(width);
        self.buf.extend_from_slice(&src[..take]);
        self.buf.resize(self.buf.len() + (width - take), 0);
        self
    }

    /// Consume the writer and return the assembled bytes.
    pub(crate) fn finish(self) -> Vec<u8> {
        self.buf
    }
}

/// Convenience: build a `Vec<u8>` from a closure operating on a [`Writer`].
pub(crate) fn wr(f: impl FnOnce(&mut Writer)) -> Vec<u8> {
    let mut w = Writer::new();
    f(&mut w);
    w.finish()
}

/// Read a little-endian `u16` from the start of `b` (panics only on a
/// programming error where `b` is too short, which callers guard against).
pub(crate) fn u16_le(b: &[u8]) -> u16 {
    u16::from_le_bytes([b[0], b[1]])
}

/// Read a little-endian `u32` from the start of `b`.
pub(crate) fn u32_le(b: &[u8]) -> u32 {
    u32::from_le_bytes([b[0], b[1], b[2], b[3]])
}

/// Compute the packet checksum.
///
/// Ported from `zkemsdk.c` (`createChkSum`) via `pyzk`. The algorithm sums the
/// payload as a sequence of little-endian 16-bit words, folding any carry above
/// [`USHRT_MAX`](constants::USHRT_MAX) back in, adds a trailing odd byte if the
/// length is odd, then returns the one's-complement of the result.
pub(crate) fn checksum(payload: &[u8]) -> u16 {
    let mut sum: u32 = 0;
    let mut i = 0;
    let n = payload.len();

    // Sum complete 16-bit little-endian words.
    while i + 1 < n {
        sum += u16::from_le_bytes([payload[i], payload[i + 1]]) as u32;
        if sum > constants::USHRT_MAX as u32 {
            sum -= constants::USHRT_MAX as u32;
        }
        i += 2;
    }

    // A leftover odd byte is added on its own.
    if i < n {
        sum += payload[n - 1] as u32;
    }

    // Fold any remaining overflow.
    while sum > constants::USHRT_MAX as u32 {
        sum -= constants::USHRT_MAX as u32;
    }

    // Complement. NOTE: this is *not* a plain bitwise NOT. The protocol wraps
    // with USHRT_MAX (65535) rather than 65536, an intentional quirk preserved
    // from `zkemsdk.c`/`pyzk`: compute `~sum` then add 65535 until non-negative.
    let mut chk: i64 = !(sum as i64);
    while chk < 0 {
        chk += constants::USHRT_MAX as i64;
    }
    chk as u16
}

/// Build a command packet header (without TCP framing).
///
/// Layout (`pyzk` `__create_header`): `<H H H H>` = command, checksum,
/// session id, reply id — followed by the `command_string` payload.
///
/// Reproduces a quirk of `pyzk`/`zkemsdk.c` precisely: the checksum is computed
/// over the packet using the *original* `reply_id` (with the checksum field
/// zeroed), but the transmitted packet carries `reply_id + 1`
/// (see [`next_reply_id`]) together with that same checksum. Devices accept this.
pub(crate) fn create_header(
    command: u16,
    command_string: &[u8],
    session_id: u16,
    reply_id: u16,
) -> Vec<u8> {
    // Checksum is taken over the buffer with the *original* reply id and a zero
    // checksum field.
    let zeroed = wr(|w| {
        w.u16(command)
            .u16(0) // checksum placeholder
            .u16(session_id)
            .u16(reply_id)
            .bytes(command_string);
    });
    let chk = checksum(&zeroed);

    // The packet we actually send advances the reply id by one.
    let rid = next_reply_id(reply_id);
    wr(|w| {
        w.u16(command)
            .u16(chk)
            .u16(session_id)
            .u16(rid)
            .bytes(command_string);
    })
}

/// Advance a reply id, wrapping at [`USHRT_MAX`](constants::USHRT_MAX) exactly as
/// `pyzk` does (`reply_id += 1; if reply_id >= USHRT_MAX { reply_id -= USHRT_MAX }`).
pub(crate) fn next_reply_id(reply_id: u16) -> u16 {
    // Done in u32 so the wrap matches pyzk exactly: the protocol folds at
    // USHRT_MAX (65535), NOT at 65536, so a plain u16 `wrapping_add` would be
    // off-by-one for reply ids at the very top of the range.
    let mut next = reply_id as u32 + 1;
    if next >= constants::USHRT_MAX as u32 {
        next -= constants::USHRT_MAX as u32;
    }
    next as u16
}

/// Wrap a packet in the 8-byte TCP "top" header used on stream connections.
///
/// Layout (`<H H I>`): magic1, magic2, payload length.
pub(crate) fn create_tcp_top(packet: &[u8]) -> Vec<u8> {
    wr(|w| {
        w.u16(constants::MACHINE_PREPARE_DATA_1)
            .u16(constants::MACHINE_PREPARE_DATA_2)
            .u32(packet.len() as u32)
            .bytes(packet);
    })
}

/// Validate a TCP "top" header and return the declared payload size.
///
/// Returns `0` if the buffer is too short or the magic does not match — same
/// contract as `pyzk`'s `__test_tcp_top`.
pub(crate) fn test_tcp_top(packet: &[u8]) -> u32 {
    if packet.len() <= 8 {
        return 0;
    }
    let m1 = u16_le(&packet[0..2]);
    let m2 = u16_le(&packet[2..4]);
    if m1 == constants::MACHINE_PREPARE_DATA_1 && m2 == constants::MACHINE_PREPARE_DATA_2 {
        u32_le(&packet[4..8])
    } else {
        0
    }
}

/// Scramble a numeric password and session id into the 4-byte comm-key the device
/// expects for `CMD_AUTH`.
///
/// Faithful port of `commpro.c` `MakeKey` via `pyzk`'s `make_commkey`:
/// 1. Reverse the low 32 bits of `key` bit-by-bit.
/// 2. Add `session_id`.
/// 3. XOR the four bytes with the ASCII of `Z`, `K`, `S`, `O`.
/// 4. Byte-swap as two 16-bit words (`HH` -> reversed).
/// 5. XOR again with `ticks` in a fixed pattern.
pub(crate) fn make_commkey(key: u32, session_id: u32, ticks: u8) -> [u8; 4] {
    // Step 1: bit-reverse `key` across 32 positions.
    let mut k: u32 = 0;
    for i in 0..32 {
        if key & (1 << i) != 0 {
            k = (k << 1) | 1;
        } else {
            k <<= 1;
        }
    }

    // Step 2: add the session id (wrapping, matching Python's 32-bit pack).
    k = k.wrapping_add(session_id);

    // Step 3: XOR each byte with "ZKSO".
    let mut b = k.to_le_bytes();
    b[0] ^= b'Z';
    b[1] ^= b'K';
    b[2] ^= b'S';
    b[3] ^= b'O';

    // Step 4: interpret as two LE u16 words and swap them (`pack('HH', k[1], k[0])`).
    let w0 = u16::from_le_bytes([b[0], b[1]]);
    let w1 = u16::from_le_bytes([b[2], b[3]]);
    let swapped = wr(|w| {
        w.u16(w1).u16(w0);
    });
    let mut b = [swapped[0], swapped[1], swapped[2], swapped[3]];

    // Step 5: final XOR pattern with the low byte of `ticks`.
    let t = ticks;
    b[0] ^= t;
    b[1] ^= t;
    b[2] = t; // note: byte 2 is replaced by the tick, not XORed
    b[3] ^= t;
    b
}

/// Decode the 4-byte packed device timestamp into a calendar date-time.
///
/// Ported from `zkemsdk.c` `DecodeTime`. The device stores time as a single
/// integer of "seconds" using fixed 31-day months / 12-month years; we unwind it
/// field by field. `year` is offset from 2000.
pub(crate) fn decode_time(packed: u32) -> chrono::NaiveDateTime {
    let mut t = packed;
    let second = t % 60;
    t /= 60;
    let minute = t % 60;
    t /= 60;
    let hour = t % 24;
    t /= 24;
    let day = t % 31 + 1;
    t /= 31;
    let month = t % 12 + 1;
    t /= 12;
    let year = t + 2000;

    naive(year as i32, month, day, hour, minute, second)
}

/// Decode the 6-byte "time hex" form used by real-time events: one byte each for
/// year-2000, month, day, hour, minute, second.
pub(crate) fn decode_timehex(b: &[u8; 6]) -> chrono::NaiveDateTime {
    naive(
        2000 + b[0] as i32,
        b[1] as u32,
        b[2] as u32,
        b[3] as u32,
        b[4] as u32,
        b[5] as u32,
    )
}

/// Encode a calendar date-time into the packed integer the device's clock uses.
///
/// Inverse of [`decode_time`], ported from `zkemsdk.c` `EncodeTime`.
pub(crate) fn encode_time(dt: &chrono::NaiveDateTime) -> u32 {
    use chrono::{Datelike, Timelike};
    let year = (dt.year() % 100) as u32;
    let month = dt.month();
    let day = dt.day();
    let hour = dt.hour();
    let minute = dt.minute();
    let second = dt.second();

    ((year * 12 * 31 + (month - 1) * 31 + day - 1) * (24 * 60 * 60))
        + (hour * 60 + minute) * 60
        + second
}

/// Build a `NaiveDateTime`, clamping obviously-invalid field combinations so a
/// corrupt packet can never panic (the device's fixed 31-day months can yield,
/// e.g., "Feb 31"). Invalid dates fall back to the Unix epoch start of that day.
fn naive(
    year: i32,
    month: u32,
    day: u32,
    hour: u32,
    minute: u32,
    second: u32,
) -> chrono::NaiveDateTime {
    use chrono::NaiveDate;
    let date = NaiveDate::from_ymd_opt(year, month.clamp(1, 12), day.clamp(1, 31))
        // Walk the day back until valid (handles 31-day-month artifacts).
        .or_else(|| {
            (1..=31)
                .rev()
                .filter_map(|d| NaiveDate::from_ymd_opt(year, month.clamp(1, 12), d))
                .next()
        })
        .unwrap_or_else(|| NaiveDate::from_ymd_opt(2000, 1, 1).expect("constant date is valid"));
    date.and_hms_opt(hour.min(23), minute.min(59), second.min(59))
        .unwrap_or_else(|| date.and_hms_opt(0, 0, 0).expect("midnight is valid"))
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::NaiveDate;

    // All expected values below were produced by running the reference `pyzk`
    // Python implementation, guaranteeing byte-for-byte parity.

    #[test]
    fn commkey_matches_pyzk() {
        assert_eq!(make_commkey(0, 0, 50), [97, 125, 50, 121]);
        assert_eq!(make_commkey(123456, 57, 50), [38, 127, 50, 249]);
        assert_eq!(make_commkey(0, 1234, 99), [48, 44, 99, 44]);
    }

    #[test]
    fn checksum_matches_pyzk() {
        // A CMD_CONNECT-shaped header (reply id 0xfffe, zero checksum field).
        assert_eq!(checksum(&[0xe8, 0x03, 0, 0, 0, 0, 0xfe, 0xff]), 64535);
        // An odd-length buffer exercises the trailing-byte path.
        assert_eq!(checksum(&[1, 2, 3, 4, 5]), 63989);
    }

    #[test]
    fn time_encode_matches_pyzk_and_roundtrips() {
        let when = NaiveDate::from_ymd_opt(2024, 5, 6)
            .unwrap()
            .and_hms_opt(7, 8, 9)
            .unwrap();
        let encoded = encode_time(&when);
        assert_eq!(encoded, 782_550_489);
        // Decoding the encoded value returns the original instant.
        assert_eq!(decode_time(encoded), when);
    }

    #[test]
    fn timehex_decodes() {
        // year-2000, month, day, hour, minute, second
        let when = decode_timehex(&[24, 5, 6, 7, 8, 9]);
        assert_eq!(
            when,
            NaiveDate::from_ymd_opt(2024, 5, 6)
                .unwrap()
                .and_hms_opt(7, 8, 9)
                .unwrap()
        );
    }

    #[test]
    fn reply_id_wraps_at_ushrt_max() {
        // 65534 -> 65535 wraps to 0 (first packet after connect).
        assert_eq!(next_reply_id(constants::USHRT_MAX - 1), 0);
        assert_eq!(next_reply_id(0), 1);
    }

    #[test]
    fn tcp_top_roundtrips() {
        let payload = [9u8, 8, 7, 6];
        let framed = create_tcp_top(&payload);
        // 8-byte magic header + payload.
        assert_eq!(framed.len(), 12);
        assert_eq!(test_tcp_top(&framed), payload.len() as u32);
    }

    #[test]
    fn header_increments_sent_reply_id_but_checksums_original() {
        // With reply_id = 65534 the transmitted packet carries reply_id 0.
        let pkt = create_header(constants::CMD_CONNECT, &[], 0, constants::USHRT_MAX - 1);
        assert_eq!(u16_le(&pkt[6..8]), 0);
        // And the command field is preserved.
        assert_eq!(u16_le(&pkt[0..2]), constants::CMD_CONNECT);
    }
}