mirror of
https://github.com/PR0M3TH3AN/SeedPass.git
synced 2025-09-08 07:18:47 +00:00
Merge pull request #184 from PR0M3TH3AN/codex/add-support-for-importing-2fa-secrets
Add imported TOTP support
This commit is contained in:
@@ -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)
|
||||
|
||||
|
@@ -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"))
|
||||
|
@@ -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
|
||||
|
@@ -41,6 +41,7 @@ def test_handle_add_totp(monkeypatch):
|
||||
|
||||
inputs = iter(
|
||||
[
|
||||
"1", # choose derive
|
||||
"Example", # label
|
||||
"", # period
|
||||
"", # digits
|
||||
|
@@ -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)
|
||||
|
Reference in New Issue
Block a user