Consolidate logic for determining passphrase source priority

avbroot/src/cli/avb.rs

@@ -242,14 +242,12 @@ fn sign_or_clear(info: &mut AvbInfo, orig_header: &Header, key_group: &KeyGroup)
                 bail!("Need to sign new AVB header, but no private key was specified");
             };
 
-            let passphrase = if let Some(v) = &key_group.pass_env_var {
-                PassphraseSource::EnvVar(v.clone())
-            } else if let Some(p) = &key_group.pass_file {
-                PassphraseSource::File(p.clone())
-            } else {
-                PassphraseSource::Prompt(format!("Enter passphrase for {key_path:?}: "))
-            };
-            let private_key = crypto::read_pem_key_file(key_path, &passphrase)
+            let source = PassphraseSource::new(
+                key_path,
+                key_group.pass_file.as_deref(),
+                key_group.pass_env_var.as_deref(),
+            );
+            let private_key = crypto::read_pem_key_file(key_path, &source)
                 .with_context(|| format!("Failed to load key: {key_path:?}"))?;
 
             info.header.set_algo_for_key(&private_key)?;

avbroot/src/cli/key.rs

@@ -18,29 +18,27 @@ use crate::{
     format::avb,
 };
 
-fn get_passphrase(group: &PassphraseGroup, key_path: &Path) -> PassphraseSource {
-    if let Some(v) = &group.pass_env_var {
-        PassphraseSource::EnvVar(v.clone())
-    } else if let Some(p) = &group.pass_file {
-        PassphraseSource::File(p.clone())
-    } else {
-        PassphraseSource::Prompt(format!("Enter passphrase for {key_path:?}: "))
-    }
+fn get_passphrase_source(group: &PassphraseGroup, key_path: &Path) -> PassphraseSource {
+    PassphraseSource::new(
+        key_path,
+        group.pass_file.as_deref(),
+        group.pass_env_var.as_deref(),
+    )
 }
 
 pub fn key_main(cli: &KeyCli) -> Result<()> {
     match &cli.command {
         KeyCommand::GenerateKey(c) => {
-            let passphrase = get_passphrase(&c.passphrase, &c.output);
+            let source = get_passphrase_source(&c.passphrase, &c.output);
             let private_key =
                 crypto::generate_rsa_key_pair().context("Failed to generate RSA keypair")?;
 
-            crypto::write_pem_key_file(&c.output, &private_key, &passphrase)
+            crypto::write_pem_key_file(&c.output, &private_key, &source)
                 .with_context(|| format!("Failed to write private key: {:?}", c.output))?;
         }
         KeyCommand::GenerateCert(c) => {
-            let passphrase = get_passphrase(&c.passphrase, &c.key);
-            let private_key = crypto::read_pem_key_file(&c.key, &passphrase)
+            let source = get_passphrase_source(&c.passphrase, &c.key);
+            let private_key = crypto::read_pem_key_file(&c.key, &source)
                 .with_context(|| format!("Failed to load key: {:?}", c.key))?;
 
             let validity = Duration::from_secs(c.validity * 24 * 60 * 60);
@@ -52,7 +50,7 @@ pub fn key_main(cli: &KeyCli) -> Result<()> {
         }
         KeyCommand::ExtractAvb(c) => {
             let public_key = if let Some(p) = &c.input.key {
-                let passphrase = get_passphrase(&c.passphrase, p);
+                let passphrase = get_passphrase_source(&c.passphrase, p);
                 let private_key = crypto::read_pem_key_file(p, &passphrase)
                     .with_context(|| format!("Failed to load key: {p:?}"))?;
 

avbroot/src/cli/ota.rs

@@ -846,24 +846,20 @@ pub fn patch_subcommand(cli: &PatchCli, cancel_signal: &AtomicBool) -> Result<()
         Cow::Borrowed,
     );
 
-    let passphrase_avb = if let Some(v) = &cli.pass_avb_env_var {
-        PassphraseSource::EnvVar(v.clone())
-    } else if let Some(p) = &cli.pass_avb_file {
-        PassphraseSource::File(p.clone())
-    } else {
-        PassphraseSource::Prompt(format!("Enter passphrase for {:?}: ", cli.key_avb))
-    };
-    let passphrase_ota = if let Some(v) = &cli.pass_ota_env_var {
-        PassphraseSource::EnvVar(v.clone())
-    } else if let Some(p) = &cli.pass_ota_file {
-        PassphraseSource::File(p.clone())
-    } else {
-        PassphraseSource::Prompt(format!("Enter passphrase for {:?}: ", cli.key_ota))
-    };
+    let source_avb = PassphraseSource::new(
+        &cli.key_avb,
+        cli.pass_avb_file.as_deref(),
+        cli.pass_avb_env_var.as_deref(),
+    );
+    let source_ota = PassphraseSource::new(
+        &cli.key_ota,
+        cli.pass_ota_file.as_deref(),
+        cli.pass_ota_env_var.as_deref(),
+    );
 
-    let key_avb = crypto::read_pem_key_file(&cli.key_avb, &passphrase_avb)
+    let key_avb = crypto::read_pem_key_file(&cli.key_avb, &source_avb)
         .with_context(|| format!("Failed to load key: {:?}", cli.key_avb))?;
-    let key_ota = crypto::read_pem_key_file(&cli.key_ota, &passphrase_ota)
+    let key_ota = crypto::read_pem_key_file(&cli.key_ota, &source_ota)
         .with_context(|| format!("Failed to load key: {:?}", cli.key_ota))?;
     let cert_ota = crypto::read_pem_cert_file(&cli.cert_ota)
         .with_context(|| format!("Failed to load certificate: {:?}", cli.cert_ota))?;

avbroot/src/crypto.rs

@@ -5,7 +5,7 @@
 
 use std::{
     env::{self, VarError},
-    ffi::OsString,
+    ffi::{OsStr, OsString},
     fs::{self, File, OpenOptions},
     io::{self, BufReader, BufWriter, Read, Write},
     path::{Path, PathBuf},
@@ -75,6 +75,16 @@ pub enum PassphraseSource {
 }
 
 impl PassphraseSource {
+    pub fn new(key_file: &Path, pass_file: Option<&Path>, env_var: Option<&OsStr>) -> Self {
+        if let Some(v) = env_var {
+            Self::EnvVar(v.to_owned())
+        } else if let Some(p) = pass_file {
+            Self::File(p.to_owned())
+        } else {
+            Self::Prompt(format!("Enter passphrase for {key_file:?}: "))
+        }
+    }
+
     pub fn acquire(&self, confirm: bool) -> Result<String> {
         let passphrase = match self {
             Self::Prompt(p) => {