diff --git a/src/main.rs b/src/main.rs index a1c96cf..8132804 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,6 +17,13 @@ fn main() { // Skip the first arg which is the calling application name. let opt = Opt::from_iter(std::env::args().skip(1)); + if opt.list_chips { + for vendor in vendor_map() { + println!("{}", vendor.0); + } + return; + } + // Try and get the cargo project information. let project = cargo_project::Project::query(".").expect("Couldn't parse the Cargo.toml"); @@ -51,7 +58,7 @@ fn main() { // todo, keep as iter. difficult because we want to filter map remove two items at once. // Remove our args as cargo build does not understand them. - let flags = ["--pid", "--vid"].iter(); + let flags = ["--pid", "--vid", "--chip"].iter(); for flag in flags { if let Some(index) = args.iter().position(|x| x == flag) { args.remove(index); @@ -76,6 +83,22 @@ fn main() { let d = if let (Some(v), Some(p)) = (opt.vid, opt.pid) { open_device_with_vid_pid(v, p) .expect("Are you sure device is plugged in and in bootloader mode?") + } else if let Some(c) = opt.chip { + println!(" {} for a connected {}.", "Searching".green().bold(), c); + + let mut device: Option> = None; + + let vendor = vendor_map(); + + if let Some(products) = vendor.get(&c) { + for (v, p) in products { + if let Some(d) = open_device_with_vid_pid(*v, *p) { + device = Some(d); + break; + } + } + } + device.expect("Are you sure device is plugged in and in bootloader mode?") } else { println!( " {} for a connected device with known vid/pid pair.", @@ -90,11 +113,9 @@ fn main() { let mut device: Option> = None; - let vendor = vendor_map(); - for d in devices { - if let Some(products) = vendor.get(&d.vendor_id()) { - if products.contains(&d.product_id()) { + for vendor in vendor_map() { + if vendor.1.contains(&(d.vendor_id(), d.product_id())) { if let Some(d) = open_device_with_vid_pid(d.vendor_id(), d.product_id()) { device = Some(d); break; @@ -107,7 +128,7 @@ fn main() { println!( " {} {} {}", - "Trying ".green().bold(), + "Found ".green().bold(), d.read_manufacturer_string_ascii(&d.device().device_descriptor().unwrap()) .unwrap(), d.read_product_string_ascii(&d.device().device_descriptor().unwrap()) @@ -182,4 +203,9 @@ struct Opt { pid: Option, #[structopt(name = "vid", long = "vid", parse(try_from_str = parse_hex_16))] vid: Option, + + #[structopt(name = "chip", long = "chip")] + chip: Option, + #[structopt(name = "list-chips", long = "list-chips")] + list_chips: bool, } diff --git a/src/utils.rs b/src/utils.rs index a0d4be5..565359b 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -8,14 +8,15 @@ use std::{fs::File, io::Read}; pub enum UtilError { Elf(goblin::error::Error), Dfu(dfu::error::Error), - File(std::io::Error) + File(std::io::Error), } /// Returns a contiguous bin with 0s between non-contiguous sections and starting address from an elf. pub fn elf_to_bin(path: PathBuf) -> Result<(Vec, u32), UtilError> { let mut file = File::open(path).map_err(|e| UtilError::File(e))?; let mut buffer = vec![]; - file.read_to_end(&mut buffer).map_err(|e| UtilError::File(e))?; + file.read_to_end(&mut buffer) + .map_err(|e| UtilError::File(e))?; let binary = goblin::elf::Elf::parse(buffer.as_slice()).map_err(|e| UtilError::Elf(e))?; @@ -58,9 +59,11 @@ pub fn flash_bin( address: u32, d: &rusb::Device, ) -> Result<(), UtilError> { - let mut dfu = dfu::Dfu::from_bus_device(d.bus_number(), d.address(), 0_u32, 0_u32).map_err(|e| UtilError::Dfu(e))?; + let mut dfu = dfu::Dfu::from_bus_device(d.bus_number(), d.address(), 0_u32, 0_u32) + .map_err(|e| UtilError::Dfu(e))?; if binary.len() < 2048 { - dfu.write_flash_from_slice(address, binary).map_err(|e| UtilError::Dfu(e))?; + dfu.write_flash_from_slice(address, binary) + .map_err(|e| UtilError::Dfu(e))?; } else { // hacky bug workaround std::fs::write("target/out.bin", binary).map_err(|e| UtilError::File(e))?; @@ -71,15 +74,17 @@ pub fn flash_bin( .map_err(|e| UtilError::File(e))?, address, None, - ).map_err(|e| UtilError::Dfu(e))?; + ) + .map_err(|e| UtilError::Dfu(e))?; std::fs::remove_file("target/out.bin").map_err(|e| UtilError::File(e))?; } Ok(()) } -pub fn vendor_map() -> std::collections::HashMap> { +pub fn vendor_map() -> std::collections::HashMap> { maplit::hashmap! { - 0x0483 => vec![0xdf11], + "stm32f4".to_string() => vec![(0x0483, 0xdf11)], + "gd32vf103".to_string() => vec![(0x28e9, 0x0189)], } }