use nom::bytes::complete::take; use nom::error::{Error, ErrorKind}; use nom::multi::many0; use nom::number::complete; use nom::IResult; use nom::{Err, InputTake}; use nom::sequence::Tuple; use num_enum::{FromPrimitive, IntoPrimitive, TryFromPrimitive}; use std::convert::TryFrom; use std::fmt::Display; use bitflags::bitflags; use byteorder::{LittleEndian, WriteBytesExt}; use strum_macros::{EnumIter, IntoStaticStr}; #[repr(u8)] #[derive(Copy, Clone, TryFromPrimitive, IntoPrimitive, Debug)] pub enum ReportId { CommandOut = 1, DataOut = 2, CommandIn = 3, DataIn = 4, } #[repr(u8)] #[derive(Copy, Clone, TryFromPrimitive, IntoPrimitive, Debug, EnumIter)] pub enum CommandTag { FlashEraseAll = 0x01, FlashEraseRegion = 0x02, ReadMemory = 0x03, WriteMemory = 0x04, FillMemory = 0x05, GetProperty = 0x07, ReceiveSbFile = 0x08, Execute = 0x09, Call = 0x0a, Reset = 0x0b, SetProperty = 0x0c, ConfigureMemory = 0x11, KeyProvision = 0x15, } #[repr(u8)] #[derive(Copy, Clone, TryFromPrimitive, IntoPrimitive, Debug)] pub enum ResponseTag { GenericResponse = 0xa0, ReadMemoryResponse = 0xa3, GetPropertyResponse = 0xa7, FlashReadOnceResponse = 0xaf, KeyProvisionResponse = 0xb5, } #[repr(u32)] #[derive(Copy, Clone, PartialEq, IntoStaticStr, FromPrimitive, IntoPrimitive, Debug)] pub enum StatusCode { Success = 0, StatusFlashAlignmentError = 101, StatusFlashEccError = 116, StatusMemoryRangeInvalid = 10200, StatusMemoryBlankPageReadDisallowed = 10211, #[num_enum(catch_all)] Unknown(u32), } impl Display for StatusCode { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { StatusCode::Success => f.write_str("Success"), StatusCode::Unknown(code) => f.write_fmt(format_args!("Error({})", code)), code => { let name: &'static str = code.into(); f.write_fmt(format_args!("{}({})", name, u32::from(code.to_owned()))) } } } } #[derive(Debug)] pub enum ResponseParameters { GenericResponse(GenericResponseParams), ReadMemoryResponse(ReadMemoryResponseParams), GetPropertyResponse(GetPropertyResponseParams), FlashReadOnceResponse(FlashReadOnceResponseParams), KeyProvisionResponse(KeyProvisionResponseParams), } #[derive(Debug)] pub struct GenericResponseParams { pub status: StatusCode, pub command: CommandTag, } #[derive(Debug)] pub struct ReadMemoryResponseParams { pub status: StatusCode, pub data_bytes: u32, } #[derive(Debug)] pub struct GetPropertyResponseParams { pub status: StatusCode, pub properties: Vec, } #[derive(Debug)] pub struct FlashReadOnceResponseParams { status: StatusCode, byte_count: u32, read_data: [u8; 20], } #[derive(Debug)] pub struct KeyProvisionResponseParams { status: StatusCode, data_bytes: u32, } bitflags! { pub struct CommandFlags: u8 { const DATA_FOLLOWS = 0b00000001; } } impl Default for CommandFlags { fn default() -> Self { CommandFlags { bits: 0 } } } #[derive(Debug, Clone)] pub struct CommandPacket { pub tag: CommandTag, pub flags: CommandFlags, pub reserved: u8, pub param_count: u8, pub params: Vec, } #[derive(Debug)] pub struct ResponsePacket { pub tag: ResponseTag, pub flags: CommandFlags, pub reserved: u8, pub param_count: u8, pub params: ResponseParameters, } #[derive(Debug)] pub struct DataPacket { pub bytes: Vec, } pub struct UsbPacket { report_id: ReportId, _padding: u8, packet_length: u16, packet: Packet, } #[derive(Debug)] //TODO: Refactor this to use traits and borrows pub enum Packet { CommandPacket(CommandPacket), ResponsePacket(ResponsePacket), DataPacket(DataPacket), } fn status_code(input: &[u8]) -> IResult<&[u8], StatusCode> { let (input, code) = complete::le_u32(input)?; Ok((input, StatusCode::from(code))) } fn response_tag(input: &[u8]) -> IResult<&[u8], ResponseTag> { let res = take(1u8)(input); match res { Ok(ir) => match ResponseTag::try_from(ir.1[0]) { Ok(val) => Ok((ir.0, val)), Err(e) => { println!("{:?}", e); unimplemented!() } }, Err(e) => Err(e), } } fn command_tag(input: &[u8]) -> IResult<&[u8], CommandTag> { let res = take(1u8)(input); match res { Ok(ir) => match CommandTag::try_from(ir.1[0]) { Ok(val) => Ok((ir.0, val)), Err(e) => { println!("{:?}", e); unimplemented!() } }, Err(e) => Err(e), } } fn generic_response_params(input: &[u8]) -> IResult<&[u8], ResponseParameters> { let (input, (status, command)) = (status_code, complete::u8).parse(input)?; Ok(( input, ResponseParameters::GenericResponse(GenericResponseParams { status, command: CommandTag::try_from(command).unwrap(), }), )) } fn get_property_response_params(input: &[u8]) -> IResult<&[u8], ResponseParameters> { let (input, (status, properties)) = (status_code, many0(complete::le_u32)).parse(input)?; Ok(( input, ResponseParameters::GetPropertyResponse(GetPropertyResponseParams { status, properties }), )) } fn get_read_memory_response_params(input: &[u8]) -> IResult<&[u8], ResponseParameters> { let (input, (status, properties)) = (status_code, many0(complete::le_u32)).parse(input)?; assert!(properties.len() == 1); Ok(( input, ResponseParameters::ReadMemoryResponse(ReadMemoryResponseParams { status, data_bytes: properties[0], }), )) } fn report_id(input: &[u8]) -> IResult<&[u8], ReportId> { let (input, id) = complete::u8(input)?; match ReportId::try_from(id) { Ok(val) => Ok((input, val)), Err(_) => Err(Err::Error(Error::new(input, ErrorKind::Fail))), } } pub fn usb_packet(input: &[u8]) -> IResult<&[u8], Packet> { let (input, (reportid, _, packet_length)) = (report_id, complete::u8, complete::le_u16).parse(input)?; let (input, buf) = take(packet_length)(input)?; let pack = match reportid { ReportId::CommandOut => command_packet(buf)?, ReportId::DataOut => unimplemented!(), ReportId::CommandIn => response_packet(buf)?, ReportId::DataIn => data_packet(buf, packet_length)?, }; Ok((input, pack.1)) } pub fn response_packet(input: &[u8]) -> IResult<&[u8], Packet> { let (input, (tag, flags, reserved, param_count)) = (response_tag, complete::u8, complete::u8, complete::u8).parse(input)?; let (input, params) = match tag { ResponseTag::GenericResponse => generic_response_params(input)?, ResponseTag::GetPropertyResponse => get_property_response_params(input)?, ResponseTag::ReadMemoryResponse => get_read_memory_response_params(input)?, ResponseTag::FlashReadOnceResponse => unimplemented!(), ResponseTag::KeyProvisionResponse => unimplemented!(), }; Ok(( input, Packet::ResponsePacket(ResponsePacket { tag, flags: CommandFlags::from_bits_truncate(flags), reserved, param_count, params, }), )) } pub fn data_packet(input: &[u8], packet_length: u16) -> IResult<&[u8], Packet> { let data = input.take(packet_length as usize); Ok(( input, Packet::DataPacket(DataPacket { bytes: data.to_vec(), }), )) } pub fn command_packet(input: &[u8]) -> IResult<&[u8], Packet> { let (input, (tag, flags, reserved, param_count)) = (command_tag, complete::u8, complete::u8, complete::u8).parse(input)?; Ok(( input, Packet::CommandPacket(CommandPacket { tag, flags: CommandFlags::from_bits_truncate(flags), reserved, param_count, params: Vec::new(), }), )) } impl Packet { pub fn write(&self, w: &mut W) -> Result<(), std::io::Error> { //TODO: Refactor to a trait match self { Packet::CommandPacket(p) => { w.write(&[ p.tag as u8, p.flags.bits() as u8, p.reserved as u8, p.param_count as u8, ])?; for param in &p.params { w.write_u32::(*param)?; } } Packet::DataPacket(p) => { w.write(p.bytes.as_slice())?; } Packet::ResponsePacket(p) => { unimplemented!("Serializing response packets is not implemented"); } } Ok(()) } pub fn length(&self) -> u16 { let header_len = 4u16; match self { Packet::CommandPacket(p) => header_len + (p.param_count as u16) * 4, Packet::ResponsePacket(p) => header_len + (p.param_count as u16) * 4, Packet::DataPacket(p) => p.bytes.len() as u16, } } } impl UsbPacket { pub fn new(report_id: ReportId, packet: Packet) -> UsbPacket { UsbPacket { report_id, _padding: 0, packet_length: packet.length(), packet, } } pub fn length(&self) -> usize { self.packet_length as usize + 4 } pub fn write(&self, w: &mut W) -> Result<(), std::io::Error> { w.write(&[self.report_id as u8, 0])?; w.write_u16::(self.packet_length)?; self.packet.write(w)?; Ok(()) } }