Coverage for silkaj/public_key.py: 97%

1# Copyright 2016-2024 Maël Azimi <m.a@moul.re> 

2# 

3# Silkaj is free software: you can redistribute it and/or modify 

4# it under the terms of the GNU Affero General Public License as published by 

5# the Free Software Foundation, either version 3 of the License, or 

6# (at your option) any later version. 

7# 

8# Silkaj is distributed in the hope that it will be useful, 

9# but WITHOUT ANY WARRANTY; without even the implied warranty of 

10# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

11# GNU Affero General Public License for more details. 

12# 

13# You should have received a copy of the GNU Affero General Public License 

14# along with Silkaj. If not, see <https://www.gnu.org/licenses/>. 

15 

16import hashlib 

17import re 

18from typing import Any, Optional, Union 

19 

20import base58 

21 

22from silkaj.constants import PUBKEY_PATTERN, SHORT_PUBKEY_SIZE 

23from silkaj.tools import message_exit 

24 

25PUBKEY_DELIMITED_PATTERN = f"^{PUBKEY_PATTERN}$" 

26CHECKSUM_SIZE = 3 

27CHECKSUM_PATTERN = f"[1-9A-HJ-NP-Za-km-z]{ {CHECKSUM_SIZE}} " 

28PUBKEY_CHECKSUM_PATTERN = f"^{PUBKEY_PATTERN}:{CHECKSUM_PATTERN}$" 

29 

30 

31def is_pubkey_and_check(pubkey: str) -> Union[str, bool]: 

32 """ 

33 Checks if the given argument contains a pubkey. 

34 If so, verifies the checksum if needed and returns the pubkey. 

35 Exits if the checksum is wrong. 

36 Else, return False 

37 """ 

38 if re.search(re.compile(PUBKEY_PATTERN), pubkey): 

39 if check_pubkey_format(pubkey, True): 

40 return validate_checksum(pubkey) 

41 return pubkey 

42 return False 

43 

44 

45def check_pubkey_format(pubkey: str, display_error: bool = True) -> Optional[bool]: 

46 """ 

47 Checks if a pubkey has a checksum. 

48 Exits if the pubkey is invalid. 

49 """ 

50 if re.search(re.compile(PUBKEY_DELIMITED_PATTERN), pubkey): 

51 return False 

52 if re.search(re.compile(PUBKEY_CHECKSUM_PATTERN), pubkey): 

53 return True 

54 if display_error: 

55 message_exit(f"Error: bad format for following public key: {pubkey}") 

56 return None 

57 

58 

59def validate_checksum(pubkey_checksum: str) -> Any: 

60 """ 

61 Check pubkey checksum after the pubkey, delimited by ":". 

62 If check pass: return pubkey 

63 Else: exit. 

64 """ 

65 pubkey, checksum = pubkey_checksum.split(":") 

66 if checksum == gen_checksum(pubkey): 

67 return pubkey 

68 message_exit( 

69 f"Error: public key '{pubkey}' does not match checksum '{checksum} \ 

70 '.\n\ 

71Please verify the public key.", 

72 ) 

73 return None 

74 

75 

76def gen_checksum(pubkey: str) -> str: 

77 """ 

78 Returns the checksum of the input pubkey (encoded in b58) 

79 """ 

80 pubkey_byte = base58.b58decode(pubkey) 

81 _hash = hashlib.sha256(hashlib.sha256(pubkey_byte).digest()).digest() 

82 return str(base58.b58encode(_hash)[:3].decode("utf-8")) 

83 

84 

85def gen_pubkey_checksum( 

86 pubkey: str, 

87 short: Optional[bool] = False, 

88 length: Optional[int] = SHORT_PUBKEY_SIZE, 

89) -> str: 

90 """ 

91 Returns "<pubkey>:<checksum>" in full form. 

92 returns `length` first chars of pubkey and checksum in short form. 

93 `length` defaults to SHORT_PUBKEY_SIZE. 

94 """ 

95 short_pubkey = f"{pubkey[:length]}…" if short else pubkey 

96 return f"{short_pubkey}:{gen_checksum(pubkey)}"