lsdk 0.0.3

//! High level device handle with the dashboard commands.

use crate::apdu::{Apdu, Response, StatusWord};
use crate::error::{Error, Result};
use crate::os::{
    self, AppAndVersion, BatteryStatus, BatteryStatusKind, EndorsementInfo, EndorsementSlot,
    InstalledApp, LanguagePack, OnboardIdentity, OsVersion, ScreenSaver, ScreenSaverAnimation,
    ScreenSaverSpeed,
};
use crate::transport::Transport;

/// A device reachable through a [`Transport`].
pub struct Device<T: Transport> {
    transport: T,
}

impl<T: Transport> Device<T> {
    pub fn new(transport: T) -> Self {
        Self { transport }
    }

    pub fn transport(&self) -> &T {
        &self.transport
    }

    pub fn into_inner(self) -> T {
        self.transport
    }

    /// Send raw bytes and get the raw answer (data + SW).
    pub async fn exchange_raw(&self, apdu: &[u8]) -> Result<Vec<u8>> {
        self.transport.exchange(apdu).await
    }

    /// Send an APDU and get the parsed response (whatever the status).
    pub async fn exchange(&self, apdu: &Apdu) -> Result<Response> {
        let raw = self.transport.exchange(&apdu.encode()?).await?;
        Response::from_raw(raw)
    }

    /// Send an APDU and return its data, failing on a non `0x9000` status.
    pub async fn send(&self, apdu: &Apdu) -> Result<Vec<u8>> {
        self.exchange(apdu).await?.ok()
    }

    /// Firmware / hardware information (`GET_VERSION`). Must be on the dashboard.
    pub async fn get_version(&self) -> Result<OsVersion> {
        let r = self.exchange(&os::GET_VERSION).await?;
        match r.sw {
            StatusWord::DEVICE_NOT_ONBOARDED | StatusWord::DEVICE_NOT_ONBOARDED_2 => {
                Err(Error::App("device is not onboarded".into()))
            }
            StatusWord::CLA_NOT_SUPPORTED | StatusWord::INS_NOT_SUPPORTED => {
                Err(Error::DashboardExpected)
            }
            _ => OsVersion::parse(&r.ok()?),
        }
    }

    /// Currently running app name and version (works in any app).
    pub async fn get_app_and_version(&self) -> Result<AppAndVersion> {
        AppAndVersion::parse(&self.send(&os::GET_APP_AND_VERSION).await?)
    }

    /// Whether the device is on the dashboard (or in bootloader).
    pub async fn is_on_dashboard(&self) -> Result<bool> {
        let r = self.exchange(&os::GET_APP_AND_VERSION).await?;
        match r.sw {
            StatusWord::OK => Ok(AppAndVersion::parse(&r.data)?.is_dashboard()),
            StatusWord::CLA_NOT_SUPPORTED | StatusWord::CLA_NOT_SUPPORTED_BOOTLOADER => Ok(true),
            StatusWord::INS_NOT_SUPPORTED => Ok(false),
            sw => Err(sw.into()),
        }
    }

    /// Ask the running app to exit back to the dashboard.
    ///
    /// Note: on most models the USB device re-enumerates when switching apps,
    /// so the transport may need to be reopened afterwards.
    pub async fn quit_app(&self) -> Result<()> {
        self.send(&os::QUIT_APP).await.map(|_| ())
    }

    /// Open an app by name (from the dashboard). Requires user confirmation on
    /// some firmwares. See [`Device::quit_app`] about re-enumeration.
    pub async fn open_app(&self, name: &str) -> Result<()> {
        let r = self.exchange(&os::open_app(name)).await?;
        match r.sw {
            StatusWord::OK => Ok(()),
            StatusWord::APP_NOT_FOUND | StatusWord(0x6984) => {
                Err(Error::NotFound(format!("app {name:?} is not installed")))
            }
            sw if sw.is_refused() => Err(Error::UserRefused),
            sw => Err(sw.into()),
        }
    }

