improv-setup/src/improv.rs

use anyhow::{
    anyhow,
    bail,
    Context,
    Error,
    Result,
};

pub const IMPROV_HEADER: [u8; 6] = [
    'I' as u8,
    'M' as u8,
    'P' as u8,
    'R' as u8,
    'O' as u8,
    'V' as u8,
];

pub const IMPROV_VERSION: u8 = 0x01;

#[derive(Clone, PartialEq)]
#[repr(u8)]
pub enum PacketType {
    CurrentState = 0x01,
    ErrorState = 0x02,
    RPCCommand = 0x03,
    RPCResult = 0x04,
}

impl TryFrom<&u8> for PacketType {
    type Error = Error;

    fn try_from(b: &u8) -> Result<Self, Self::Error> {
        match b {
            0x01 => Ok(Self::CurrentState),
            0x02 => Ok(Self::ErrorState),
            0x03 => Ok(Self::RPCCommand),
            0x04 => Ok(Self::RPCResult),
            _ => Err(anyhow!("Cannot convert to packet type")),
        }
    }
}

impl std::fmt::Display for PacketType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::CurrentState => write!(f, "Current State"),
            Self::ErrorState => write!(f, "Error State"),
            Self::RPCCommand => write!(f, "RPC Command"),
            Self::RPCResult => write!(f, "RPC Result"),
        }
    }
}

#[derive(Clone)]
#[repr(u8)]
pub enum CurrentState {
    Ready = 0x02,
    Provisioning = 0x03,
    Provisioned = 0x04,
}


impl TryFrom<&u8> for CurrentState {
    type Error = Error;

    fn try_from(b: &u8) -> Result<Self, Self::Error> {
        match b {
            0x02 => Ok(Self::Ready),
            0x03 => Ok(Self::Provisioning),
            0x04 => Ok(Self::Provisioned),
            _ => Err(anyhow!("Cannot convert to current state")),
        }
    }
}

impl std::fmt::Display for CurrentState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Ready => write!(f, "Ready"),
            Self::Provisioning => write!(f, "Provisioning"),
            Self::Provisioned => write!(f, "Provisioned"),
        }
    }
}

#[derive(Clone)]
#[repr(u8)]
pub enum ErrorState {
    NoError = 0x00,
    InvalidRPCPacket = 0x01,
    UnknownRPCCommand = 0x02,
    UnableToConnect = 0x03,
    UnknownError = 0xFF,
}

impl TryFrom<&u8> for ErrorState {
    type Error = Error;

    fn try_from(b: &u8) -> Result<Self, Self::Error> {
        match b {
            0x00 => Ok(Self::NoError),
            0x01 => Ok(Self::InvalidRPCPacket),
            0x02 => Ok(Self::UnknownRPCCommand),
            0x03 => Ok(Self::UnableToConnect),
            0x04 => Ok(Self::UnknownError),
            _ => Err(anyhow!("Cannot convert to error type")),
        }
    }
}

impl std::fmt::Display for ErrorState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NoError => write!(f, "No Error"),
            Self::InvalidRPCPacket => write!(f, "Invalid RPC Packet"),
            Self::UnknownRPCCommand => write!(f, "Unknown RPC Command"),
            Self::UnableToConnect => write!(f, "Unable To Connect"),
            Self::UnknownError => write!(f, "Unknown Error"),
        }
    }
}

#[derive(Clone)]
#[repr(u8)]
pub enum RPCCommand {
    SendWiFiSettings = 0x01,
    RequestCurrentState = 0x02,
    RequestDeviceInformation = 0x03,
    RequestScannedWiFiNetworks = 0x04,
}

impl TryFrom<&u8> for RPCCommand {
    type Error = Error;

    fn try_from(b: &u8) -> Result<Self, Self::Error> {
        match b {
            0x01 => Ok(Self::SendWiFiSettings),
            0x02 => Ok(Self::RequestCurrentState),
            0x03 => Ok(Self::RequestDeviceInformation),
            0x04 => Ok(Self::RequestScannedWiFiNetworks),
            _ => Err(anyhow!("Cannot convert to RPC command")),
        }
    }
}


pub fn calculate_checksum(data: &[u8]) -> u8 {
    // Pass data as full packet, with header, but without checksum byte

    let mut checksum: u8 = 0x00;

    for e in data {
        checksum = checksum.wrapping_add(e.clone());
    }

    return checksum;
}

