From 9f966de69a5e5ca6704cd1005356cb1eef09a64e Mon Sep 17 00:00:00 2001
From: Vincent Breitmoser <look@my.amazin.horse>
Date: Wed, 5 Jun 2019 21:04:11 +0200
Subject: [PATCH] hagridctl: working "regenerate" command

---
 hagridctl/src/main.rs       |  6 +++
 hagridctl/src/regenerate.rs | 73 +++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)
 create mode 100644 hagridctl/src/regenerate.rs

diff --git a/hagridctl/src/main.rs b/hagridctl/src/main.rs
index 7daefc9..bf8bf49 100644
--- a/hagridctl/src/main.rs
+++ b/hagridctl/src/main.rs
@@ -11,6 +11,7 @@ extern crate hagrid_database as database;
 extern crate serde_derive;
 extern crate toml;
 extern crate indicatif;
+extern crate walkdir;
 
 use std::fs;
 use std::path::PathBuf;
@@ -21,6 +22,7 @@ use failure::Fallible as Result;
 use clap::{Arg, App, SubCommand};
 
 mod import;
+mod regenerate;
 
 #[derive(Deserialize)]
 pub struct HagridConfigs {
@@ -60,6 +62,8 @@ fn main() -> Result<()> {
                                .takes_value(true)
                                .default_value("prod")
                                .possible_values(&["dev","stage","prod"]))
+                          .subcommand(SubCommand::with_name("regenerate")
+                                      .about("Regenerate symlink directory"))
                           .subcommand(SubCommand::with_name("import")
                                       .about("Import keys into Hagrid")
                                       .arg(Arg::with_name("dry run")
@@ -91,6 +95,8 @@ fn main() -> Result<()> {
             .map(|arg| PathBuf::from_str(arg).unwrap())
             .collect();
         import::do_import(&config, dry_run, keyrings)?;
+    } else if let Some(_matches) = matches.subcommand_matches("regenerate") {
+        regenerate::do_regenerate(&config)?;
     } else {
         println!("{}", matches.usage());
     }
diff --git a/hagridctl/src/regenerate.rs b/hagridctl/src/regenerate.rs
new file mode 100644
index 0000000..adb65ce
--- /dev/null
+++ b/hagridctl/src/regenerate.rs
@@ -0,0 +1,73 @@
+use failure::Fallible as Result;
+
+use std::path::Path;
+
+use walkdir::WalkDir;
+use indicatif::{ProgressBar,ProgressStyle};
+
+use HagridConfig;
+use database::{Database,KeyDatabase};
+use database::types::Fingerprint;
+
+pub fn do_regenerate(config: &HagridConfig) -> Result<()> {
+    let db = KeyDatabase::new_internal(
+        config.keys_internal_dir.as_ref().unwrap(),
+        config.keys_external_dir.as_ref().unwrap(),
+        config.tmp_dir.as_ref().unwrap(),
+        false,
+    )?;
+
+    let published_dir = config.keys_external_dir.as_ref().unwrap().join("published");
+    let dirs: Vec<_> = WalkDir::new(published_dir)
+        .min_depth(1)
+        .max_depth(1)
+        .into_iter()
+        .flatten()
+        .map(|entry| entry.into_path())
+        .collect();
+
+    let progress_bar = ProgressBar::new(dirs.len() as u64);
+    progress_bar
+        .set_style(ProgressStyle::default_bar()
+            .template("[{elapsed_precise}] {bar:40.cyan/blue} {msg}")
+            .progress_chars("##-"));
+
+    for dir in dirs {
+        progress_bar.inc(1);
+        for dir2 in WalkDir::new(dir)
+            .min_depth(1)
+            .max_depth(1)
+            .into_iter()
+            .flatten()
+            .map(|entry| entry.into_path()) {
+
+            let prefix2 = dir2.file_name().unwrap().to_string_lossy();
+            let prefix1 = dir2.parent().unwrap().file_name().unwrap().to_string_lossy();
+            progress_bar.set_message(&format!("Regenerating keys with prefix {}{}",
+                                    prefix1, prefix2));
+            regenerate_dir(&db, &dir2)?;
+        }
+    }
+
+    Ok(())
+}
+
+fn regenerate_dir(db: &KeyDatabase, dir: &Path) -> Result<()> {
+    for path in WalkDir::new(dir)
+        .min_depth(1)
+        .max_depth(1)
+        .into_iter()
+        .flatten()
+        .map(|entry| entry.into_path()) {
+
+        let suffix = path.file_name().unwrap().to_string_lossy();
+        let prefix2 = path.parent().unwrap().file_name().unwrap().to_string_lossy();
+        let prefix1 = path.parent().unwrap().parent().unwrap().file_name().unwrap().to_string_lossy();
+        let fpr_str = format!("{}{}{}", prefix1, prefix2, suffix);
+
+        let fpr = fpr_str.parse::<Fingerprint>()?;
+        db.regenerate_links(&fpr)?;
+    }
+
+    Ok(())
+}