diff --git a/silkaj/crypto_tools.py b/silkaj/crypto_tools.py index bf3d011471d19428b13b394e6f810b565d520c11..ba183a84a6ee50cf5ab44ba88f9054b47d358a42 100644 --- a/silkaj/crypto_tools.py +++ b/silkaj/crypto_tools.py @@ -15,6 +15,7 @@ import hashlib import re +from typing import Any, Optional, Union import base58 @@ -27,7 +28,7 @@ CHECKSUM_PATTERN = f"[1-9A-HJ-NP-Za-km-z]{{{CHECKSUM_SIZE}}}" PUBKEY_CHECKSUM_PATTERN = f"^{PUBKEY_PATTERN}:{CHECKSUM_PATTERN}$" -def is_pubkey_and_check(pubkey): +def is_pubkey_and_check(pubkey: str) -> Union[str, bool]: """ Checks if the given argument contains a pubkey. If so, verifies the checksum if needed and returns the pubkey. @@ -41,7 +42,7 @@ def is_pubkey_and_check(pubkey): return False -def check_pubkey_format(pubkey, display_error=True): +def check_pubkey_format(pubkey: str, display_error: bool = True) -> Optional[bool]: """ Checks if a pubkey has a checksum. Exits if the pubkey is invalid. @@ -52,10 +53,10 @@ def check_pubkey_format(pubkey, display_error=True): return True elif display_error: message_exit(f"Error: bad format for following public key: {pubkey}") - return + return None -def validate_checksum(pubkey_checksum): +def validate_checksum(pubkey_checksum: str) -> Any: """ Check pubkey checksum after the pubkey, delimited by ":". If check pass: return pubkey @@ -68,12 +69,13 @@ def validate_checksum(pubkey_checksum): f"Error: public key '{pubkey}' does not match checksum '{checksum}'.\n\ Please verify the public key." ) + return None -def gen_checksum(pubkey): +def gen_checksum(pubkey: str) -> str: """ Returns the checksum of the input pubkey (encoded in b58) """ pubkey_byte = base58.b58decode(pubkey) hash = hashlib.sha256(hashlib.sha256(pubkey_byte).digest()).digest() - return base58.b58encode(hash)[:3].decode("utf-8") + return str(base58.b58encode(hash)[:3].decode("utf-8"))