pub trait ImprovDataPacketType {
    const packet_type: PacketType;
}

pub trait ImprovDataToPacket: ImprovDataPacketType {
    fn to_bytes(&self) -> Vec<u8>;

    fn to_raw_packet(&self) -> RawPacket {
        return RawPacket {
            version: IMPROV_VERSION,
            r#type: Self::packet_type,
            data: self.to_bytes().to_vec(),
        }

    }
}

pub trait ImprovDataFromPacket: ImprovDataPacketType + Sized {
    type Error;

    fn try_from_raw_packet(raw_packet: &RawPacket) -> Result<Self, Self::Error>;
}

pub struct RawPacket {
    pub version: u8,
    pub r#type: PacketType,
    pub data: Vec<u8>,
}

impl RawPacket {
    pub fn try_from_bytes(bytes: &Vec<u8>) -> Result<Self> {
        if bytes.len() < 11 {
            bail!("Packet too small");
        }

        for i in 0..5 {
            if bytes[i] != IMPROV_HEADER[i] {
                bail!("Improv header not found");
            }
        }

        let length: usize = bytes[8].into();

        if bytes.len() != length + 10 {
            bail!("Packet with invalid length");
        }

        return Ok(Self {
            version: bytes[6],
            r#type: PacketType::try_from(&bytes[7])?,
            data: bytes[9..(bytes.len()-1)].to_vec(),
        });
    }

    pub fn to_bytes(self: &Self) -> Vec<u8> {
        let mut out = Vec::new();

        for b in IMPROV_HEADER.clone() {
            out.push(b);
        }

        out.push(self.version.clone());

        out.push(self.r#type.clone() as u8);

        let length: u8 = self.data.len().try_into().unwrap();

        out.push(length.clone());

        for b in self.data.clone() {
            out.push(b);
        }

        let checksum: u8 = calculate_checksum(&out);

        out.push(checksum.clone());

        return out;
    }
}

pub enum ImprovPacket {
    CurrentState(CurrentStatePacket),
    ErrorState(ErrorStatePacket),
    RequestCurrentState(RequestCurrentStatePacket),
    RequestDeviceInformation(RequestDeviceInformationPacket),
    RequestScannedWiFiNetworks(RequestScannedWiFiNetworksPacket),
    RPCResult(RPCResultPacket),
}

impl ImprovPacket {
    pub fn try_from_raw_packet(raw_packet: &RawPacket) -> Result<Self> {
        match raw_packet.r#type {
            PacketType::CurrentState => Ok(ImprovPacket::CurrentState(CurrentStatePacket::try_from_raw_packet(raw_packet)?)),
            PacketType::ErrorState => Ok(ImprovPacket::ErrorState(ErrorStatePacket::try_from_raw_packet(raw_packet)?)),
            //PacketType::RPCCommand => _,
            PacketType::RPCResult => Ok(ImprovPacket::RPCResult(RPCResultPacket::try_from_raw_packet(raw_packet)?)),
            _ => Err(anyhow!("Conversion into packet type not implemented")),
        }
    }
}

pub struct CurrentStatePacket {
    pub current_state: CurrentState,
}

impl ImprovDataPacketType for CurrentStatePacket {
    const packet_type: PacketType = PacketType::CurrentState;
}

impl ImprovDataToPacket for CurrentStatePacket {
    fn to_bytes(self: &Self) -> Vec<u8> {
        let mut out = Vec::new();
        out.push(self.current_state.clone() as u8);
        return out;
    }
}

impl ImprovDataFromPacket for CurrentStatePacket {
    type Error = Error;

    fn try_from_raw_packet(raw_packet: &RawPacket) -> Result<Self, Self::Error>{
        if raw_packet.r#type != Self::packet_type {
            return Err(anyhow!("Packet is not CurrentStatePacket"));
        }

        return Ok(Self {
            current_state: CurrentState::try_from(&raw_packet.data[0])?,
        })
    }
}

pub struct ErrorStatePacket {
    pub error_state: ErrorState,
}

impl ImprovDataPacketType for ErrorStatePacket {
    const packet_type: PacketType = PacketType::ErrorState;
}

impl ImprovDataFromPacket for ErrorStatePacket {
    type Error = Error;

