Add partition-table subcommand

Author: svenstaro

Cargo.lock

@@ -6,7 +6,7 @@ authors = [
     "Jesse Braham <[email protected]>",
 ]
 edition = "2021"
-rust-version = "1.56"
+rust-version = "1.58"
 description = "A command-line tool for flashing Espressif devices over serial"
 repository = "https://github.com/esp-rs/espflash"
 license = "GPL-2.0"
@@ -35,6 +35,7 @@ path = "src/main.rs"
 binread = "2.2"
 bytemuck = { version = "1.9", features = ["derive"] }
 clap = { version = "3.1", features = ["derive"] }
+comfy-table = "5"
 crossterm = "0.23"
 csv = "1.1"
 dialoguer = "0.10"

espflash/Cargo.toml

@@ -346,6 +346,18 @@ pub enum PartitionTableError {
     #[error(transparent)]
     #[diagnostic(transparent)]
     UnalignedPartitionError(#[from] UnalignedPartitionError),
+    #[error(transparent)]
+    #[diagnostic(transparent)]
+    LengthNotMultipleOf32(#[from] LengthNotMultipleOf32),
+    #[error(transparent)]
+    #[diagnostic(transparent)]
+    InvalidChecksum(#[from] InvalidChecksum),
+    #[error(transparent)]
+    #[diagnostic(transparent)]
+    NoEndMarker(#[from] NoEndMarker),
+    #[error(transparent)]
+    #[diagnostic(transparent)]
+    InvalidPartitionTable(#[from] InvalidPartitionTable),
 }
 
 #[derive(Debug, Error, Diagnostic)]
@@ -535,6 +547,26 @@ impl UnalignedPartitionError {
     }
 }
 
+#[derive(Debug, Error, Diagnostic)]
+#[error("Partition table length not a multiple of 32")]
+#[diagnostic(code(espflash::partition_table::invalid_length))]
+pub struct LengthNotMultipleOf32;
+
+#[derive(Debug, Error, Diagnostic)]
+#[error("Checksum invalid")]
+#[diagnostic(code(espflash::partition_table::invalid_checksum))]
+pub struct InvalidChecksum;
+
+#[derive(Debug, Error, Diagnostic)]
+#[error("No end marker found")]
+#[diagnostic(code(espflash::partition_table::no_end_marker))]
+pub struct NoEndMarker;
+
+#[derive(Debug, Error, Diagnostic)]
+#[error("Invalid partition table")]
+#[diagnostic(code(espflash::partition_table::invalid_partition_table))]
+pub struct InvalidPartitionTable;
+
 #[derive(Debug, Error)]
 #[error("{0}")]
 pub struct ElfError(&'static str);

espflash/src/error.rs

@@ -1,7 +1,7 @@
 pub use chip::Chip;
 pub use cli::config::Config;
 pub use elf::{FlashFrequency, FlashMode};
-pub use error::Error;
+pub use error::{Error, InvalidPartitionTable};
 pub use flasher::{FlashSize, Flasher};
 pub use image_format::ImageFormatId;
 pub use partition_table::PartitionTable;

espflash/src/lib.rs

@@ -1,12 +1,12 @@
-use std::{fs, mem::swap, path::PathBuf, str::FromStr};
+use std::{fs, io::Write, mem::swap, path::PathBuf, str::FromStr};
 
 use clap::{IntoApp, Parser};
 use espflash::{
     cli::{
         board_info, connect, flash_elf_image, monitor::monitor, save_elf_as_image, ConnectOpts,
         FlashConfigOpts, FlashOpts,
     },
-    Chip, Config, ImageFormatId,
+    Chip, Config, ImageFormatId, InvalidPartitionTable, PartitionTable,
 };
 use miette::{IntoDiagnostic, Result, WrapErr};
 
@@ -34,6 +34,8 @@ pub enum SubCommand {
     BoardInfo(ConnectOpts),
     /// Save the image to disk instead of flashing to device
     SaveImage(SaveImageOpts),
+    /// Operations for partitions tables
+    PartitionTable(PartitionTableOpts),
 }
 
 #[derive(Parser)]
@@ -60,13 +62,34 @@ pub struct SaveImageOpts {
     pub partition_table: Option<PathBuf>,
 }
 
+#[derive(Parser)]
+pub struct PartitionTableOpts {
+    /// Convert CSV parition table to binary representation
+    #[clap(long, required_unless_present_any = ["info", "to-csv"])]
+    to_binary: bool,
+    /// Convert binary partition table to CSV representation
+    #[clap(long, required_unless_present_any = ["info", "to-binary"])]
+    to_csv: bool,
+    /// Show information on partition table
+    #[clap(short, long, required_unless_present_any = ["to-binary", "to-csv"])]
+    info: bool,
+    /// Input partition table
+    partition_table: PathBuf,
+    /// Optional output file name, if unset will output to stdout
+    #[clap(short, long)]
+    output: Option<PathBuf>,
+}
+
 fn main() -> Result<()> {
     miette::set_panic_hook();
 
     let mut opts = Opts::parse();
     let config = Config::load()?;
 
-    if !matches!(opts.subcommand, Some(SubCommand::BoardInfo(..))) {
+    if !matches!(
+        opts.subcommand,
+        Some(SubCommand::BoardInfo(..) | SubCommand::PartitionTable(..)),
+    ) {
         // If neither the IMAGE nor SERIAL arguments have been provided, print the
         // help message and exit.
         if opts.image.is_none() && opts.connect_opts.serial.is_none() {
@@ -89,6 +112,7 @@ fn main() -> Result<()> {
         match subcommand {
             BoardInfo(opts) => board_info(opts, config),
             SaveImage(opts) => save_image(opts),
+            PartitionTable(opts) => partition_table(opts),
         }
     } else {
         flash(opts, config)
@@ -167,3 +191,50 @@ fn save_image(opts: SaveImageOpts) -> Result<()> {
 
     Ok(())
 }
+
+fn partition_table(opts: PartitionTableOpts) -> Result<()> {
+    if opts.to_binary {
+        let input = fs::read(&opts.partition_table).into_diagnostic()?;
+        let part_table = PartitionTable::try_from_str(String::from_utf8(input).into_diagnostic()?)
+            .into_diagnostic()?;
+
+        // Use either stdout or a file if provided for the output.
+        let mut writer: Box<dyn Write> = if let Some(output) = opts.output {
+            Box::new(fs::File::create(output).into_diagnostic()?)
+        } else {
+            Box::new(std::io::stdout())
+        };
+        part_table.save_bin(&mut writer).into_diagnostic()?;
+    } else if opts.to_csv {
+        let input = fs::read(&opts.partition_table).into_diagnostic()?;
+        let part_table = PartitionTable::try_from_bytes(input).into_diagnostic()?;
+
+        // Use either stdout or a file if provided for the output.
+        let mut writer: Box<dyn Write> = if let Some(output) = opts.output {
+            Box::new(fs::File::create(output).into_diagnostic()?)
+        } else {
+            Box::new(std::io::stdout())
+        };
+        part_table.save_csv(&mut writer).into_diagnostic()?;
+    } else if opts.info {
+        let input = fs::read(&opts.partition_table).into_diagnostic()?;
+
+        // Try getting the partition table from either the csv or the binary representation and
+        // fail otherwise.
+        let part_table = if let Ok(part_table) =
+            PartitionTable::try_from_bytes(input.clone()).into_diagnostic()
+        {
+            part_table
+        } else if let Ok(part_table) =
+            PartitionTable::try_from_str(String::from_utf8(input).into_diagnostic()?)
+        {
+            part_table
+        } else {
+            return Err((InvalidPartitionTable {}).into());
+        };
+
+        part_table.pretty_print();
+    }
+
+    Ok(())
+}