foxmark3 0.1.1

//! [`Proxmark`] interface for sending/receiving [`request`]s/[`response`]s.

use std::{
    io,
    time::{Duration, Instant},
};

use bytemuck::from_bytes;
use thiserror::Error;
use tracing::{debug, error, info, trace, warn};

use crate::raw::common::SZ_DATA;

use super::{
    Response,
    common::{Command, FrameOld, NgPayload},
    request::{self, Request},
    response::{self, ResponseNgFrame},
};

// [include/usart_defs.h:20] #define USART_BAUD_RATE 115200
const BAUD_RATE: u32 = 115_200;

const SZ_BUF_MAX: usize = if request::SZ_BUF > response::SZ_BUF {
    request::SZ_BUF
} else {
    response::SZ_BUF
};

/// Creates a new [`Proxmark`] from its serial port path.
pub fn new<'a>(path: impl Into<std::borrow::Cow<'a, str>>) -> Result<Proxmark, DirtyError> {
    let serial = serialport::new(path, BAUD_RATE)
        .data_bits(serialport::DataBits::Eight)
        .stop_bits(serialport::StopBits::One)
        .parity(serialport::Parity::None)
        .open()?;

    let mut proxmark = Proxmark {
        serial: serial.into(),
        buffer: Box::new([0u8; SZ_BUF_MAX]).into(),
        debug: true,
    };

    if let Ok(resp) = proxmark.response(Duration::from_secs(0)) {
        return Err(DirtyError::Dirty(Box::new(resp)));
    }

    Ok(proxmark)
}

/// Finds the only Proxmark3 connected.
///
/// # Errors
///
/// - If zero devices are found: [`FindError::NoProxmark`].
/// - If two or more devices are found: [`FindError::TwoOrMoreProxmarks`] with the first two
///   devices found.
///
/// Other underlying errors may occur and can be found [`FindError`].
#[cfg(target_os = "linux")]
pub fn find_path() -> Result<String, FindError> {
    let subsystem = "tty";

    let mut enumerator = udev::Enumerator::new()?;
    enumerator.match_subsystem(subsystem)?;
    enumerator.match_property("ID_VENDOR", "proxmark.org")?;

    let mut iter = enumerator.scan_devices()?;

    let first = iter.next();
    if let Some(d1) = first {
        let second = iter.next();

        if let Some(d2) = second {
            Err(FindError::TwoOrMoreProxmarks(
                d1.sysname().to_os_string(),
                d2.sysname().to_os_string(),
            ))
        } else {
            let sysname = d1.sysname();

            if let Some(n) = sysname.to_str() {
                Ok(format!("/dev/{n}"))
            } else {
                Err(FindError::InvalidSysname(sysname.to_os_string()))
            }
        }
    } else {
        Err(FindError::NoProxmark)
    }
}

/// Finds the only Proxmark3 connected. This is a convenience function for:
///
/// ```no_run
/// # fn main() -> Result<(), FindError> {
/// use foxmark3::raw;
///
/// raw::new(raw::find_path()?)?
/// # ;
/// # Ok(())
/// # }
/// ```
#[cfg(target_os = "linux")]
pub fn find() -> Result<Proxmark, FindError> {
    let ret = new(find_path()?)?;
    Ok(ret)
}

/// A Proxmark3 device which supports sending/receiving [`request`]/[`response`] frames.
///
/// This may be upconverted to [`commands::Proxmark`](crate::commands::Proxmark) using
/// [`commands::Proxmark::from_raw`](crate::commands::Proxmark::from_raw).
#[must_use]
#[derive(Debug)]
pub struct Proxmark {
    /// The serial port
    serial: debug_ignore::DebugIgnore<Box<dyn serialport::SerialPort>>,
    /// Where serialized requests live, prior to writing them to the serial port
    buffer: debug_ignore::DebugIgnore<Box<[u8; SZ_BUF_MAX]>>,
    /// When true, the following commands will be logged and NOT returned by `response` functions:
    ///
    /// - `0x100: DEBUG_PRINT_STRINGS`
    /// - `0x102: DEBUG_PRINT_INTEGERS`
    /// - `0x103: DEBUG_PRINT_BYTES`
    ///
    /// Some commands e.g. `0x108: STATUS` reply using debug commands, so disabling this flag may
    /// be required in rare situations.
    debug: bool,
}

impl Drop for Proxmark {
    fn drop(&mut self) {
        // Send QUIT_SESSION before we die!

        // Ignore result
        if let Err(err) = self.request(request::ng(Command::QUIT_SESSION, [])) {
            warn!(?err, "failed to send QUIT_SESSION");
        } else {
            debug!("sent QUIT_SESSION");
        }
    }
}

