Coverage for ledger/pin.py: 87%

1# The MIT License (MIT) 

2# 

3# Copyright (c) 2021 RSK Labs Ltd 

4# 

5# Permission is hereby granted, free of charge, to any person obtaining a copy of 

6# this software and associated documentation files (the "Software"), to deal in 

7# the Software without restriction, including without limitation the rights to 

8# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 

9# of the Software, and to permit persons to whom the Software is furnished to do 

10# so, subject to the following conditions: 

11# 

12# The above copyright notice and this permission notice shall be included in all 

13# copies or substantial portions of the Software. 

14# 

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

21# SOFTWARE. 

22 

23import os 

24import random 

25import string 

26import logging 

27 

28 

29class PinError(RuntimeError): 

30 pass 

31 

32 

33class BasePin: 

34 PIN_LENGTH = 8 

35 POSSIBLE_CHARS = string.ascii_letters + string.digits 

36 ALPHA_CHARS = string.ascii_letters 

37 

38 @classmethod 

39 def generate_pin(cls): 

40 random.seed() 

41 pin = None 

42 while pin is None or not cls.is_valid(pin.encode()): 

43 pin = "" 

44 for i in range(cls.PIN_LENGTH): 

45 pin += random.choice(cls.POSSIBLE_CHARS) 

46 return pin.encode() 

47 

48 @classmethod 

49 def is_valid(cls, pin, any_pin=False): 

50 if type(pin) != bytes: 

51 return False 

52 

53 if not all(map(lambda c: chr(c) in cls.POSSIBLE_CHARS, pin)): 

54 return False 

55 

56 if any_pin: 

57 return True 

58 

59 if len(pin) != cls.PIN_LENGTH: 

60 return False 

61 

62 return any(map(lambda c: chr(c) in cls.ALPHA_CHARS, pin)) 

63 

64 

65# Handles PIN initialization and loading 

66class FileBasedPin(BasePin): 

67 @classmethod 

68 def new(cls, path): 

69 # Generate a new random pin and save to the given path. Return the 

70 # generated pin 

71 with open(path, 'wb') as file: 

72 pin = cls.generate_pin() 

73 file.write(pin) 

74 return pin 

75 

76 def __init__(self, path, default_pin, force_change=False): 

77 self.logger = logging.getLogger("pin") 

78 self._path = path 

79 

80 # Check pin file existence 

81 try: 

82 self.logger.info("Checking whether pin file '%s' exists", path) 

83 pin_file_exists = os.path.isfile(path) 

84 except Exception as e: 

85 self._error("Error checking file existence: %s" % format(e)) 

86 

87 # Load pin 

88 if pin_file_exists: 

89 try: 

90 self.logger.info("Loading PIN from '%s'", path) 

91 with open(path, 'rb') as file: 

92 self._pin = file.read().strip() 

93 except Exception as e: 

94 self._error("Error loading pin: %s" % format(e)) 

95 else: 

96 self.logger.info("Using default PIN") 

97 self._pin = default_pin 

98 

99 # Make sure pin is valid 

100 if not self.is_valid(self._pin): 

101 self._error("Invalid pin: %s" % self._pin) 

102 

103 # Pin needs to be changed if no pin file found or a pin change is forced 

104 self._needs_change = force_change or not pin_file_exists 

105 self._changing = False 

106 

107 # Returns the current pin 

108 def get_pin(self): 

109 return self._pin 

110 

111 def get_new_pin(self): 

112 if not self._changing: 

113 return None 

114 

115 return self._new_pin 

116 

117 # Indicates whether the current pin needs changing 

118 def needs_change(self): 

119 return self._needs_change 

120 

121 # Starts a pin change and generates a new pin 

122 def start_change(self): 

123 self.logger.info("Starting PIN change") 

124 

125 if self._changing or not self._needs_change: 

126 return 

127 

128 self._changing = True 

129 self._new_pin = self.generate_pin() 

130 

131 self.logger.info("New PIN generated") 

132 

133 # Finalizes a pin change by commiting the changes to disk 

134 def commit_change(self): 

135 self.logger.info("Commiting PIN change to disk") 

136 if not self._changing: 

137 return 

138 

139 try: 

140 with open(self._path, "wb") as file: 

141 file.write(self._new_pin) 

142 except Exception as e: 

143 self._error("Error commiting: %s" % format(e)) 

144 

145 self._pin = self._new_pin 

146 self._new_pin = None 

147 self._changing = False 

148 self._needs_change = False 

149 

150 self.logger.info("PIN change committed to disk") 

151 

152 # Aborts a pin change 

153 def abort_change(self): 

154 self.logger.info("Aborting PIN change") 

155 

156 if not self._changing: 

157 return 

158 

159 self._new_pin = None 

160 self._changing = False 

161 

162 self.logger.info("PIN change aborted") 

163 

164 def _error(self, msg): 

165 self.logger.error(msg) 

166 raise PinError(msg)