    /// List installed apps (requires user consent on the device unless the
    /// secure channel was already allowed).
    pub async fn list_apps(&self) -> Result<Vec<InstalledApp>> {
        let mut apps = Vec::new();
        let mut r = self.exchange(&os::LIST_APPS).await?;
        loop {
            match r.sw {
                StatusWord::OK => {}
                sw if sw.is_refused() => return Err(Error::UserRefused),
                sw => return Err(sw.into()),
            }
            if r.data.is_empty() {
                break;
            }
            apps.extend(os::parse_list_apps(&r.data)?);
            r = self.exchange(&os::LIST_APPS_CONTINUE).await?;
        }
        Ok(apps)
    }

    /// Device name (as shown over Bluetooth). The device asks the user to
    /// allow it, so avoid calling this implicitly.
    pub async fn get_device_name(&self) -> Result<String> {
        let data = self.send(&os::GET_DEVICE_NAME).await?;
        Ok(String::from_utf8_lossy(&data).into_owned())
    }

    /// Rename the device (requires user confirmation).
    pub async fn set_device_name(&self, name: &str) -> Result<()> {
        if name.is_empty() || name.len() > 32 {
            return Err(Error::InvalidInput("name must be 1..=32 bytes".into()));
        }
        match self.exchange(&os::set_device_name(name)).await?.sw {
            StatusWord::OK => Ok(()),
            sw if sw.is_refused() => Err(Error::UserRefused),
            sw => Err(sw.into()),
        }
    }

    fn screen_saver_err(sw: StatusWord) -> Error {
        match sw {
            sw if sw.is_refused() => Error::UserRefused,
            // The dashboard answers 0x6d02 outside recovery mode and for a
            // malformed request (verified on a Nano S Plus).
            StatusWord(0x6d00..=0x6dff) => Error::App(
                "screensaver command rejected: only accepted in recovery mode (Nano X / Nano S Plus), with a valid request"
                    .into(),
            ),
            sw => sw.into(),
        }
    }

    /// Current custom screensaver, `None` for the default one. Recovery mode only.
    pub async fn get_screen_saver(&self) -> Result<Option<ScreenSaver>> {
        let r = self.exchange(&os::get_screen_saver()).await?;
        if !r.sw.is_ok() {
            return Err(Self::screen_saver_err(r.sw));
        }
        ScreenSaver::parse(&r.data)
    }

    /// Set a custom scrolling-text screensaver (user confirmation required).
    /// Recovery mode only.
    pub async fn set_screen_saver(
        &self,
        text: &str,
        animation: ScreenSaverAnimation,
        speed: ScreenSaverSpeed,
    ) -> Result<()> {
        let r = self
            .exchange(&os::set_screen_saver(text, animation, speed)?)
            .await?;
        if r.sw.is_ok() {
            Ok(())
        } else {
            Err(Self::screen_saver_err(r.sw))
        }
    }

    /// Restore the default screensaver (user confirmation required). Recovery mode only.
    pub async fn reset_screen_saver(&self) -> Result<()> {
        let r = self.exchange(&os::reset_screen_saver()).await?;
        if r.sw.is_ok() {
            Ok(())
        } else {
            Err(Self::screen_saver_err(r.sw))
        }
    }

    fn check(&self, r: crate::apdu::Response) -> Result<Vec<u8>> {
        match r.sw {
            StatusWord::OK => Ok(r.data),
            sw if sw.is_refused() => Err(Error::UserRefused),
            sw => Err(sw.into()),
        }
    }

    /// Restart the dashboard (`RESET`).
    pub async fn reset(&self) -> Result<()> {
        match self.exchange(&os::RESET).await {
            Ok(r) => self.check(r).map(|_| ()),
            Err(Error::Disconnected) => Ok(()),
            Err(e) => Err(e),
        }
    }

    /// Bluetooth MAC address (BLE models), as `aa:bb:cc:dd:ee:ff`.
    pub async fn get_device_mac(&self) -> Result<String> {
        let d = self.check(self.exchange(&os::GET_DEVICE_MAC).await?)?;
        Ok(d.iter()
            .map(|b| format!("{b:02x}"))
            .collect::<Vec<_>>()
            .join(":"))
    }

    /// Set the battery charge limit (percent, models with a battery).
    pub async fn set_battery_charge_limit(&self, percent: u8) -> Result<()> {
        self.check(
            self.exchange(&os::set_battery_charge_limit(percent)?)
                .await?,
        )
        .map(|_| ())
    }

