tkeyclient 0.1.0

use blake2::{Blake2s256, Digest, digest::FixedOutput};
use serialport::SerialPort;
use std::time::Duration;
mod error;
mod types;

pub use error::*;
pub use types::*;

pub const APP_MAX_SIZE: usize = 0x20000;

pub struct TKey {
    port: Box<dyn SerialPort>,
}

#[cfg(unix)]
const DEFAULT_TTY: &str = "/dev/ttyACM0";
#[cfg(windows)]
const DEFAULT_TTY: &str = "COM6";

const DEFAULT_BAUD: u32 = 62500;

const STATUS_OK: u8 = 0x00;
// const STATUS_BAD: u8 = 0x01;

const CMD_GET_NAME_VERSION: Cmd = Cmd {
    code: 0x01,
    name: "cmdGetNameVersion",
    len: CmdLen::CmdLen1,
    is_app: false,
};
const RSP_GET_NAME_VERSION: Cmd = Cmd {
    code: 0x02,
    name: "rspGetNameVersion",
    len: CmdLen::CmdLen32,
    is_app: false,
};
const CMD_LOAD_APP: Cmd = Cmd {
    code: 0x03,
    name: "cmdLoadApp",
    len: CmdLen::CmdLen128,
    is_app: false,
};
const RSP_LOAD_APP: Cmd = Cmd {
    code: 0x04,
    name: "rspLoadApp",
    len: CmdLen::CmdLen4,
    is_app: false,
};
const CMD_LOAD_APP_DATA: Cmd = Cmd {
    code: 0x05,
    name: "cmdLoadAppData",
    len: CmdLen::CmdLen128,
    is_app: false,
};
const RSP_LOAD_APP_DATA: Cmd = Cmd {
    code: 0x06,
    name: "rspLoadAppData",
    len: CmdLen::CmdLen4,
    is_app: false,
};
const RSP_LOAD_APP_DATA_READY: Cmd = Cmd {
    code: 0x07,
    name: "rspLoadAppDataReady",
    len: CmdLen::CmdLen128,
    is_app: false,
};
#[allow(dead_code)]
const CMD_GET_UDI: Cmd = Cmd {
    code: 0x08,
    name: "cmdGetUDI",
    len: CmdLen::CmdLen1,
    is_app: false,
};
#[allow(dead_code)]
const RSP_GET_UDI: Cmd = Cmd {
    code: 0x09,
    name: "rspGetUDI",
    len: CmdLen::CmdLen32,
    is_app: false,
};

// 4 chars each
const WANT_FW_NAME0: &'static str = "tk1 ";
const WANT_FW_NAME1: &'static str = "mkdf";

struct Cmd {
    code: u8,
    #[allow(dead_code)]
    name: &'static str,
    len: CmdLen,
    is_app: bool,
}

impl TKey {
    pub fn connect(serialport_path: Option<&str>) -> Result<Self> {
        let port = serialport::new(serialport_path.unwrap_or(DEFAULT_TTY), DEFAULT_BAUD)
            .timeout(Duration::from_secs(0))
            .open()?;
        Ok(Self { port })
    }

    pub fn is_firmware_mode(&mut self) -> Result<bool> {
        let name_version = self.get_name_version()?;
        Ok(name_version.name0 == WANT_FW_NAME0 && name_version.name1 == WANT_FW_NAME1)
    }

    pub fn get_name_version(&mut self) -> Result<NameVersion> {
        let tx = TKey::new_frame_buf(CMD_GET_NAME_VERSION, 2)?;
        self.port.write(&tx)?;

        let frame = self.read_frame(RSP_GET_NAME_VERSION, 2)?;
        let version = u32::from_le_bytes(frame.data[8..12].try_into().unwrap());
        Ok(NameVersion {
            name0: String::from_utf8(frame.data[0..4].to_vec()).unwrap(),
            name1: String::from_utf8(frame.data[4..8].to_vec()).unwrap(),
            version,
        })
    }

