Merge pull request #163 from PR0M3TH3AN/codex/replace-sys.exit-with-custom-exceptions

Add custom exception handling

src/local_bip85/bip85.py

@@ -31,6 +31,12 @@ from cryptography.hazmat.backends import default_backend
 logger = logging.getLogger(__name__)
 
 
+class Bip85Error(Exception):
+    """Exception raised for BIP85-related errors."""
+
+    pass
+
+
 class BIP85:
     def __init__(self, seed_bytes: bytes | str):
         """Initialize from BIP39 seed bytes or BIP32 xprv string."""
@@ -43,7 +49,7 @@ class BIP85:
         except Exception as e:
             logging.error(f"Error initializing BIP32 context: {e}", exc_info=True)
             print(f"{Fore.RED}Error initializing BIP32 context: {e}")
-            sys.exit(1)
+            raise Bip85Error(f"Error initializing BIP32 context: {e}")
 
     def derive_entropy(
         self, index: int, bytes_len: int, app_no: int = 39, words_len: int | None = None
@@ -90,21 +96,23 @@ class BIP85:
                 print(
                     f"{Fore.RED}Error: Derived entropy length is {len(entropy)} bytes; expected {bytes_len} bytes."
                 )
-                sys.exit(1)
+                raise Bip85Error(
+                    f"Derived entropy length is {len(entropy)} bytes; expected {bytes_len} bytes."
+                )
 
             logging.debug(f"Derived entropy: {entropy.hex()}")
             return entropy
         except Exception as e:
             logging.error(f"Error deriving entropy: {e}", exc_info=True)
             print(f"{Fore.RED}Error deriving entropy: {e}")
-            sys.exit(1)
+            raise Bip85Error(f"Error deriving entropy: {e}")
 
     def derive_mnemonic(self, index: int, words_num: int) -> str:
         bytes_len = {12: 16, 18: 24, 24: 32}.get(words_num)
         if not bytes_len:
             logging.error(f"Unsupported number of words: {words_num}")
             print(f"{Fore.RED}Error: Unsupported number of words: {words_num}")
-            sys.exit(1)
+            raise Bip85Error(f"Unsupported number of words: {words_num}")
 
         entropy = self.derive_entropy(
             index=index, bytes_len=bytes_len, app_no=39, words_len=words_num
@@ -118,7 +126,7 @@ class BIP85:
         except Exception as e:
             logging.error(f"Error generating mnemonic: {e}", exc_info=True)
             print(f"{Fore.RED}Error generating mnemonic: {e}")
-            sys.exit(1)
+            raise Bip85Error(f"Error generating mnemonic: {e}")
 
     def derive_symmetric_key(self, index: int = 0, app_no: int = 2) -> bytes:
         """Derive 32 bytes of entropy for symmetric key usage."""
@@ -129,4 +137,4 @@ class BIP85:
         except Exception as e:
             logging.error(f"Error deriving symmetric key: {e}", exc_info=True)
             print(f"{Fore.RED}Error deriving symmetric key: {e}")
-            sys.exit(1)
+            raise Bip85Error(f"Error deriving symmetric key: {e}")

src/main.py

@@ -17,6 +17,8 @@ import traceback
 from password_manager.manager import PasswordManager
 from nostr.client import NostrClient
 from constants import INACTIVITY_TIMEOUT
+from utils.password_prompt import PasswordPromptError
+from local_bip85.bip85 import Bip85Error
 
 
 colorama_init()
@@ -631,6 +633,10 @@ if __name__ == "__main__":
     try:
         password_manager = PasswordManager()
         logger.info("PasswordManager initialized successfully.")
+    except (PasswordPromptError, Bip85Error) as e:
+        logger.error(f"Failed to initialize PasswordManager: {e}", exc_info=True)
+        print(colored(f"Error: Failed to initialize PasswordManager: {e}", "red"))
+        sys.exit(1)
     except Exception as e:
         logger.error(f"Failed to initialize PasswordManager: {e}", exc_info=True)
         print(colored(f"Error: Failed to initialize PasswordManager: {e}", "red"))
@@ -677,6 +683,16 @@ if __name__ == "__main__":
             logging.error(f"Error during shutdown: {e}")
             print(colored(f"Error during shutdown: {e}", "red"))
         sys.exit(0)
+    except (PasswordPromptError, Bip85Error) as e:
+        logger.error(f"A user-related error occurred: {e}", exc_info=True)
+        print(colored(f"Error: {e}", "red"))
+        try:
+            password_manager.nostr_client.close_client_pool()
+            logging.info("NostrClient closed successfully.")
+        except Exception as close_error:
+            logging.error(f"Error during shutdown: {close_error}")
+            print(colored(f"Error during shutdown: {close_error}", "red"))
+        sys.exit(1)
     except Exception as e:
         logger.error(f"An unexpected error occurred: {e}", exc_info=True)
         print(colored(f"Error: An unexpected error occurred: {e}", "red"))

src/password_manager/manager.py

@@ -55,7 +55,7 @@ import gzip
 import bcrypt
 from pathlib import Path
 
-from local_bip85.bip85 import BIP85
+from local_bip85.bip85 import BIP85, Bip85Error
 from bip_utils import Bip39SeedGenerator, Bip39MnemonicGenerator, Bip39Languages
 from datetime import datetime
 
@@ -645,6 +645,10 @@ class PasswordManager:
             bip85 = BIP85(master_seed)
             mnemonic = bip85.derive_mnemonic(index=0, words_num=12)
             return mnemonic
+        except Bip85Error as e:
+            logging.error(f"Failed to generate BIP-85 seed: {e}", exc_info=True)
+            print(colored(f"Error: Failed to generate BIP-85 seed: {e}", "red"))
+            sys.exit(1)
         except Exception as e:
             logging.error(f"Failed to generate BIP-85 seed: {e}", exc_info=True)
             print(colored(f"Error: Failed to generate BIP-85 seed: {e}", "red"))

src/tests/test_bip85_vectors.py

@@ -4,7 +4,7 @@ import pytest
 
 sys.path.append(str(Path(__file__).resolve().parents[1]))
 
-from local_bip85.bip85 import BIP85
+from local_bip85.bip85 import BIP85, Bip85Error
 
 MASTER_XPRV = "xprv9s21ZrQH143K2LBWUUQRFXhucrQqBpKdRRxNVq2zBqsx8HVqFk2uYo8kmbaLLHRdqtQpUm98uKfu3vca1LqdGhUtyoFnCNkfmXRyPXLjbKb"
 
@@ -33,7 +33,7 @@ def test_bip85_symmetric_key(bip85):
 
 
 def test_invalid_params(bip85):
-    with pytest.raises(SystemExit):
+    with pytest.raises(Bip85Error):
         bip85.derive_mnemonic(index=0, words_num=15)
-    with pytest.raises(SystemExit):
+    with pytest.raises(Bip85Error):
         bip85.derive_mnemonic(index=-1, words_num=12)

src/utils/password_prompt.py

@@ -29,6 +29,12 @@ colorama_init()
 logger = logging.getLogger(__name__)
 
 
+class PasswordPromptError(Exception):
+    """Exception raised for password prompt errors."""
+
+    pass
+
+
 def prompt_new_password() -> str:
     """
     Prompts the user to enter and confirm a new password for encrypting the parent seed.
@@ -40,7 +46,7 @@ def prompt_new_password() -> str:
         str: The confirmed password entered by the user.
 
     Raises:
-        SystemExit: If the user fails to provide a valid password after multiple attempts.
+        PasswordPromptError: If the user fails to provide a valid password after multiple attempts.
     """
     max_retries = 5
     attempts = 0
@@ -87,7 +93,7 @@ def prompt_new_password() -> str:
         except KeyboardInterrupt:
             print(colored("\nOperation cancelled by user.", "yellow"))
             logging.info("Password prompt interrupted by user.")
-            sys.exit(0)
+            raise PasswordPromptError("Operation cancelled by user")
         except Exception as e:
             logging.error(
                 f"Unexpected error during password prompt: {e}", exc_info=True
@@ -97,7 +103,7 @@ def prompt_new_password() -> str:
 
     print(colored("Maximum password attempts exceeded. Exiting.", "red"))
     logging.error("User failed to provide a valid password after multiple attempts.")
-    sys.exit(1)
+    raise PasswordPromptError("Maximum password attempts exceeded")
 
 
 def prompt_existing_password(prompt_message: str = "Enter your password: ") -> str:
@@ -113,7 +119,7 @@ def prompt_existing_password(prompt_message: str = "Enter your password: ") -> s
         str: The password entered by the user.
 
     Raises:
-        SystemExit: If the user interrupts the operation.
+        PasswordPromptError: If the user interrupts the operation.
     """
     try:
         password = getpass.getpass(prompt=prompt_message).strip()
@@ -121,7 +127,7 @@ def prompt_existing_password(prompt_message: str = "Enter your password: ") -> s
         if not password:
             print(colored("Error: Password cannot be empty.", "red"))
             logging.warning("User attempted to enter an empty password.")
-            sys.exit(1)
+            raise PasswordPromptError("Password cannot be empty")
 
         # Normalize the password to NFKD form
         normalized_password = unicodedata.normalize("NFKD", password)
@@ -131,13 +137,13 @@ def prompt_existing_password(prompt_message: str = "Enter your password: ") -> s
     except KeyboardInterrupt:
         print(colored("\nOperation cancelled by user.", "yellow"))
         logging.info("Existing password prompt interrupted by user.")
-        sys.exit(0)
+        raise PasswordPromptError("Operation cancelled by user")
     except Exception as e:
         logging.error(
             f"Unexpected error during existing password prompt: {e}", exc_info=True
         )
         print(colored(f"Error: {e}", "red"))
-        sys.exit(1)
+        raise PasswordPromptError(str(e))
 
 
 def confirm_action(
@@ -154,7 +160,7 @@ def confirm_action(
         bool: True if the user confirms the action, False otherwise.
 
     Raises:
-        SystemExit: If the user interrupts the operation.
+        PasswordPromptError: If the user interrupts the operation.
     """
     try:
         while True:
@@ -171,13 +177,13 @@ def confirm_action(
     except KeyboardInterrupt:
         print(colored("\nOperation cancelled by user.", "yellow"))
         logging.info("Action confirmation interrupted by user.")
-        sys.exit(0)
+        raise PasswordPromptError("Operation cancelled by user")
     except Exception as e:
         logging.error(
             f"Unexpected error during action confirmation: {e}", exc_info=True
         )
         print(colored(f"Error: {e}", "red"))
-        sys.exit(1)
+        raise PasswordPromptError(str(e))
 
 
 def prompt_for_password() -> str: