From 591faf2bf51e63dd842b8ae218f8fd7b0cd6e474 Mon Sep 17 00:00:00 2001
From: Kai Lueke <kailuke@microsoft.com>
Date: Tue, 16 Jan 2024 19:46:03 +0100
Subject: [PATCH] Add parameters to specify target filename

When a match pattern has to be used, the resulting filename in the
output directory is unknown.
Add parameters to take only the first match and use a particular name
for the extracted target file.
---
 src/bin/download_sysext.rs | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/bin/download_sysext.rs b/src/bin/download_sysext.rs
index 592faaa..4756e07 100644
--- a/src/bin/download_sysext.rs
+++ b/src/bin/download_sysext.rs
@@ -268,7 +268,7 @@ where
     })
 }
 
-fn do_download_verify(pkg: &mut Package<'_>, output_dir: &Path, unverified_dir: &Path, pubkey_file: &str, client: &Client, print_progress: bool) -> Result<()> {
+fn do_download_verify(pkg: &mut Package<'_>, output_filename: Option<String>, output_dir: &Path, unverified_dir: &Path, pubkey_file: &str, client: &Client, print_progress: bool) -> Result<()> {
     pkg.check_download(unverified_dir)?;
 
     pkg.download(unverified_dir, client, print_progress).context(format!("unable to download \"{:?}\"", pkg.name))?;
@@ -276,7 +276,7 @@ fn do_download_verify(pkg: &mut Package<'_>, output_dir: &Path, unverified_dir:
     // Unverified payload is stored in e.g. "output_dir/.unverified/oem.gz".
     // Verified payload is stored in e.g. "output_dir/oem.raw".
     let pkg_unverified = unverified_dir.join(&*pkg.name);
-    let pkg_verified = output_dir.join(pkg_unverified.with_extension("raw").file_name().unwrap_or_default());
+    let pkg_verified = output_dir.join(output_filename.as_ref().map(OsStr::new).unwrap_or(pkg_unverified.with_extension("raw").file_name().unwrap_or_default()));
 
     let datablobspath = pkg.verify_signature_on_disk(&pkg_unverified, pubkey_file).context(format!("unable to verify signature \"{}\"", pkg.name))?;
 
@@ -295,6 +295,10 @@ struct Args {
     #[argh(option, short = 'o')]
     output_dir: String,
 
+    /// target filename in directory
+    #[argh(option, short = 'n')]
+    target_filename: Option<String>,
+
     /// path to the Omaha XML file, or - to read from stdin
     #[argh(option, short = 'i')]
     input_xml: Option<String>,
@@ -312,6 +316,10 @@ struct Args {
     #[argh(option, short = 'm')]
     image_match: Vec<String>,
 
+    /// only take the first matching entry
+    #[argh(switch, short = 't')]
+    take_first_match: bool,
+
     /// report download progress
     #[argh(switch, short = 'v')]
     print_progress: bool,
@@ -338,6 +346,10 @@ fn main() -> Result<(), Box<dyn Error>> {
     let args: Args = argh::from_env();
     println!("{:?}", args);
 
+    if args.payload_url.is_none() && !args.take_first_match && args.target_filename.is_some() {
+        return Err("--target-filename can only be specified with --take-first-match".into());
+    }
+
     let glob_set = args.image_match_glob_set()?;
 
     let output_dir = Path::new(&*args.output_dir);
@@ -391,6 +403,7 @@ fn main() -> Result<(), Box<dyn Error>> {
             )?;
             do_download_verify(
                 &mut pkg_fake,
+                args.target_filename.clone(),
                 output_dir,
                 unverified_dir.as_path(),
                 args.pubkey_file.as_str(),
@@ -424,12 +437,16 @@ fn main() -> Result<(), Box<dyn Error>> {
     for pkg in pkgs_to_dl.iter_mut() {
         do_download_verify(
             pkg,
+            args.target_filename.clone(),
             output_dir,
             unverified_dir.as_path(),
             args.pubkey_file.as_str(),
             &client,
             args.print_progress,
         )?;
+        if args.take_first_match {
+            break;
+        }
     }
 
     // clean up data