From b1879ca91ef3f5ed9a487c8039eb16599d1055a0 Mon Sep 17 00:00:00 2001 From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com> Date: Fri, 4 Jul 2025 16:26:05 -0400 Subject: [PATCH] Add SSH key entry support --- src/password_manager/entry_management.py | 36 ++++++++++++++++++--- src/password_manager/manager.py | 28 ++++++++++++++++ src/password_manager/password_generation.py | 27 +++++++++++++++- src/tests/test_entry_add.py | 9 ++++-- src/tests/test_ssh_entry.py | 30 +++++++++++++++++ 5 files changed, 121 insertions(+), 9 deletions(-) create mode 100644 src/tests/test_ssh_entry.py diff --git a/src/password_manager/entry_management.py b/src/password_manager/entry_management.py index 098ede5..15819cd 100644 --- a/src/password_manager/entry_management.py +++ b/src/password_manager/entry_management.py @@ -209,16 +209,42 @@ class EntryManager: logger.error(f"Failed to generate otpauth URI: {e}") raise - def add_ssh_key(self, notes: str = "") -> int: - """Placeholder for adding an SSH key entry.""" - index = self.get_next_index() + def add_ssh_key( + self, parent_seed: str, index: int | None = None, notes: str = "" + ) -> int: + """Add a new SSH key pair entry. + + The provided ``index`` serves both as the vault entry identifier and + derivation index for the key. If not supplied, the next available index + is used. Only metadata is stored – keys are derived on demand. + """ + + if index is None: + index = self.get_next_index() + data = self.vault.load_index() data.setdefault("entries", {}) - data["entries"][str(index)] = {"type": EntryType.SSH.value, "notes": notes} + data["entries"][str(index)] = { + "type": EntryType.SSH.value, + "index": index, + "notes": notes, + } self._save_index(data) self.update_checksum() self.backup_manager.create_backup() - raise NotImplementedError("SSH key entry support not implemented yet") + return index + + def get_ssh_key_pair(self, index: int, parent_seed: str) -> tuple[str, str]: + """Return the PEM formatted SSH key pair for the given entry.""" + + entry = self.retrieve_entry(index) + if not entry or entry.get("type") != EntryType.SSH.value: + raise ValueError("Entry is not an SSH key entry") + + from password_manager.password_generation import derive_ssh_key_pair + + key_index = int(entry.get("index", index)) + return derive_ssh_key_pair(parent_seed, key_index) def add_seed(self, notes: str = "") -> int: """Placeholder for adding a seed entry.""" diff --git a/src/password_manager/manager.py b/src/password_manager/manager.py index c20e0aa..01ff8f4 100644 --- a/src/password_manager/manager.py +++ b/src/password_manager/manager.py @@ -1093,6 +1093,34 @@ class PasswordManager: logging.error(f"Error generating TOTP code: {e}", exc_info=True) print(colored(f"Error: Failed to generate TOTP code: {e}", "red")) return + if entry_type == EntryType.SSH.value: + notes = entry.get("notes", "") + try: + priv_pem, pub_pem = self.entry_manager.get_ssh_key_pair( + index, self.parent_seed + ) + if self.secret_mode_enabled: + copy_to_clipboard(priv_pem, self.clipboard_clear_delay) + print( + colored( + f"[+] SSH private key copied to clipboard. Will clear in {self.clipboard_clear_delay} seconds.", + "green", + ) + ) + print(colored("Public Key:", "cyan")) + print(pub_pem) + else: + print(colored("\n[+] Retrieved SSH Key Pair:\n", "green")) + print(colored("Public Key:", "cyan")) + print(pub_pem) + print(colored("Private Key:", "cyan")) + print(priv_pem) + if notes: + print(colored(f"Notes: {notes}", "cyan")) + except Exception as e: + logging.error(f"Error deriving SSH key pair: {e}", exc_info=True) + print(colored(f"Error: Failed to derive SSH keys: {e}", "red")) + return website_name = entry.get("website") length = entry.get("length") diff --git a/src/password_manager/password_generation.py b/src/password_manager/password_generation.py index 52970ca..41d7a93 100644 --- a/src/password_manager/password_generation.py +++ b/src/password_manager/password_generation.py @@ -25,8 +25,10 @@ from termcolor import colored from pathlib import Path import shutil from cryptography.hazmat.primitives.kdf.hkdf import HKDF -from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ed25519 from cryptography.hazmat.backends import default_backend +from bip_utils import Bip39SeedGenerator from local_bip85.bip85 import BIP85 @@ -340,6 +342,29 @@ def derive_ssh_key(bip85: BIP85, idx: int) -> bytes: return bip85.derive_entropy(index=idx, bytes_len=32, app_no=32) +def derive_ssh_key_pair(parent_seed: str, index: int) -> tuple[str, str]: + """Derive an Ed25519 SSH key pair from the seed phrase and index.""" + + seed_bytes = Bip39SeedGenerator(parent_seed).Generate() + bip85 = BIP85(seed_bytes) + entropy = derive_ssh_key(bip85, index) + + private_key = ed25519.Ed25519PrivateKey.from_private_bytes(entropy) + priv_pem = private_key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode() + + public_key = private_key.public_key() + pub_pem = public_key.public_bytes( + serialization.Encoding.PEM, + serialization.PublicFormat.SubjectPublicKeyInfo, + ).decode() + + return priv_pem, pub_pem + + def derive_seed_phrase(bip85: BIP85, idx: int, words: int = 24) -> str: """Derive a new BIP39 seed phrase using BIP85.""" return bip85.derive_mnemonic(index=idx, words_num=words) diff --git a/src/tests/test_entry_add.py b/src/tests/test_entry_add.py index b1f625d..5540851 100644 --- a/src/tests/test_entry_add.py +++ b/src/tests/test_entry_add.py @@ -61,9 +61,12 @@ def test_round_trip_entry_types(method, expected_type): entry_mgr.add_totp("example", TEST_SEED) index = 0 else: - with pytest.raises(NotImplementedError): - getattr(entry_mgr, method)() - index = 0 + if method == "add_ssh_key": + index = entry_mgr.add_ssh_key(TEST_SEED) + else: + with pytest.raises(NotImplementedError): + getattr(entry_mgr, method)() + index = 0 entry = entry_mgr.retrieve_entry(index) assert entry["type"] == expected_type diff --git a/src/tests/test_ssh_entry.py b/src/tests/test_ssh_entry.py new file mode 100644 index 0000000..d69b4d3 --- /dev/null +++ b/src/tests/test_ssh_entry.py @@ -0,0 +1,30 @@ +import sys +from pathlib import Path +from tempfile import TemporaryDirectory + +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.vault import Vault +from password_manager.config_manager import ConfigManager + + +def test_add_and_retrieve_ssh_key_pair(): + with TemporaryDirectory() as tmpdir: + tmp_path = Path(tmpdir) + vault, enc_mgr = create_vault(tmp_path, TEST_SEED, TEST_PASSWORD) + cfg_mgr = ConfigManager(vault, tmp_path) + backup_mgr = BackupManager(tmp_path, cfg_mgr) + entry_mgr = EntryManager(vault, backup_mgr) + + index = entry_mgr.add_ssh_key(TEST_SEED) + entry = entry_mgr.retrieve_entry(index) + assert entry == {"type": "ssh", "index": index, "notes": ""} + + priv1, pub1 = entry_mgr.get_ssh_key_pair(index, TEST_SEED) + priv2, pub2 = entry_mgr.get_ssh_key_pair(index, TEST_SEED) + assert priv1 == priv2 + assert pub1 == pub2