/// An error caused by opening a [`Proxmark`].
#[must_use]
#[derive(Debug, Error)]
pub enum DirtyError {
    /// The error came from the serial port.
    #[error("serialport error")]
    Serialport(#[from] serialport::Error),
    /// Upon opening a the [`Proxmark`], a response was found buffered, which indicates the
    /// [`Proxmark`] was not idle and opening cannot proceed.
    #[error("device was dirty with response: {0:?}")]
    Dirty(Box<Response>),
}

/// An error caused by finding a [`Proxmark`].
#[cfg(target_os = "linux")]
#[must_use]
#[derive(Debug, Error)]
pub enum FindError {
    /// The error came from opening the [`Proxmark`].
    #[error(transparent)]
    Dirty(#[from] DirtyError),
    /// The error came from udev.
    #[error("udev error")]
    Udev(#[from] io::Error),
    #[error("proxmark had invalid sysname: {}", .0.display())]
    /// The udev sysname was unable to be cleanly converted to a UTF-8 string. This requirement is
    /// a restriction of the underlying `serialport` crate, which requires a UTF-8 string to open a
    /// serialport.
    InvalidSysname(std::ffi::OsString),
    /// No device was found
    #[error("no proxmark found")]
    NoProxmark,
    /// Two or more devices were found, so it was unknown which device to open
    #[error("two or more Proxmarks found")]
    TwoOrMoreProxmarks(std::ffi::OsString, std::ffi::OsString),
}

/// An error caused by sending a [`request`].
#[must_use]
#[derive(Debug, Error)]
#[error(transparent)]
pub struct RequestError(#[from] std::io::Error);

pub use response::Error as ResponseError;

/// An error caused by general use of a [`Proxmark`].
#[must_use]
#[derive(Debug, Error)]
pub enum Error {
    /// The error came from a request.
    #[error("error while writing request")]
    Request(#[from] RequestError),
    /// The error came from a response.
    #[error("error while reading response")]
    Response(#[from] ResponseError),
    /// The error came from opening the [`Proxmark`].
    #[error("error while checking if the device was dirty")]
    Dirty(#[from] DirtyError),
}

fn debug_print_strings(flag: u16, data: &[u8]) {
    bitflags::bitflags! {
        #[derive(PartialEq, Eq)]
        struct Flags: u16 {
            const Log = 1 << 0;
            const Newline = 1 << 1;
            const InPlace = 1 << 2;
        }
    }

    let msg = String::from_utf8_lossy(data);

    let conv = Flags::from_bits_retain(flag);
    if conv == Flags::Log {
        info!("Log from PM3: {}", msg);
    } else {
        let newline = if conv.contains(Flags::Newline) {
            "\n"
        } else {
            ""
        };

        if conv.contains(Flags::InPlace) {
            debug!("message(+) from PM3: {msg}{newline}");
        } else {
            debug!("message from PM3: {msg}{newline}");
        }

        let rem = conv.difference(Flags::all());
        if !rem.is_empty() {
            warn!("PM3 message contained unknown flags: {}", rem.bits());
        }
    }
}

impl Proxmark {
    /// Disables debug mode for the duration of the closure. This is used to parse debug
    /// messages as data, rather than having them intercepted as logs.
    pub fn debug_disable<F, T, E>(&mut self, f: F) -> Result<T, E>
    where
        F: FnOnce(&mut Self) -> Result<T, E>,
    {
        self.debug = false;
        let ret = f(self);
        self.debug = true;
        ret
    }

    /// Send a [`Request`] to the device.
    pub fn request(&mut self, request: impl Request) -> Result<(), RequestError> {
        trace!(?request, "sending request");

        let written = request.byte_fold_to(&mut self.buffer[..request::SZ_BUF]);

        self.serial.write_all(written)?;
        self.serial.flush()?;

        Ok(())
    }

    /// Fetch a [`Response`] from the device with a specified timeout
    pub fn response(&mut self, timeout: Duration) -> Result<Response, ResponseError> {
        let ret = TimeoutProxmark::new(self, timeout).response()?;

        Ok(ret)
    }

    /// Blocks awaiting a [`Response`] with a certain [`Command`] for a specified duration.
    /// Responses without the specified [`Command`] are ignored.
    pub fn response_of(
        &mut self,
        cmd: Command,
        timeout: Duration,
    ) -> Result<Response, ResponseError> {
        let mut timed = TimeoutProxmark::new(self, timeout);

        loop {
            let resp = timed.response()?;

            if resp.cmd() == cmd {
                break Ok(resp);
            }

            warn!(?cmd, ?resp, "received unrelated response");
            // Let loop again to get relevant response?
        }
    }

    /// Sends a [`Request`] and blocks awaiting a [`Response`] with the same [`Command`] as the request,
    /// for a specified duration. Responses without the same [`Command`] are ignored.
    pub fn request_response(
        &mut self,
        request: impl Request,
        timeout: Duration,
    ) -> Result<Response, Error> {
        let cmd = request.cmd();

        self.request(request)?;
        let ret = self.response_of(cmd, timeout)?;

        Ok(ret)
    }

    fn read_response(&mut self, timeout: Duration) -> Result<Response, ResponseError> {
        // SAFETY: buffer is SZ_BUF_MAX which is >= response::SZ_BUF, so the slice will always be
        // the correct size to form the array
        let buf: &mut [u8; response::SZ_BUF] = unsafe {
            (&mut self.buffer[..response::SZ_BUF])
                .try_into()
                .unwrap_unchecked()
        };

        self.serial.set_timeout(timeout)?;
        let ret = response::read_from(self.serial.as_mut(), &mut *buf)?;

        trace!(?ret, "received response");

        Ok(ret)
    }

    /// Returns a tuple of:
    ///
    /// 1. Was this packet an overhead packet (i.e. don't return it?)
    /// 2. How long should a timeout (if present)?
    fn parse_overhead(&self, response: &Response) -> (bool, Option<Duration>) {
        match response {
            // WTX can only be sent as NG
            Response::Ng(ResponseNgFrame {
                cmd: Command::WTX,
                shim,
                ..
            }) => {
                if shim.payload().len() == 2 {
                    let ms: u16 = *from_bytes(shim.payload());
                    debug!("received Waiting Time eXtension (WTX) for {ms} ms");

                    (true, Some(Duration::from_millis(ms.into())))
                } else {
                    error!(
                        "received malformed Waiting Time eXtension (WTX) (expected length 2, got {})",
                        shim.payload().len()
                    );

                    (true, None)
                }
            }
            // [src/comms.c:298-339] START
            Response::Ng(ResponseNgFrame {
                cmd: Command::DEBUG_PRINT_STRINGS,
                shim,
                ..
            }) => {
                if shim.payload().len() < 2 {
                    error!(
                        "received malformed debug string message (NG) (expected length >=2, got {})",
                        shim.payload().len()
                    );
                } else {
                    let flag: u16 = *from_bytes(&shim.payload()[..2]);
                    let data = &shim.payload()[2..];

                    if self.debug {
                        debug_print_strings(flag, data);
                    }
                }

                (self.debug, None)
            }
            Response::Old(FrameOld {
                cmd: Command::DEBUG_PRINT_STRINGS,
                args: [len, flag, _],
                payload,
            }) => {
                let data = if *len > SZ_DATA as u64 {
                    error!(
                        "received malformed debug string message (MIX). It purports to have {len} bytes, but the frame has only {SZ_DATA} bytes!",
                    );

                    payload.as_slice()
                } else {
                    // We check above that len <= SZ_DATA = 512
                    #[allow(clippy::cast_possible_truncation)]
                    let len = *len as usize;

                    &payload[..len]
                };

                // maybe more debug flags get added later
                #[allow(clippy::cast_possible_truncation)]
                let flag = *flag as u16;

                if self.debug {
                    debug_print_strings(flag, data);
                }

                (self.debug, None)
            }
            // [src/comms.c:298-339] END
            // [src/comms.c:340-348] START
            Response::Old(FrameOld {
                cmd: Command::DEBUG_PRINT_INTEGERS,
                args: [a, b, c],
                ..
            }) => {
                if self.debug {
                    debug!("integers from PM3: {a:016X}, {b:016X}, {c:016X}");
                }

                (self.debug, None)
            }
            _ => (false, None),
        }
    }
}

struct TimeoutProxmark<'a> {
    proxmark: &'a mut Proxmark,
    original: Duration,
    start: Instant,
    end: Instant,
}

impl<'a> TimeoutProxmark<'a> {
    fn new(proxmark: &'a mut Proxmark, timeout: Duration) -> Self {
        let start = Instant::now();

        Self {
            proxmark,
            original: timeout,
            start,
            end: start + timeout,
        }
    }

    fn response(&mut self) -> Result<Response, ResponseError> {
        loop {
            let res = self.proxmark.read_response(self.end - Instant::now());

            match res {
                Ok(resp) => {
                    let (ignore, extend) = self.proxmark.parse_overhead(&resp);

                    if !ignore {
                        break Ok(resp);
                    } else if let Some(x) = extend {
                        self.end += x;
                    }
                }
                Err(ResponseError::Io(err)) if err.kind() == io::ErrorKind::TimedOut => {
                    break Err(ResponseError::Timeout {
                        original: self.original,
                        extensions: (self.end - self.start).saturating_sub(self.original),
                    });
                }
                Err(err) => break Err(err),
            }
        }
    }
}