From 6a98df4e564f453b851718524000151db13937ad Mon Sep 17 00:00:00 2001 From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com> Date: Thu, 3 Jul 2025 10:31:06 -0400 Subject: [PATCH] Add TOTP export option --- src/main.py | 13 ++++--- src/password_manager/manager.py | 58 +++++++++++++++++++++++++++++ src/tests/test_export_totp_codes.py | 55 +++++++++++++++++++++++++++ 3 files changed, 121 insertions(+), 5 deletions(-) create mode 100644 src/tests/test_export_totp_codes.py diff --git a/src/main.py b/src/main.py index cecfd63..14aedc5 100644 --- a/src/main.py +++ b/src/main.py @@ -503,9 +503,10 @@ def handle_settings(password_manager: PasswordManager) -> None: print("5. Backup Parent Seed") print("6. Export database") print("7. Import database") - print("8. Set inactivity timeout") - print("9. Lock Vault") - print("10. Back") + print("8. Export 2FA codes") + print("9. Set inactivity timeout") + print("10. Lock Vault") + print("11. Back") choice = input("Select an option: ").strip() if choice == "1": handle_profiles_menu(password_manager) @@ -524,12 +525,14 @@ def handle_settings(password_manager: PasswordManager) -> None: if path: password_manager.handle_import_database(Path(path)) elif choice == "8": - handle_set_inactivity_timeout(password_manager) + password_manager.handle_export_totp_codes() elif choice == "9": + handle_set_inactivity_timeout(password_manager) + elif choice == "10": password_manager.lock_vault() print(colored("Vault locked. Please re-enter your password.", "yellow")) password_manager.unlock_vault() - elif choice == "10": + elif choice == "11": break else: print(colored("Invalid choice.", "red")) diff --git a/src/password_manager/manager.py b/src/password_manager/manager.py index a3cc557..473a51f 100644 --- a/src/password_manager/manager.py +++ b/src/password_manager/manager.py @@ -38,6 +38,7 @@ from utils.checksum import calculate_checksum, verify_checksum from utils.password_prompt import ( prompt_for_password, prompt_existing_password, + prompt_new_password, confirm_action, ) from utils.memory_protection import InMemorySecret @@ -1440,6 +1441,63 @@ class PasswordManager: logging.error(f"Failed to import database: {e}", exc_info=True) print(colored(f"Error: Failed to import database: {e}", "red")) + def handle_export_totp_codes(self) -> Path | None: + """Export all 2FA codes to a JSON file for other authenticator apps.""" + try: + data = self.entry_manager.vault.load_index() + entries = data.get("entries", {}) + + totp_entries = [] + for entry in entries.values(): + if entry.get("type") == EntryType.TOTP.value: + label = entry.get("label", "") + period = int(entry.get("period", 30)) + digits = int(entry.get("digits", 6)) + if "secret" in entry: + secret = entry["secret"] + else: + idx = int(entry.get("index", 0)) + secret = TotpManager.derive_secret(self.parent_seed, idx) + uri = TotpManager.make_otpauth_uri(label, secret, period, digits) + totp_entries.append( + { + "label": label, + "secret": secret, + "period": period, + "digits": digits, + "uri": uri, + } + ) + + if not totp_entries: + print(colored("No 2FA codes to export.", "yellow")) + return None + + dest_str = input( + "Enter destination file path (default: totp_export.json): " + ).strip() + dest = Path(dest_str) if dest_str else Path("totp_export.json") + + json_data = json.dumps({"entries": totp_entries}, indent=2) + + if confirm_action("Encrypt export with a password? (Y/N): "): + password = prompt_new_password() + key = derive_key_from_password(password) + enc_mgr = EncryptionManager(key, dest.parent) + data_bytes = enc_mgr.encrypt_data(json_data.encode("utf-8")) + dest = dest.with_suffix(dest.suffix + ".enc") + dest.write_bytes(data_bytes) + else: + dest.write_text(json_data) + + os.chmod(dest, 0o600) + print(colored(f"2FA codes exported to '{dest}'.", "green")) + return dest + except Exception as e: + logging.error(f"Failed to export TOTP codes: {e}", exc_info=True) + print(colored(f"Error: Failed to export 2FA codes: {e}", "red")) + return None + def handle_backup_reveal_parent_seed(self) -> None: """ Handles the backup and reveal of the parent seed. diff --git a/src/tests/test_export_totp_codes.py b/src/tests/test_export_totp_codes.py new file mode 100644 index 0000000..0c40936 --- /dev/null +++ b/src/tests/test_export_totp_codes.py @@ -0,0 +1,55 @@ +import sys +from pathlib import Path +from tempfile import TemporaryDirectory + +import json + +from helpers import create_vault, TEST_SEED, TEST_PASSWORD + +sys.path.append(str(Path(__file__).resolve().parents[1])) + +from password_manager.entry_management import EntryManager +from password_manager.backup import BackupManager +from password_manager.manager import PasswordManager, EncryptionMode +from password_manager.totp import TotpManager + + +class FakeNostrClient: + def publish_snapshot(self, data: bytes): + return None, "abcd" + + +def test_handle_export_totp_codes(monkeypatch, tmp_path): + vault, enc_mgr = create_vault(tmp_path, TEST_SEED, TEST_PASSWORD) + entry_mgr = EntryManager(vault, tmp_path) + backup_mgr = BackupManager(tmp_path) + + pm = PasswordManager.__new__(PasswordManager) + pm.encryption_mode = EncryptionMode.SEED_ONLY + pm.encryption_manager = enc_mgr + pm.vault = vault + pm.entry_manager = entry_mgr + pm.backup_manager = backup_mgr + pm.parent_seed = TEST_SEED + pm.nostr_client = FakeNostrClient() + pm.fingerprint_dir = tmp_path + + # add totp entries + entry_mgr.add_totp("Example", TEST_SEED) + entry_mgr.add_totp("Imported", TEST_SEED, secret="JBSWY3DPEHPK3PXP") + + export_path = tmp_path / "out.json" + monkeypatch.setattr("builtins.input", lambda *a, **k: str(export_path)) + monkeypatch.setattr( + "password_manager.manager.confirm_action", lambda *_a, **_k: False + ) + + pm.handle_export_totp_codes() + + data = json.loads(export_path.read_text()) + assert len(data["entries"]) == 2 + labels = {e["label"] for e in data["entries"]} + assert {"Example", "Imported"} == labels + # check URI format + uri = data["entries"][0]["uri"] + assert uri.startswith("otpauth://totp/")