Add TOTP export option

This commit is contained in:
thePR0M3TH3AN
2025-07-03 10:31:06 -04:00
parent c014acc2e4
commit 6a98df4e56
3 changed files with 121 additions and 5 deletions

View File

@@ -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"))

View File

@@ -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.

View File

@@ -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/")