diff --git a/src/password_manager/entry_management.py b/src/password_manager/entry_management.py index 70f1832..87285ca 100644 --- a/src/password_manager/entry_management.py +++ b/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)] = { - "type": EntryType.TOTP.value, - "label": label, - "index": index, - "period": period, - "digits": digits, - } + 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) diff --git a/src/password_manager/manager.py b/src/password_manager/manager.py index 009efd9..9eec67f 100644 --- a/src/password_manager/manager.py +++ b/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() - if not label: - print(colored("Error: Label cannot be empty.", "red")) - return - - 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")) + 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")) + 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() + entry_id = self.entry_manager.get_next_index() + uri = self.entry_manager.add_totp( + label, + self.parent_seed, + index=totp_index, + period=int(period), + digits=int(digits), + ) + secret = TotpManager.derive_secret(self.parent_seed, totp_index) + self.is_dirty = True + self.last_update = time.time() + 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() + 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 - 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, - digits=digits, - ) - - self.is_dirty = True - self.last_update = time.time() - - secret = TotpManager.derive_secret(self.parent_seed, totp_index) - - 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, - ) - + 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")) diff --git a/src/password_manager/totp.py b/src/password_manager/totp.py index 4e50e22..a6a88b6 100644 --- a/src/password_manager/totp.py +++ b/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 diff --git a/src/tests/test_manager_add_totp.py b/src/tests/test_manager_add_totp.py index 0540927..6ee18b9 100644 --- a/src/tests/test_manager_add_totp.py +++ b/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 diff --git a/src/tests/test_totp_entry.py b/src/tests/test_totp_entry.py index a6d946a..d79aa41 100644 --- a/src/tests/test_totp_entry.py +++ b/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)