    fn read_frame(&mut self, cmd: Cmd, id: u8) -> Result<TkeyFrame> {
        self.port.set_timeout(Duration::from_secs(2))?;

        let res = {
            let mut header_buf = [0u8; 1];
            self.port.read(&mut header_buf)?;
            let header = TKeyFrameHeader::parse(header_buf[0])?;
            if header.id != id {
                return tkey_err("wrong header id");
            }
            if header.cmd_len.byte_len() != cmd.len.byte_len() {
                return tkey_err("wrong cmd len");
            }

            let mut buf: Vec<u8> = vec![0; 0 + header.cmd_len.byte_len() as usize];
            self.port.read_exact(buf.as_mut_slice())?;
            let code = buf[0];
            if code != cmd.code {
                return tkey_err("wrong cmd code");
            }

            Ok(TkeyFrame {
                header,
                code,
                data: buf.drain(1..).collect(),
            })
        };
        self.port.set_timeout(Duration::from_secs(0))?;
        res
    }

    fn new_frame_buf(cmd: Cmd, id: u8) -> Result<Vec<u8>> {
        if id > 3 {
            return tkey_err("frame ID must be 0..3");
        }

        let size = 1 + cmd.len.byte_len() as usize;
        let mut tx: Vec<u8> = vec![0u8; size];

        let endpoint: u8 = if cmd.is_app {
            Endpoint::DestApp
        } else {
            Endpoint::DestFW
        } as u8;

        tx[0] = (id << 5) | (endpoint << 3) | cmd.len as u8;
        tx[1] = cmd.code;
        Ok(tx)
    }

    pub fn load_app(&mut self, bin: &[u8], uss: Option<&[u8; 32]>) -> Result<()> {
        let bin_len = bin.len();
        if bin_len > APP_MAX_SIZE {
            return tkey_err("file too big");
        }

        // 1. tell the tkey how large the app is + send USS
        self.load_app_meta(bin_len, uss)?;

        // 2. stream the binary
        let mut offset = 0usize;
        let mut device_digest = [0u8; 32];

        while offset < bin_len {
            let last_chunk = bin_len - offset <= CMD_LOAD_APP_DATA.len.byte_len() as usize - 1;
            let (digest, nsent) = self.load_app_data(&bin[offset..], last_chunk)?;
            if last_chunk {
                device_digest = digest;
            }
            offset += nsent;
        }

        if offset > bin_len {
            return tkey_err("transmitted more than expected");
        }

        // 3. locally compute digest of the whole binary
        let digest = {
            let mut hasher = Blake2s256::new();
            hasher.update(bin);
            let res = hasher.finalize_fixed();
            let mut arr = [0u8; 32];
            arr.copy_from_slice(&res);
            arr
        };

        if device_digest != digest {
            return tkey_err("different digests");
        }

        Ok(())
    }

    fn load_app_meta(&mut self, size: usize, uss: Option<&[u8; 32]>) -> Result<()> {
        let id = 2;
        let mut tx = Self::new_frame_buf(CMD_LOAD_APP, id)?;

        tx[2..6].copy_from_slice(&(size as u32).to_le_bytes());

        if uss.is_none() {
            tx[6] = 0;
        } else {
            tx[6] = 1;
            tx[6..6 + 32].copy_from_slice(uss.unwrap());
        }

        self.port.write(&tx)?;

        let rx = self.read_frame(RSP_LOAD_APP, id)?;

        if rx.data.get(0).copied() != Some(STATUS_OK) {
            return tkey_err("cmdLoadApp failed");
        }
        Ok(())
    }

    /// Send one chunk of the binary.
    ///
    /// Returns `(digest_from_device, bytes_sent)`.
    fn load_app_data(&mut self, content: &[u8], last: bool) -> Result<([u8; 32], usize)> {
        let id = 2;
        let mut tx = Self::new_frame_buf(CMD_LOAD_APP_DATA, id)?;

        let payload_len = CMD_LOAD_APP_DATA.len.byte_len() as usize - 1; // -1 because rx[2] is status
        let mut payload = vec![0u8; payload_len];

        let copied = std::cmp::min(content.len(), payload_len);
        payload[..copied].copy_from_slice(&content[..copied]);
        tx[2..2 + payload_len].copy_from_slice(&payload);

        self.port.write(&tx)?;

        let expected_rsp = if last {
            RSP_LOAD_APP_DATA_READY
        } else {
            RSP_LOAD_APP_DATA
        };

        let rx = self.read_frame(expected_rsp, id)?;
        if rx.data.get(0).copied() != Some(STATUS_OK) {
            return tkey_err("load_app_data failed");
        }

        if last {
            let mut digest = [0u8; 32];
            digest.copy_from_slice(&rx.data[1..33]);
            return Ok((digest, copied));
        }

        Ok(([0u8; 32], copied))
    }
}