Merge pull request #184 from PR0M3TH3AN/codex/add-support-for-importing-2fa-secrets

Add imported TOTP support

src/password_manager/entry_management.py

@@ -168,31 +168,44 @@ class EntryManager:
         self,
         label: str,
         parent_seed: str,
+        *,
+        secret: str | None = None,
         index: int | None = None,
         period: int = 30,
         digits: int = 6,
     ) -> str:
         """Add a new TOTP entry and return the provisioning URI."""
         entry_id = self.get_next_index()
-        if index is None:
-            index = self.get_next_totp_index()
         data = self.vault.load_index()
         data.setdefault("entries", {})
 
-        data["entries"][str(entry_id)] = {
+        if secret is None:
+            if index is None:
+                index = self.get_next_totp_index()
+            secret = TotpManager.derive_secret(parent_seed, index)
+            entry = {
                 "type": EntryType.TOTP.value,
                 "label": label,
                 "index": index,
                 "period": period,
                 "digits": digits,
             }
+        else:
+            entry = {
+                "type": EntryType.TOTP.value,
+                "label": label,
+                "secret": secret,
+                "period": period,
+                "digits": digits,
+            }
+
+        data["entries"][str(entry_id)] = entry
 
         self._save_index(data)
         self.update_checksum()
         self.backup_index_file()
 
         try:
-            secret = TotpManager.derive_secret(parent_seed, index)
             return TotpManager.make_otpauth_uri(label, secret, period, digits)
         except Exception as e:
             logger.error(f"Failed to generate otpauth URI: {e}")
@@ -221,13 +234,16 @@ class EntryManager:
         raise NotImplementedError("Seed entry support not implemented yet")
 
     def get_totp_code(
-        self, index: int, parent_seed: str, timestamp: int | None = None
+        self, index: int, parent_seed: str | None = None, timestamp: int | None = None
     ) -> str:
         """Return the current TOTP code for the specified entry."""
         entry = self.retrieve_entry(index)
         if not entry or entry.get("type") != EntryType.TOTP.value:
             raise ValueError("Entry is not a TOTP entry")
-
+        if "secret" in entry:
+            return TotpManager.current_code_from_secret(entry["secret"], timestamp)
+        if parent_seed is None:
+            raise ValueError("Seed required for derived TOTP")
         totp_index = int(entry.get("index", 0))
         return TotpManager.current_code(parent_seed, totp_index, timestamp)
 

src/password_manager/manager.py

@@ -864,59 +864,96 @@ class PasswordManager:
             print(colored(f"Error: Failed to generate password: {e}", "red"))
 
     def handle_add_totp(self) -> None:
-        """Prompt for details and add a new TOTP entry."""
+        """Add a TOTP entry either derived from the seed or imported."""
         try:
-            label = input("Enter the account label: ").strip()
+            while True:
+                print("\nAdd TOTP:")
+                print("1. Make 2FA (derive from seed)")
+                print("2. Import 2FA (paste otpauth URI or secret)")
+                print("3. Back")
+                choice = input("Select option: ").strip()
+                if choice == "1":
+                    label = input("Label: ").strip()
                     if not label:
                         print(colored("Error: Label cannot be empty.", "red"))
-                return
+                        continue
-
+                    period = input("Period (default 30): ").strip() or "30"
+                    digits = input("Digits (default 6): ").strip() or "6"
+                    if not period.isdigit() or not digits.isdigit():
+                        print(
+                            colored("Error: Period and digits must be numbers.", "red")
+                        )
+                        continue
                     totp_index = self.entry_manager.get_next_totp_index()
-
-            period_input = input("TOTP period in seconds (default 30): ").strip()
-            period = 30
-            if period_input:
-                if not period_input.isdigit():
-                    print(colored("Error: Period must be a number.", "red"))
-                    return
-                period = int(period_input)
-
-            digits_input = input("Number of digits (default 6): ").strip()
-            digits = 6
-            if digits_input:
-                if not digits_input.isdigit():
-                    print(colored("Error: Digits must be a number.", "red"))
-                    return
-                digits = int(digits_input)
-
                     entry_id = self.entry_manager.get_next_index()
                     uri = self.entry_manager.add_totp(
                         label,
                         self.parent_seed,
                         index=totp_index,
-                period=period,
+                        period=int(period),
-                digits=digits,
+                        digits=int(digits),
                     )
-
+                    secret = TotpManager.derive_secret(self.parent_seed, totp_index)
                     self.is_dirty = True
                     self.last_update = time.time()
-
+                    print(
-            secret = TotpManager.derive_secret(self.parent_seed, totp_index)
+                        colored(
-
+                            f"\n[+] TOTP entry added with ID {entry_id}.\n", "green"
-            print(colored(f"\n[+] TOTP entry added with ID {entry_id}.\n", "green"))
+                        )
+                    )
                     print(colored("Add this URI to your authenticator app:", "cyan"))
                     print(colored(uri, "yellow"))
                     print(colored(f"Secret: {secret}\n", "cyan"))
-
                     try:
                         self.sync_vault()
-                logging.info("Encrypted index posted to Nostr after TOTP add.")
                     except Exception as nostr_error:
                         logging.error(
                             f"Failed to post updated index to Nostr: {nostr_error}",
                             exc_info=True,
                         )
-
+                    break
+                elif choice == "2":
+                    raw = input("Paste otpauth URI or secret: ").strip()
+                    try:
+                        if raw.lower().startswith("otpauth://"):
+                            label, secret, period, digits = TotpManager.parse_otpauth(
+                                raw
+                            )
+                        else:
+                            label = input("Label: ").strip()
+                            secret = raw.upper()
+                            period = int(input("Period (default 30): ").strip() or 30)
+                            digits = int(input("Digits (default 6): ").strip() or 6)
+                        entry_id = self.entry_manager.get_next_index()
+                        uri = self.entry_manager.add_totp(
+                            label,
+                            self.parent_seed,
+                            secret=secret,
+                            period=period,
+                            digits=digits,
+                        )
+                        self.is_dirty = True
+                        self.last_update = time.time()
+                        print(
+                            colored(
+                                f"\nImported \u2714  Codes for {label} are now stored in SeedPass.",
+                                "green",
+                            )
+                        )
+                        try:
+                            self.sync_vault()
+                        except Exception as nostr_error:
+                            logging.error(
+                                f"Failed to post updated index to Nostr: {nostr_error}",
+                                exc_info=True,
+                            )
+                        break
+                    except ValueError as err:
+                        print(colored(f"Error: {err}", "red"))
+                elif choice == "3":
+                    return
+                else:
+                    print(colored("Invalid choice.", "red"))
         except Exception as e:
             logging.error(f"Error during TOTP setup: {e}", exc_info=True)
             print(colored(f"Error: Failed to add TOTP: {e}", "red"))

src/password_manager/totp.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import sys
 import time
 from urllib.parse import quote
+from urllib.parse import urlparse, parse_qs, unquote
 
 import pyotp
 
@@ -28,6 +29,27 @@ class TotpManager:
             return totp.now()
         return totp.at(timestamp)
 
+    @staticmethod
+    def current_code_from_secret(secret: str, timestamp: int | None = None) -> str:
+        """Return the TOTP code for a raw secret."""
+        totp = pyotp.TOTP(secret)
+        return totp.now() if timestamp is None else totp.at(timestamp)
+
+    @staticmethod
+    def parse_otpauth(uri: str) -> tuple[str, str, int, int]:
+        """Parse an otpauth URI and return (label, secret, period, digits)."""
+        if not uri.startswith("otpauth://"):
+            raise ValueError("Not an otpauth URI")
+        parsed = urlparse(uri)
+        label = unquote(parsed.path.lstrip("/"))
+        qs = parse_qs(parsed.query)
+        secret = qs.get("secret", [""])[0].upper()
+        period = int(qs.get("period", ["30"])[0])
+        digits = int(qs.get("digits", ["6"])[0])
+        if not secret:
+            raise ValueError("Missing secret in URI")
+        return label, secret, period, digits
+
     @staticmethod
     def make_otpauth_uri(
         label: str, secret: str, period: int = 30, digits: int = 6

src/tests/test_manager_add_totp.py

@@ -41,6 +41,7 @@ def test_handle_add_totp(monkeypatch):
 
         inputs = iter(
             [
+                "1",  # choose derive
                 "Example",  # label
                 "",  # period
                 "",  # digits

src/tests/test_totp_entry.py

@@ -12,6 +12,7 @@ sys.path.append(str(Path(__file__).resolve().parents[1]))
 from password_manager.entry_management import EntryManager
 from password_manager.vault import Vault
 from password_manager.totp import TotpManager
+import pyotp
 
 
 def test_add_totp_and_get_code():
@@ -47,3 +48,20 @@ def test_totp_time_remaining(monkeypatch):
         monkeypatch.setattr(TotpManager, "time_remaining", lambda period: 7)
         remaining = entry_mgr.get_totp_time_remaining(0)
         assert remaining == 7
+
+
+def test_add_totp_imported(tmp_path):
+    vault, enc = create_vault(tmp_path, TEST_SEED, TEST_PASSWORD)
+    em = EntryManager(vault, tmp_path)
+    secret = "JBSWY3DPEHPK3PXP"
+    em.add_totp("Imported", TEST_SEED, secret=secret)
+    entry = em.retrieve_entry(0)
+    assert entry == {
+        "type": "totp",
+        "label": "Imported",
+        "secret": secret,
+        "period": 30,
+        "digits": 6,
+    }
+    code = em.get_totp_code(0, timestamp=0)
+    assert code == pyotp.TOTP(secret).at(0)