diff --git a/native/src/base/misc.rs b/native/src/base/misc.rs index 07877e380..c270cf7b5 100644 --- a/native/src/base/misc.rs +++ b/native/src/base/misc.rs @@ -1,5 +1,5 @@ use crate::{Utf8CStr, cstr, ffi}; -use argh::EarlyExit; +use argh::{EarlyExit, MissingRequirements}; use libc::c_char; use std::{ fmt, @@ -102,6 +102,51 @@ impl EarlyExitExt for Result { } } +pub struct PositionalArgParser<'a>(pub slice::Iter<'a, &'a str>); + +impl PositionalArgParser<'_> { + pub fn required(&mut self, field_name: &'static str) -> Result { + if let Some(next) = self.0.next() { + Ok(next.to_string()) + } else { + let mut missing = MissingRequirements::default(); + missing.missing_positional_arg(field_name); + missing.err_on_any()?; + unreachable!() + } + } + + pub fn optional(&mut self) -> Option { + self.0.next().map(ToString::to_string) + } + + pub fn last_required(&mut self, field_name: &'static str) -> Result { + let r = self.required(field_name)?; + self.ensure_end()?; + Ok(r) + } + + pub fn last_optional(&mut self) -> Result, EarlyExit> { + let r = self.optional(); + if r.is_none() { + return Ok(r); + } + self.ensure_end()?; + Ok(r) + } + + fn ensure_end(&mut self) -> Result<(), EarlyExit> { + if self.0.len() == 0 { + Ok(()) + } else { + Err(EarlyExit::from(format!( + "Unrecognized argument: {}\n", + self.0.next().unwrap() + ))) + } + } +} + pub struct FmtAdaptor<'a, T>(pub &'a mut T) where T: Write; diff --git a/native/src/boot/cli.rs b/native/src/boot/cli.rs index cab1a2612..cfb4fc2f8 100644 --- a/native/src/boot/cli.rs +++ b/native/src/boot/cli.rs @@ -5,10 +5,10 @@ use crate::ffi::{BootImage, FileFormat, cleanup, repack, split_image_dtb, unpack use crate::patch::hexpatch; use crate::payload::extract_boot_from_payload; use crate::sign::{sha1_hash, sign_boot_image}; -use argh::FromArgs; +use argh::{CommandInfo, EarlyExit, FromArgs, SubCommand}; use base::{ - CmdArgs, EarlyExitExt, LoggedResult, MappedFile, ResultExt, Utf8CStr, WriteExt, - cmdline_logging, cstr, libc, libc::umask, log_err, + CmdArgs, EarlyExitExt, LoggedResult, MappedFile, PositionalArgParser, ResultExt, Utf8CStr, + WriteExt, cmdline_logging, cstr, libc, libc::umask, log_err, }; use std::ffi::c_char; use std::io::{Seek, SeekFrom, Write}; @@ -78,13 +78,28 @@ struct Sign { args: Vec, } -#[derive(FromArgs)] -#[argh(subcommand, name = "extract")] struct Extract { - #[argh(positional)] payload: String, - #[argh(positional)] - args: Vec, + partition: Option, + outfile: Option, +} + +impl FromArgs for Extract { + fn from_args(_command_name: &[&str], args: &[&str]) -> Result { + let mut parse = PositionalArgParser(args.iter()); + Ok(Extract { + payload: parse.required("payload.bin")?, + partition: parse.optional(), + outfile: parse.last_optional()?, + }) + } +} + +impl SubCommand for Extract { + const COMMAND: &'static CommandInfo = &CommandInfo { + name: "extract", + description: "", + }; } #[derive(FromArgs)] @@ -136,26 +151,61 @@ struct Sha1 { #[argh(subcommand, name = "cleanup")] struct Cleanup {} -#[derive(FromArgs)] -#[argh(subcommand, name = "compress")] struct Compress { - #[argh(option, short = 'f', default = r#""gzip".to_string()"#)] - format: String, - #[argh(positional)] + format: FileFormat, file: String, - #[argh(positional)] out: Option, } -#[derive(FromArgs)] -#[argh(subcommand, name = "decompress")] +impl FromArgs for Compress { + fn from_args(command_name: &[&str], args: &[&str]) -> Result { + let cmd = command_name.last().copied().unwrap_or_default(); + let fmt = cmd.strip_prefix("compress=").unwrap_or("gzip"); + + let Ok(fmt) = FileFormat::from_str(fmt) else { + return Err(EarlyExit::from(format!( + "Unsupported or unknown compression format: {fmt}\n" + ))); + }; + + let mut iter = PositionalArgParser(args.iter()); + Ok(Compress { + format: fmt, + file: iter.required("infile")?, + out: iter.last_optional()?, + }) + } +} + +impl SubCommand for Compress { + const COMMAND: &'static CommandInfo = &CommandInfo { + name: "compress", + description: "", + }; +} + struct Decompress { - #[argh(positional)] file: String, - #[argh(positional)] out: Option, } +impl FromArgs for Decompress { + fn from_args(_command_name: &[&str], args: &[&str]) -> Result { + let mut iter = PositionalArgParser(args.iter()); + Ok(Decompress { + file: iter.required("infile")?, + out: iter.last_optional()?, + }) + } +} + +impl SubCommand for Decompress { + const COMMAND: &'static CommandInfo = &CommandInfo { + name: "decompress", + description: "", + }; +} + fn print_usage(cmd: &str) { eprintln!( r#"MagiskBoot - Boot Image Modification Tool @@ -307,17 +357,20 @@ fn boot_main(cmds: CmdArgs) -> LoggedResult { cmds[1] = &cmds[1][2..]; } - if let Some(fmt) = str::strip_prefix(cmds[1], "compress=") { - cmds.insert(1, "compress"); - cmds.insert(2, "-f"); - cmds[3] = fmt; + let mut cli = if cmds[1].starts_with("compress=") { + // Skip the main parser, directly parse the subcommand + Compress::from_args(&cmds[..2], &cmds[2..]).map(|compress| Cli { + action: Action::Compress(compress), + }) + } else { + Cli::from_args(&[cmds[0]], &cmds[1..]) } - - let mut cli = Cli::from_args(&[cmds[0]], &cmds[1..]).on_early_exit(|| match cmds.get(1) { - Some(&"dtb") => print_dtb_usage(), - Some(&"cpio") => print_cpio_usage(), + .on_early_exit(|| match cmds[1] { + "dtb" => print_dtb_usage(), + "cpio" => print_cpio_usage(), _ => print_usage(cmds[0]), }); + match cli.action { Action::Unpack(Unpack { no_decompress, @@ -358,16 +411,13 @@ fn boot_main(cmds: CmdArgs) -> LoggedResult { iter.next().map(Utf8CStr::from_string), )?; } - Action::Extract(Extract { payload, args }) => { - if args.len() > 2 { - log_err!("Too many arguments")?; - } - extract_boot_from_payload( - &payload, - args.first().map(|x| x.as_str()), - args.get(1).map(|x| x.as_str()), - ) - .log_with_msg(|w| w.write_str("Failed to extract from payload"))?; + Action::Extract(Extract { + payload, + partition, + outfile, + }) => { + extract_boot_from_payload(&payload, partition.as_deref(), outfile.as_deref()) + .log_with_msg(|w| w.write_str("Failed to extract from payload"))?; } Action::HexPatch(HexPatch { mut file, @@ -417,15 +467,11 @@ fn boot_main(cmds: CmdArgs) -> LoggedResult { decompress_cmd(&mut file, out.as_mut())?; } Action::Compress(Compress { - ref mut file, - ref format, - ref mut out, + format, + mut file, + mut out, }) => { - compress_cmd( - FileFormat::from_str(format).unwrap_or(FileFormat::UNKNOWN), - file, - out.as_mut(), - )?; + compress_cmd(format, &mut file, out.as_mut())?; } } Ok(0) diff --git a/native/src/boot/compress.rs b/native/src/boot/compress.rs index f91bf153a..9af777489 100644 --- a/native/src/boot/compress.rs +++ b/native/src/boot/compress.rs @@ -1,6 +1,6 @@ use crate::ffi::{FileFormat, check_fmt}; use base::libc::{O_RDONLY, O_TRUNC, O_WRONLY}; -use base::{Chunker, FileOrStd, LoggedResult, Utf8CStr, Utf8CString, WriteExt, error, log_err}; +use base::{Chunker, FileOrStd, LoggedResult, Utf8CStr, Utf8CString, WriteExt, log_err}; use bytemuck::bytes_of_mut; use bzip2::{Compression as BzCompression, write::BzDecoder, write::BzEncoder}; use flate2::{Compression as GzCompression, write::GzEncoder, write::MultiGzDecoder}; @@ -480,10 +480,6 @@ pub(crate) fn compress_cmd( infile: &mut String, outfile: Option<&mut String>, ) -> LoggedResult<()> { - if method == FileFormat::UNKNOWN { - error!("Unsupported compression format"); - } - let infile = Utf8CStr::from_string(infile); let outfile = outfile.map(Utf8CStr::from_string);