    fn try_from_raw_packet(raw_packet: &RawPacket) -> Result<Self, Self::Error>{
        if raw_packet.r#type != Self::packet_type {
            return Err(anyhow!("Packet is not ErrorState"));
        }

        return Ok(Self {
            error_state: ErrorState::try_from(&raw_packet.data[0])?,
        })
    }
}

pub struct SendWiFiSettingsPacket {
    pub ssid: String,
    pub password: String,
}

impl ImprovDataPacketType for SendWiFiSettingsPacket {
    const packet_type: PacketType = PacketType::RPCCommand;
}

impl ImprovDataToPacket for SendWiFiSettingsPacket {
    fn to_bytes(self: &Self) -> Vec<u8> {
        let ssid_bytes = self.ssid.as_bytes();
        let ssid_len: u8 = ssid_bytes.len().try_into().unwrap();

        let password_bytes = self.password.as_bytes();
        let password_len: u8 = password_bytes.len().try_into().unwrap();

        let mut out = Vec::new();
        out.push(RPCCommand::SendWiFiSettings as u8);
        out.push(1 + ssid_len + 1 + password_len); // data len
        out.push(ssid_len);
        out.extend_from_slice(ssid_bytes);
        out.push(password_len);
        out.extend_from_slice(password_bytes);
        return out;
    }
}

pub struct RequestCurrentStatePacket {
}

impl ImprovDataPacketType for RequestCurrentStatePacket {
    const packet_type: PacketType = PacketType::RPCCommand;
}

impl ImprovDataToPacket for RequestCurrentStatePacket {
    fn to_bytes(self: &Self) -> Vec<u8> {
        let mut out = Vec::new();
        out.push(RPCCommand::RequestCurrentState as u8);
        out.push(0x00); // rpc command payload length
        return out;
    }
}

pub struct RequestDeviceInformationPacket {
}

impl ImprovDataPacketType for RequestDeviceInformationPacket {
    const packet_type: PacketType = PacketType::RPCCommand;
}

impl ImprovDataToPacket for RequestDeviceInformationPacket {
    fn to_bytes(self: &Self) -> Vec<u8> {
        let mut out = Vec::new();
        out.push(RPCCommand::RequestDeviceInformation as u8);
        out.push(0x00); // rpc command payload length
        return out;
    }
}

pub struct RequestScannedWiFiNetworksPacket {
}

impl ImprovDataPacketType for RequestScannedWiFiNetworksPacket {
    const packet_type: PacketType = PacketType::RPCCommand;
}

impl ImprovDataToPacket for RequestScannedWiFiNetworksPacket {
    fn to_bytes(self: &Self) -> Vec<u8> {
        let mut out = Vec::new();
        out.push(RPCCommand::RequestScannedWiFiNetworks as u8);
        out.push(0x00); // rpc command payload length
        return out;
    }
}

pub struct RPCResultPacket {
    pub command_responded_to: RPCCommand,
    pub results: Vec<String>,
}

impl ImprovDataPacketType for RPCResultPacket {
    const packet_type: PacketType = PacketType::RPCResult;
}

impl ImprovDataFromPacket for RPCResultPacket {
    type Error = Error;

    fn try_from_raw_packet(raw_packet: &RawPacket) -> Result<Self, Self::Error>{
        if raw_packet.r#type != Self::packet_type {
            return Err(anyhow!("Packet is not RPCResult"));
        }

        let mut results: Vec<String> = Vec::new();

        let mut data_position: usize = 2;
        loop {
            if data_position >= raw_packet.data.len() {
                break;
            }

            // find string bounds
            let current_string_len: usize = usize::from(raw_packet.data[data_position]);
            let current_string_begin = data_position + 1;
            let current_string_end = data_position + current_string_len;

            // load string and append to results
            let string_bytes = &raw_packet.data[current_string_begin..current_string_end+1];
            let string = std::str::from_utf8(string_bytes).context("Failed to convert response to UTF-8 string")?.to_string();
            results.push(string);

            // next data position
            data_position = data_position + current_string_len + 1;
        }

        return Ok(Self {
            command_responded_to: RPCCommand::try_from(&raw_packet.data[0]).context("Failed to determine RPC command this packet responds to")?,
            results: results,
        })
    }
}