improv-setup/src/improv.rs

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= &'static str;

    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("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= &'static str;

    fn try_from(b: &u8) -> Result<Self, Self::Error> {
        match b {
            0x02 => Ok(Self::Ready),
            0x03 => Ok(Self::Provisioning),
            0x04 => Ok(Self::Provisioned),
            _ => Err("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= &'static str;

    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("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,
}

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, &str> {
        if bytes.len() < 11 {
            return Err("Packet too small");
        }

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

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

        if bytes.len() != length + 10 {
            return Err("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 {
    CurrentStateResponse(CurrentStateResponse),
    ErrorState(ErrorStatePacket),
    RequestCurrentStateCommand(RequestCurrentStateCommand),
}

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

pub struct CurrentStateResponse {
    pub current_state: CurrentState,
}

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

impl ImprovDataToPacket for CurrentStateResponse {
    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 CurrentStateResponse {
    type Error = &'static str;

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

        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 = &'static str;

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

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

pub struct RequestCurrentStateCommand {
}

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

impl ImprovDataToPacket for RequestCurrentStateCommand {
    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;
    }
}