    /// Installed language packs.
    pub async fn list_language_packs(&self) -> Result<Vec<LanguagePack>> {
        let mut out = Vec::new();
        let mut first = true;
        loop {
            let d = self.check(self.exchange(&os::list_language_packs(first)).await?)?;
            match os::parse_language_pack(&d)? {
                Some(p) => out.push(p),
                None => break,
            }
            first = false;
            if out.len() > 32 {
                break;
            }
        }
        Ok(out)
    }

    /// Replay an APDU script (one hex APDU per line, `#` comments allowed),
    /// e.g. a signed language pack. Stops at the first non-`9000` status.
    /// `on_progress(done, total)` counts APDUs.
    pub async fn run_apdu_script(
        &self,
        script: &str,
        on_progress: &mut dyn FnMut(usize, usize),
    ) -> Result<()> {
        let apdus: Vec<Vec<u8>> = script
            .lines()
            .map(|l| l.split('#').next().unwrap_or_default().trim())
            .filter(|l| !l.is_empty())
            .map(|l| hex::decode(l.replace(' ', "")))
            .collect::<core::result::Result<_, _>>()
            .map_err(|e| Error::InvalidInput(format!("APDU script: {e}")))?;
        if apdus.is_empty() {
            return Err(Error::InvalidInput("empty APDU script".into()));
        }
        let total = apdus.len();
        for (i, apdu) in apdus.iter().enumerate() {
            let r = crate::apdu::Response::from_raw(self.exchange_raw(apdu).await?)?;
            match r.sw.0 {
                0x9000 => {}
                0x6985 | 0x5501 => return Err(Error::UserRefused),
                0x6a84 | 0x5102 => return Err(Error::App("not enough space on the device".into())),
                _ => {
                    return Err(Error::App(format!(
                        "APDU {} of {total} failed: {}",
                        i + 1,
                        r.sw
                    )));
                }
            }
            on_progress(i + 1, total);
        }
        Ok(())
    }

    /// Install a language pack from its signed APDU script, replacing the
    /// installed packs (one at a time is supported).
    pub async fn install_language_script(
        &self,
        script: &str,
        on_progress: &mut dyn FnMut(usize, usize),
    ) -> Result<()> {
        for installed in self.list_language_packs().await? {
            self.delete_language_pack(installed.id as u8).await?;
        }
        self.run_apdu_script(script, on_progress).await
    }

    /// Remove a language pack (the device may ask for confirmation).
    pub async fn delete_language_pack(&self, language: u8) -> Result<()> {
        self.check(self.exchange(&os::delete_language_pack(language)).await?)
            .map(|_| ())
    }

    /// Endorsement key and certificate of a slot.
    pub async fn endorsement_info(&self, slot: EndorsementSlot) -> Result<EndorsementInfo> {
        let r = self.exchange(&os::endorsement_info(slot)).await?;
        if !r.sw.is_ok() && !r.sw.is_refused() && r.sw != StatusWord::LOCKED_DEVICE {
            return Err(Error::NotFound(format!(
                "endorsement slot {} is not set up ({})",
                slot as u8, r.sw
            )));
        }
        EndorsementInfo::parse(&self.check(r)?)
    }

    /// Start an endorsement: the device generates the slot key and returns
    /// `(public_key, device_signature)`; the signature (by the device
    /// attestation key) covers `sha256(0xFF || public_key)`.
    pub async fn endorse_start(&self, slot: EndorsementSlot) -> Result<(Vec<u8>, Vec<u8>)> {
        let d = self.check(self.exchange(&os::endorse_set_start(slot)).await?)?;
        if d.len() < 65 {
            return Err(Error::InvalidResponse("short endorsement key".into()));
        }
        Ok((d[..65].to_vec(), d[65..].to_vec()))
    }

    /// Finish an endorsement with the owner certificate (DER signature over
    /// `sha256(0xFE || public_key)`).
    pub async fn endorse_commit(&self, certificate: &[u8]) -> Result<()> {
        self.check(self.exchange(&os::endorse_set_commit(certificate)).await?)
            .map(|_| ())
    }

    /// Onboard from the host. **Sends the PIN and recovery phrase in clear
    /// over USB**: test devices only.
    pub async fn onboard(
        &self,
        identity: OnboardIdentity,
        pin: &str,
        prefix: &str,
        passphrase: &str,
        words: &str,
    ) -> Result<()> {
        self.check(
            self.exchange(&os::onboard(identity, pin, prefix, passphrase, words)?)
                .await?,
        )
        .map(|_| ())
    }

