use crate::sc64; use chrono::Local; use colored::Colorize; use encoding_rs::EUC_JP; use std::{ fs::File, io::{stdin, Read, Write}, path::PathBuf, sync::mpsc::{channel, Receiver, Sender}, thread::spawn, }; pub enum Encoding { UTF8, EUCJP, } pub struct Handler { header: Option>, line_rx: Receiver, encoding: Encoding, } enum DataType { Text, RawBinary, Header, Screenshot, Heartbeat, Unknown, } impl From for DataType { fn from(value: u8) -> Self { match value { 0x01 => Self::Text, 0x02 => Self::RawBinary, 0x03 => Self::Header, 0x04 => Self::Screenshot, 0x05 => Self::Heartbeat, _ => Self::Unknown, } } } impl From for u8 { fn from(value: DataType) -> Self { match value { DataType::Text => 0x01, DataType::RawBinary => 0x02, DataType::Header => 0x03, DataType::Screenshot => 0x04, DataType::Heartbeat => 0x05, DataType::Unknown => 0xFF, } } } impl From for u32 { fn from(value: DataType) -> Self { u8::from(value) as u32 } } #[derive(Clone, Copy)] enum ScreenshotPixelFormat { Rgba16, Rgba32, } impl TryFrom for ScreenshotPixelFormat { type Error = String; fn try_from(value: u32) -> Result { Ok(match value { 2 => Self::Rgba16, 4 => Self::Rgba32, _ => return Err("Invalid pixel format for screenshot metadata".into()), }) } } impl From for u32 { fn from(value: ScreenshotPixelFormat) -> Self { match value { ScreenshotPixelFormat::Rgba16 => 2, ScreenshotPixelFormat::Rgba32 => 4, } } } struct ScreenshotMetadata { format: ScreenshotPixelFormat, width: u32, height: u32, } impl TryFrom> for ScreenshotMetadata { type Error = String; fn try_from(value: Vec) -> Result { if value.len() != 16 { return Err("Invalid header length for screenshot metadata".into()); } if u32::from_be_bytes(value[0..4].try_into().unwrap()) != DataType::Screenshot.into() { return Err("Invalid header datatype for screenshot metadata".into()); } let format = u32::from_be_bytes(value[4..8].try_into().unwrap()); let width = u32::from_be_bytes(value[8..12].try_into().unwrap()); let height = u32::from_be_bytes(value[12..16].try_into().unwrap()); if width > 4096 || height > 4096 { return Err("Invalid width or height for screenshot metadata".into()); } Ok(ScreenshotMetadata { format: format.try_into()?, width, height, }) } } struct Heartbeat { usb_protocol: u16, version: u16, } impl TryFrom<&[u8]> for Heartbeat { type Error = String; fn try_from(value: &[u8]) -> Result { if value.len() < 4 { return Err("Invalid heartbeat data length".into()); } let usb_protocol = u16::from_be_bytes(value[0..2].try_into().unwrap()); let version = u16::from_be_bytes(value[2..4].try_into().unwrap()); Ok(Heartbeat { usb_protocol, version, }) } } macro_rules! success { ($($a: tt)*) => { println!("{}", format!($($a)*).bright_blue()); }; } macro_rules! error { ($($a: tt)*) => {{ println!("{}", format!("Error: {}", format!($($a)*)).bright_red()); }}; } macro_rules! stop { ($r: expr, $($a: tt)*) => {{ error!($($a)*); $r }}; } const MAX_PACKET_LENGTH: usize = 8 * 1024 * 1024; const SUPPORTED_USB_PROTOCOL_VERSION: u16 = 2; impl Handler { pub fn set_text_encoding(&mut self, encoding: Encoding) { self.encoding = encoding; } pub fn process_user_input(&self) -> Option { let line = match self.line_rx.try_recv() { Ok(line) => { if line.len() == 0 { return None; } else { line } } Err(_) => return None, }; let token_count = line.matches("@").count(); if (token_count % 2) != 0 { return stop!(None, "Missing closing '@' token"); } let packet = if token_count == 2 && line.starts_with("@") && line.ends_with("@") { sc64::DebugPacket { datatype: DataType::RawBinary.into(), data: match load_file(line.trim_matches('@')) { Ok(data) => data, Err(error) => return stop!(None, "{error}"), }, } } else { let mut is_text = true; let mut path = String::new(); let mut character_buffer = vec![0u8; 4]; let mut data = vec![0u8; 0]; for character in line.chars() { if character == '@' { if is_text { is_text = false; } else { let mut file = match load_file(&path) { Ok(data) => data, Err(error) => return stop!(None, "{error}"), }; let length = file.len(); data.append(&mut format!("@{length}@").into_bytes()); data.append(&mut file); is_text = true; path = String::new(); } } else { if is_text { let encoded = character.encode_utf8(&mut character_buffer); data.append(&mut encoded.as_bytes().to_vec()); } else { path.push(character); } } } data.append(&mut b"\0".to_vec()); sc64::DebugPacket { datatype: DataType::Text.into(), data, } }; if packet.data.len() > MAX_PACKET_LENGTH { return stop!( None, "Debug packet size too big ({}), exceeding maximum size of {}", packet.data.len(), MAX_PACKET_LENGTH ); } Some(packet) } pub fn handle_debug_packet(&mut self, debug_packet: sc64::DebugPacket) { let sc64::DebugPacket { datatype, data } = debug_packet; match datatype.into() { DataType::Text => self.handle_datatype_text(&data), DataType::RawBinary => self.handle_datatype_raw_binary(&data), DataType::Header => self.handle_datatype_header(&data), DataType::Screenshot => self.handle_datatype_screenshot(&data), DataType::Heartbeat => self.handle_datatype_heartbeat(&data), _ => error!("Received unknown debug packet datatype: 0x{datatype:02X}"), } } pub fn handle_is_viewer_64(&self, data: &[u8]) { self.print_text(data); } pub fn handle_save_writeback( &self, save_writeback: sc64::SaveWriteback, path: &Option, ) { let filename = &if let Some(path) = path { path.to_string_lossy().to_string() } else { generate_filename( "save", match save_writeback.save { sc64::SaveType::Eeprom4k | sc64::SaveType::Eeprom16k => "eep", sc64::SaveType::Sram | sc64::SaveType::SramBanked | sc64::SaveType::Sram1m => { "srm" } sc64::SaveType::Flashram => "fla", _ => "sav", }, ) }; match File::create(filename) { Ok(mut file) => { if let Err(error) = file.write_all(&save_writeback.data) { error!("Couldn't write save [{filename}]: {error}"); } else { success!("Wrote [{}] save to [{filename}]", save_writeback.save); } } Err(error) => error!("Couldn't create save writeback file [{filename}]: {error}"), } } pub fn handle_data_flushed(&self) { error!("Debug data write dropped due to timeout"); } fn handle_datatype_text(&self, data: &[u8]) { self.print_text(data); } fn handle_datatype_raw_binary(&self, data: &[u8]) { let filename = &generate_filename("binaryout", "bin"); match File::create(filename) { Ok(mut file) => { if let Err(error) = file.write_all(data) { error!("Couldn't write raw binary [{filename}]: {error}"); } else { success!("Wrote [{}] bytes to [{filename}]", data.len()); } } Err(error) => error!("Couldn't create raw binary file [{filename}]: {error}"), } } fn handle_datatype_header(&mut self, data: &[u8]) { self.header = Some(data.to_vec()); } fn handle_datatype_screenshot(&mut self, data: &[u8]) { let header = match self.header.take() { Some(header) => header, None => return error!("Got screenshot packet without header data"), }; let ScreenshotMetadata { format, height, width, } = match header.try_into() { Ok(data) => data, Err(error) => return error!("{error}"), }; let format_size: u32 = format.into(); if data.len() as u32 != format_size * width * height { return error!("Data length did not match header data for screenshot datatype"); } let mut image = image::RgbaImage::new(width, height); for (x, y, pixel) in image.enumerate_pixels_mut() { let location = ((x + (y * width)) * format_size) as usize; let p = &data[location..location + format_size as usize]; pixel.0 = match format { ScreenshotPixelFormat::Rgba16 => { let r = ((p[0] >> 3) & 0x1F) << 3; let g = (((p[0] & 0x07) << 2) | ((p[1] >> 6) & 0x03)) << 3; let b = ((p[1] >> 1) & 0x1F) << 3; let a = ((p[1]) & 0x01) * 255; [r, g, b, a] } ScreenshotPixelFormat::Rgba32 => [p[0], p[1], p[2], p[3]], } } let filename = &generate_filename("screenshot", "png"); if let Some(error) = image.save(filename).err() { return error!("Couldn't save screenshot [{filename}]: {error}"); } success!("Wrote {width}x{height} pixels to [{filename}]"); } fn handle_datatype_heartbeat(&mut self, data: &[u8]) { let Heartbeat { usb_protocol, version, } = match data.try_into() { Ok(heartbeat) => heartbeat, Err(error) => return error!("Error while parsing heartbeat datatype: {error}"), }; if usb_protocol > SUPPORTED_USB_PROTOCOL_VERSION { return error!("Unsupported USB protocol version: {usb_protocol}"); } match version { 1 => {} _ => return error!("Unsupported USB heartbeat version: {version}"), } } fn print_text(&self, data: &[u8]) { match self.encoding { Encoding::UTF8 => print!("{}", String::from_utf8_lossy(&data)), Encoding::EUCJP => print!("{}", EUC_JP.decode(&data).0), } } } pub fn new() -> Handler { let (line_tx, line_rx) = channel::(); spawn(move || stdin_thread(line_tx)); Handler { header: None, line_rx, encoding: Encoding::UTF8, } } fn load_file(path: &str) -> Result, String> { if path.len() == 0 { return Err(format!("Couldn't open file: Specified path is empty")); } let mut file = match File::open(path) { Ok(file) => file, Err(error) => return Err(format!("Couldn't open file [{path}]: {error}")), }; let length = match file.metadata() { Ok(metadata) => metadata.len(), Err(error) => return Err(format!("Couldn't get file [{path}] length: {error}")), }; if length > MAX_PACKET_LENGTH as u64 { return Err(format!("File [{path}] size too big")); } let mut data = vec![0u8; length as usize]; match file.read_exact(&mut data) { Ok(()) => Ok(data), Err(error) => Err(format!("Couldn't read file [{path}] contents: {error}")), } } fn generate_filename(prefix: &str, extension: &str) -> String { format!( "{prefix}-{}.{extension}", Local::now().format("%y%m%d%H%M%S.%f") ) } fn stdin_thread(line_tx: Sender) { loop { let mut line = String::new(); if stdin().read_line(&mut line).is_ok() { if line_tx.send(line.trim_end().to_string()).is_err() { return; } } } }