From 06801897239edb341c032c1d01d3044a25711cd8 Mon Sep 17 00:00:00 2001
From: Awiteb <a@4rs.nl>
Date: Sun, 19 May 2024 04:27:09 +0300
Subject: [PATCH] chore: Validate the TOTP secret string

---
 src/cli/add_command.rs  | 12 +++++++++---
 src/cli/edit_command.rs | 16 +++++++++++++---
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/cli/add_command.rs b/src/cli/add_command.rs
index 9bf6f2c..e543a4b 100644
--- a/src/cli/add_command.rs
+++ b/src/cli/add_command.rs
@@ -19,7 +19,7 @@ use clap::Args;
 use crate::{
     clap_parsers,
     utils,
-    vault::{Vault, Vaults},
+    vault::{cipher, Vault, Vaults},
     LprsCommand,
     LprsError,
     LprsResult,
@@ -64,10 +64,16 @@ impl Add {
 impl LprsCommand for Add {
     fn run(mut self, mut vault_manager: Vaults) -> LprsResult<()> {
         if !self.is_empty() {
+            if let Some(totp_secret) = utils::user_secret(self.totp_secret, "TOTP Secret:", false)?
+            {
+                cipher::base32_decode(&totp_secret).map_err(|_| {
+                    LprsError::Base32("Invalid TOTP secret, must be valid base32 string".to_owned())
+                })?;
+                self.vault_info.totp_secret = Some(totp_secret);
+            }
+
             self.vault_info.name = self.vault_info.name.trim().to_string();
             self.vault_info.password = utils::user_secret(self.password, "Vault password:", false)?;
-            self.vault_info.totp_secret =
-                utils::user_secret(self.totp_secret, "TOTP Secret:", false)?;
             self.vault_info.custom_fields = self.custom_fields.into_iter().collect();
             vault_manager.add_vault(self.vault_info);
             vault_manager.try_export()?;
diff --git a/src/cli/edit_command.rs b/src/cli/edit_command.rs
index e1606c6..8319f0a 100644
--- a/src/cli/edit_command.rs
+++ b/src/cli/edit_command.rs
@@ -16,7 +16,14 @@
 
 use clap::Args;
 
-use crate::{clap_parsers, utils, vault::Vaults, LprsCommand, LprsError, LprsResult};
+use crate::{
+    clap_parsers,
+    utils,
+    vault::{cipher, Vaults},
+    LprsCommand,
+    LprsError,
+    LprsResult,
+};
 
 #[derive(Debug, Args)]
 #[command(author, version, about, long_about = None)]
@@ -79,8 +86,11 @@ impl LprsCommand for Edit {
         if self.password.is_some() {
             vault.password = utils::user_secret(self.password, "New vault password:", false)?;
         }
-        if self.totp_secret.is_some() {
-            vault.totp_secret = utils::user_secret(self.totp_secret, "TOTP Secret:", false)?;
+        if let Some(totp_secret) = utils::user_secret(self.totp_secret, "TOTP Secret:", false)? {
+            cipher::base32_decode(&totp_secret).map_err(|_| {
+                LprsError::Base32("Invalid TOTP secret, must be valid base32 string".to_owned())
+            })?;
+            vault.totp_secret = Some(totp_secret);
         }
         if let Some(new_username) = self.username {
             vault.username = Some(new_username);