    /// Factory personalization: device public key, serial and certificate
    /// request (devices not yet personalized only).
    pub async fn factory_get_device_public_key(&self) -> Result<Vec<u8>> {
        self.check(self.exchange(&os::GET_DEVICE_PUBLIC_KEY).await?)
    }

    /// Factory test command (`E0 44`, devices without serial number only).
    pub async fn factory_test(&self, p1: u8, p2: u8, data: &[u8]) -> Result<Vec<u8>> {
        self.check(self.exchange(&os::factory_test(p1, p2, data)).await?)
    }

    /// Factory: read the cryptographic port rights (`E0 C5`).
    pub async fn factory_get_cxport(&self) -> Result<Vec<u8>> {
        self.check(self.exchange(&os::GET_CXPORT).await?)
    }

    /// Factory: `SET_CERTIFICATE` (`E0 41`, raw payload).
    pub async fn factory_set_certificate(&self, p2: u8, data: &[u8]) -> Result<Vec<u8>> {
        self.check(
            self.exchange(&Apdu::with_data(0xe0, 0x41, 0x00, p2, data.to_vec()))
                .await?,
        )
    }

    /// Factory: `SET_CXPORT` (`E0 C4`, raw payload).
    pub async fn factory_set_cxport(&self, data: &[u8]) -> Result<Vec<u8>> {
        self.check(
            self.exchange(&Apdu::with_data(0xe0, 0xc4, 0x00, 0x00, data.to_vec()))
                .await?,
        )
    }

    /// Query one battery metric (raw payload).
    pub async fn get_battery_raw(&self, kind: BatteryStatusKind) -> Result<Vec<u8>> {
        self.send(&os::get_battery_status(kind)).await
    }

    /// Battery summary (models with a battery only).
    pub async fn get_battery_status(&self) -> Result<BatteryStatus> {
        let mut st = BatteryStatus::default();
        let p = self.get_battery_raw(BatteryStatusKind::Percentage).await?;
        st.percentage = p.first().copied().filter(|v| *v <= 100);
        if let Ok(v) = self.get_battery_raw(BatteryStatusKind::Voltage).await
            && v.len() >= 2
        {
            st.voltage_mv = Some(u16::from_be_bytes([v[0], v[1]]));
        }
        if let Ok(v) = self.get_battery_raw(BatteryStatusKind::Temperature).await {
            st.temperature_c = v.first().map(|b| *b as i8);
        }
        if let Ok(v) = self.get_battery_raw(BatteryStatusKind::Current).await {
            st.current_ma = v.first().map(|b| *b as i8);
        }
        if let Ok(v) = self.get_battery_raw(BatteryStatusKind::Flags).await
            && v.len() >= 4
        {
            let f = u32::from_be_bytes([v[0], v[1], v[2], v[3]]);
            st.raw_flags = Some(f);
            st.charging = f & 0x1 != 0;
            st.usb_connected = f & 0x2 != 0;
            st.ble_connected = f & 0x4 != 0;
        }
        Ok(st)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::transport::mock::MockTransport;

    fn block_on<F: core::future::Future>(f: F) -> F::Output {
        tokio::runtime::Builder::new_current_thread()
            .build()
            .unwrap()
            .block_on(f)
    }

    #[test]
    fn list_apps_paging() {
        let mut page = String::from("01");
        page += &format!("{:02x}", 70 + 3);
        page += "00100000";
        page += &"aa".repeat(32);
        page += &"bb".repeat(32);
        page += "03657468";
        let t = MockTransport::new()
            .expect("e0de000000", &format!("{page}9000"))
            .expect("e0df000000", "9000");
        let d = Device::new(t);
        let apps = block_on(d.list_apps()).unwrap();
        assert_eq!(apps.len(), 1);
        assert_eq!(apps[0].name, "eth");
        assert!(d.transport().is_done());
    }

    #[test]
    fn dashboard_detection() {
        let t = MockTransport::new().expect("b001000000", "6e01");
        assert!(block_on(Device::new(t).is_on_dashboard()).unwrap());
    }
}