From 446e78b95a2d2b74afb9ab7c39bf135ee5e8ab32 Mon Sep 17 00:00:00 2001
From: Robin Krahl <robin.krahl@ireas.org>
Date: Thu, 27 Aug 2020 13:13:54 +0200
Subject: Add otp-cache extension

---
 extensions/otp-cache/src/main.rs | 118 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 118 insertions(+)
 create mode 100644 extensions/otp-cache/src/main.rs

(limited to 'extensions/otp-cache/src/main.rs')

diff --git a/extensions/otp-cache/src/main.rs b/extensions/otp-cache/src/main.rs
new file mode 100644
index 0000000..b5f0b90
--- /dev/null
+++ b/extensions/otp-cache/src/main.rs
@@ -0,0 +1,118 @@
+mod ext;
+
+use std::collections;
+use std::fs;
+use std::io;
+use std::path;
+
+type Cache = collections::BTreeMap<ext::OtpAlgorithm, Vec<Slot>>;
+
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
+struct Slot {
+  name: String,
+  id: usize,
+}
+
+/// Access Nitrokey OTP slots by name
+#[derive(Debug, structopt::StructOpt)]
+#[structopt(bin_name = "nitrocli otp-cache")]
+struct Args {
+  /// Update the cached slot data
+  #[structopt(short, long)]
+  force_update: bool,
+
+  #[structopt(short, long, global = true, default_value = "totp")]
+  algorithm: ext::OtpAlgorithm,
+
+  #[structopt(subcommand)]
+  cmd: Command,
+
+  #[structopt(flatten)]
+  ext: ext::Args,
+}
+
+#[derive(Debug, structopt::StructOpt)]
+enum Command {
+  /// Generates a one-time passwords
+  Get {
+    /// The name of the OTP slot to generate a OTP from
+    name: String,
+  },
+  /// Lists the cached slots and their ID
+  List,
+}
+
+fn main() -> anyhow::Result<()> {
+  use structopt::StructOpt as _;
+
+  let args = Args::from_args();
+  let mut cache = get_cache(&args)?;
+  let slots = cache.remove(&args.algorithm).unwrap_or_default();
+
+  match &args.cmd {
+    Command::Get { name } => match slots.iter().find(|s| &s.name == name) {
+      Some(slot) => println!("{}", generate_otp(&args, slot.id)?),
+      None => anyhow::bail!("No OTP slot with the given name!"),
+    },
+    Command::List => {
+      println!("id\tslot");
+      for slot in slots {
+        println!("{}\t{}", slot.id, slot.name);
+      }
+    }
+  }
+
+  Ok(())
+}
+
+fn get_cache(args: &Args) -> anyhow::Result<Cache> {
+  // TODO: use model and serial number as cache file name
+  let cache_file =
+    xdg::BaseDirectories::with_prefix("nitrocli")?.place_cache_file("otp-cache/cache.toml")?;
+  if args.force_update || !cache_file.is_file() {
+    let cache = update_cache(args)?;
+    save_cache(&cache, &cache_file)?;
+    Ok(cache)
+  } else {
+    load_cache(&cache_file)
+  }
+}
+
+fn update_cache(args: &Args) -> anyhow::Result<Cache> {
+  let slots = args.ext.nitrocli().args(&["otp", "status"]).text()?;
+  let mut cache = Cache::new();
+  for line in slots.lines().skip(1) {
+    let parts: Vec<_> = line.splitn(3, "\t").collect();
+    if parts.len() == 3 {
+      let algorithm: ext::OtpAlgorithm = parts[0].parse()?;
+      let id: usize = parts[1].parse()?;
+      let name = parts[2].to_owned();
+      cache.entry(algorithm).or_default().push(Slot { name, id });
+    }
+  }
+  Ok(cache)
+}
+
+fn save_cache(cache: &Cache, path: &path::Path) -> anyhow::Result<()> {
+  use io::Write as _;
+
+  let mut f = fs::File::create(path)?;
+  f.write_all(&toml::to_vec(cache)?)?;
+  Ok(())
+}
+
+fn load_cache(path: &path::Path) -> anyhow::Result<Cache> {
+  let s = fs::read_to_string(path)?;
+  toml::from_str(&s).map_err(From::from)
+}
+
+fn generate_otp(args: &Args, slot: usize) -> anyhow::Result<String> {
+  args
+    .ext
+    .nitrocli()
+    .args(&["otp", "get"])
+    .arg(slot.to_string())
+    .arg("--algorithm")
+    .arg(args.algorithm.to_string())
+    .text()
+}
-- 
cgit v1.2.1