Add mismatched filter option (#1)

* Add mismatched filter option

* Bump version

* Update README sample output

* More soft changes

* Add recommended feature flag

* Add .envrc

.envrc

@@ -0,0 +1 @@
+use flake

.gitignore

@@ -1,6 +1,5 @@
 Cargo.lock
 .direnv
-.envrc
 target
 test.txt
 matching-edeks.txt

CHANGELOG.md

@@ -1,8 +1,12 @@
 # Changelog
 
+## 1.2.0
+
+- Added a `-m --mismatched` flag that filters the output to only EDEKs that need a rekey because the KMS config ID in their header and the leased key used to encrypt them is mismatched. Must be rekeyed with TSP 4.11.1+ to fix.
+
 ## 1.1.0
 
-- Added `-v --verbose` flag that outputs full result tuples. Lack of the flag outputs only raw identifiers, one per line.
+- Added a `-v --verbose` flag that outputs full result tuples. Lack of the flag outputs only raw identifiers, one per line.
 - Added a `broken-edeks.txt` output file for EDEKs that couldn't be parsed. Contains one line per broken EDEK of
   `("identifier", "edek_data", "error message")` if `--verbose` is enabled, only the raw identifiers otherwise.
 

Cargo.toml

@@ -1,9 +1,9 @@
 [package]
 name = "search-edeks"
-version = "1.1.0"
+version = "1.2.0"
 authors = ["IronCore Labs <[email protected]>"]
 categories = ["utilities"]
-description = "Tool to search EDEK's protobuf. Can be used to find EDEKs that need to be rekeyed from an old KMS config ID."
+description = "Tool to search IronCoreLabs Tenant Security Proxy EDEK's protobuf."
 edition = "2021"
 license = "AGPL-3.0-only"
 readme = "README.md"
@@ -13,9 +13,14 @@ repository = "https://github.com/IronCoreLabs/search-edeks"
 [dependencies]
 base64 = "~0.21"
 bytes = "1.4.0"
-clap = { version = "~3", features = ["cargo", "derive", "suggestions"] }
+clap = { version = "~4", features = [
+  "cargo",
+  "derive",
+  "suggestions",
+  "wrap_help",
+] }
 hex = "0.4.3"
-protobuf = {version = "3.2", features = ["with-bytes"]}
+protobuf = { version = "3.2", features = ["with-bytes"] }
 ron = "0.8.0"
 serde = { version = "~1.0", features = ["derive"] }
 

README.md

@@ -16,28 +16,26 @@ Check out this repo and run `cargo b --release`. The binary will be at `target/r
 
 ```console
 search-edeks --help
-search-edeks 1.1.0
-IronCore Labs <[email protected]>
-Tool to search EDEK's protobuf. Can be used to find EDEKs that need to be rekeyed from an old KMS
-config ID.
-
-USAGE:
-    search-edeks [OPTIONS] --id <VALUE> --file <FILE> <--hex|--base64>
-
-OPTIONS:
-    -b, --base64         Consume and output base64 formatted EDEKs
-    -d, --debug          Print extra debug information
-    -f, --file <FILE>    File with one `("identifier", "EDEK")` per line
-    -h, --hex            Consume and output hex formatted EDEKs
-        --help           Print help information
-    -i, --id <VALUE>     Sets the KMS config ID we're searching for
-    -v, --verbose        Output identifier and original EDEK (and error message if applicable). If
-                         not enabled, only identifiers will be output
-    -V, --version        Print version information
+Tool to search IronCoreLabs Tenant Security Proxy EDEK's protobuf.
+
+Usage: search-edeks [OPTIONS] --file <FILE> <--id <VALUE>|--mismatched> <--hex|--base64>
+
+Options:
+  -i, --id <VALUE>   Sets the KMS config ID we're searching for
+  -m, --mismatched   Searches for mismatches between the KMS config ID in the EDEK header and the leased key used to encrypt the EDEK. Resulting EDEKs must be rekeyed with TSP 4.11.1+ to repair.
+  -f, --file <FILE>  File with one `("identifier", "EDEK")` per line
+  -h, --hex          Consume and output hex formatted EDEKs
+  -b, --base64       Consume and output base64 formatted EDEKs
+  -d, --debug        Print extra debug information
+  -v, --verbose      Output identifier and original EDEK (and error message if applicable). If not enabled, only identifiers will be output
+  -h, --help         Print help
+  -V, --version      Print version
 ```
 
 For example `search-edeks --file edeks.txt --id 1201 --hex` would search `edeks.txt` for any EDEKs that were created using KMS config ID `1201`. It would output `matching-edeks.txt` with the one identifier per line for each EDEK that matched. It would output `broken-edeks.txt` with one identifier per line for each EDEK that wasn't parsable as an EDEK. If `--verbose` was enabled, the output would be tuples of the required input form (with the broken EDEKs additonally containing an error message).
 
+If multiple search filters are included, all must be present for an EDEK to match.
+
 ## Releasing
 
 * update the version in Cargo.toml according to semver before tagging for release

src/filters.rs

@@ -0,0 +1,30 @@
+use crate::{proto::transform::EncryptedDek, util::edek_from_bytes};
+
+#[derive(Clone, Debug)]
+pub(crate) enum Filter {
+    ConfigId(i32),
+    Mismatched,
+}
+
+fn execute_config_id_filter(
+    parsed_edek: &EncryptedDek,
+    config_id_to_match: i32,
+) -> Result<bool, String> {
+    Ok(parsed_edek.kmsConfigId == config_id_to_match)
+}
+fn execute_mismatched_filter(parsed_edek: &EncryptedDek) -> Result<bool, String> {
+    if !parsed_edek.encryptedLeasedKeyData.is_empty() {
+        match edek_from_bytes(&parsed_edek.encryptedLeasedKeyData) {
+            Ok(lk_edek) => Ok(lk_edek.kmsConfigId != parsed_edek.kmsConfigId),
+            Err(e) => Err(format!("Failed to parse leased key: {e}")),
+        }
+    } else {
+        Ok(false)
+    }
+}
+pub(crate) fn execute_filter(filter: &Filter, parsed_edek: &EncryptedDek) -> Result<bool, String> {
+    match filter {
+        Filter::ConfigId(config_id) => execute_config_id_filter(parsed_edek, *config_id),
+        Filter::Mismatched => execute_mismatched_filter(parsed_edek),
+    }
+}

src/main.rs

@@ -1,27 +1,37 @@
+mod filters;
 mod util;
 mod proto {
     include!(concat!(env!("OUT_DIR"), "/proto/mod.rs"));
 }
 
 use base64::{engine::general_purpose::STANDARD, Engine};
+use clap::{arg, command, value_parser, ArgAction, ArgGroup};
+use filters::{execute_filter, Filter};
 use std::{
+    convert::identity,
     fs::File,
     io::{BufRead, BufReader},
     path::{Path, PathBuf},
 };
 use util::{edek_from_bytes, write_file, EdekFileEntry};
 
-use clap::{arg, command, value_parser, ArgAction, ArgGroup};
-
 fn main() {
     let matches = command!()
         .arg(
             arg!(
                 -i --id <VALUE> "Sets the KMS config ID we're searching for"
             )
-            .required(true)
             .value_parser(value_parser!(i32)),
         )
+        .arg(arg!(
+            -m --mismatched "Searches for mismatches between the KMS config ID in the EDEK header and the leased key used to encrypt the EDEK. Resulting EDEKs must be rekeyed with TSP 4.11.1+ to repair."
+        ).action(ArgAction::SetTrue))
+        .group(
+            ArgGroup::new("search")
+                .required(true)
+                .args(&["id", "mismatched"])
+                .multiple(true),
+        )
         .arg(
             arg!(
                 -f --file <FILE> r#"File with one `("identifier", "EDEK")` per line"#
@@ -43,8 +53,15 @@ fn main() {
         .arg(arg!(-v --verbose "Output identifier and original EDEK (and error message if applicable). If not enabled, only identifiers will be output").action(ArgAction::SetTrue))
         .get_matches();
 
-    // get our required arguments
-    let config_id = matches.get_one::<i32>("id").expect("id is required");
+    // filter args
+    let config_id_option = matches.get_one::<i32>("id").map(|c| Filter::ConfigId(*c));
+    let mismatched = if matches.get_flag("mismatched") {
+        Some(Filter::Mismatched)
+    } else {
+        None
+    };
+
+    // io args
     let edek_file_path = matches
         .get_one::<PathBuf>("file")
         .expect("file is required");
@@ -53,35 +70,41 @@ fn main() {
     let debug: bool = matches.get_flag("debug");
     let verbose: bool = matches.get_flag("verbose");
 
-    // read the edek tuples line by line
+    // open the edek tuples file handle, later to be read as a buffered streaming iterator
     let edek_file = File::open(edek_file_path)
-        .map_err(|e| format!("Failed to open EDEK file: {}", e))
+        .map_err(|e| format!("Failed to open EDEK file: {e}"))
         .unwrap();
     // zip an index in so we can give line numbers
     let edek_lines = BufReader::new(edek_file).lines().zip(1..);
+    // if we ever hit memory issues while writing the whole vec out of memory at once, we could write directly to shared write handlers
     let mut found_lines: Vec<EdekFileEntry> = vec![];
     let mut found_broken: Vec<(String, String, String)> = vec![];
+    // a list of filters to attempt to match
+    let active_filters: Vec<Filter> = vec![config_id_option, mismatched]
+        .into_iter()
+        .filter_map(identity)
+        .collect();
     for line in edek_lines {
         if let (Ok(edek_entry), line_number) = line {
             if debug {
-                println!("edek string: {}", edek_entry);
+                println!("edek string: {edek_entry}");
             };
             let (identifier, edek) = ron::from_str::<EdekFileEntry>(&edek_entry)
-                .map_err(|e| format!("Unexpected error processing line {}: {}", line_number, e))
+                .map_err(|e| format!("Unexpected error processing line {line_number}: {e}"))
                 .unwrap();
             // decode the edek string in the desired format
             // if we fail, log it, but keep going in case there are lines that match
             let decode_attempt = if use_base64 {
                 STANDARD
                     .decode(edek.clone())
-                    .map_err(|e| format!("EDEK was not base64: {}", e))
+                    .map_err(|e| format!("EDEK was not base64: {e}"))
             } else if use_hex {
                 let stripped = if edek.starts_with("0x") || edek.starts_with("0X") {
                     edek.chars().skip(2).collect()
                 } else {
                     edek.clone()
                 };
-                hex::decode(stripped).map_err(|e| format!("EDEK was not hex: {}", e))
+                hex::decode(stripped).map_err(|e| format!("EDEK was not hex: {e}"))
             } else {
                 // this should've already been handled by clap, but writing again so Rust is happy
                 panic!("Base64 or Hex format must be specified.");
@@ -94,35 +117,56 @@ fn main() {
                     match parse_attempt {
                         Ok(parsed_edek) => {
                             if debug {
-                                println!("parsed proto: {}, line {}", parsed_edek, line_number);
+                                println!("parsed proto: {parsed_edek}, line {line_number}");
                             }
-                            // do the actual comparison we care about
-                            if parsed_edek.kmsConfigId == *config_id {
-                                found_lines.push((identifier, edek));
+                            let (matched_filters, failures): (Vec<_>, Vec<_>) = active_filters
+                                .iter()
+                                .map(|filter| execute_filter(filter, &parsed_edek))
+                                .partition(Result::is_ok);
+
+                            if !failures.is_empty() {
+                                let message = failures.into_iter().map(Result::unwrap_err).fold(
+                                    format!("Failure on line {line_number}: "),
+                                    |acc, failure| format!("{} {}", acc, failure),
+                                );
+                                if debug {
+                                    println!("WARNING: {message}");
+                                }
+                                // there was at least one active filter that failed in some way, add this line to the broken ones
+                                found_broken.push((identifier, edek, message))
+                            } else {
+                                if matched_filters
+                                    .into_iter()
+                                    .map(Result::unwrap)
+                                    .all(identity)
+                                {
+                                    // the current line passed all filters, add it to found
+                                    found_lines.push((identifier, edek));
+                                }
                             }
                         }
                         Err(e) => {
-                            println!(
-                                "WARNING: Encountered an unparsable EDEK on line {}: {}",
-                                line_number, e
-                            );
+                            if debug {
+                                println!(
+                                    "WARNING: Encountered an unparsable EDEK on line {line_number}: {e}"
+                                );
+                            }
                             found_broken.push((
                                 identifier,
                                 edek,
-                                format!("Encountered an unparsable EDEK: {}", e),
+                                format!("Encountered an unparsable EDEK: {e}"),
                             ));
                         }
                     }
                 }
                 Err(e) => {
                     println!(
-                        "Encountered an incorrectly formatted EDEK at line {}: {}",
-                        line_number, e
+                        "Encountered an incorrectly formatted EDEK at line {line_number}: {e}"
                     );
                     found_broken.push((
                         identifier,
                         edek,
-                        format!("Encountered an incorrectly formated EDEK: {}", e),
+                        format!("Encountered an incorrectly formated EDEK: {e}"),
                     ));
                 }
             }
@@ -139,7 +183,7 @@ fn main() {
             path.display()
         );
     } else {
-        println!("Found no EDEKs with the given config ID.");
+        println!("Found no EDEKs matching the search parameters.");
     }
     if !found_broken.is_empty() {
         let path = Path::new("broken-edeks.txt").to_path_buf();

src/util.rs

@@ -23,6 +23,8 @@ pub fn write_file<T>(
 where
     T: Serialize + GetIdentifier,
 {
+    // if we ever run into issues writing these vecs all at once (having the entire vec in memory and the concatenated
+    // string) we could write to the handler line by line or use Line/BufWriters directly avoiding vec collection
     let output_str = to_write
         .into_iter()
         .map(|line| {