Coverage for src/lib2fas/core.py: 100%

89 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2025-01-12 17:34 +0100

1""" 

2This file contains the core functionality. 

3""" 

4 

5import sys 

6import typing 

7from collections import defaultdict 

8from pathlib import Path 

9from typing import Optional 

10 

11import pyjson5 

12 

13from ._security import decrypt, keyring_manager 

14from ._types import TwoFactorAuthDetails, into_class 

15from .utils import flatten, fuzzy_match 

16 

17T_TwoFactorAuthDetails = typing.TypeVar("T_TwoFactorAuthDetails", bound=TwoFactorAuthDetails) 

18 

19 

20class TwoFactorStorage(typing.Generic[T_TwoFactorAuthDetails]): 

21 """ 

22 Container to make working with a collection of 2fas services easier. 

23 """ 

24 

25 _multidict: defaultdict[str, list[T_TwoFactorAuthDetails]] 

26 count: int 

27 

28 def __init__(self, _klass: typing.Type[T_TwoFactorAuthDetails] = None) -> None: 

29 """ 

30 Create a new instance, usually done by `new_auth_storage()`. 

31 

32 Args: 

33 _klass: _klass is purely for annotation atm 

34 """ 

35 self._multidict = defaultdict(list) # one name can map to multiple keys 

36 self.count = 0 

37 

38 def __len__(self) -> int: 

39 """ 

40 The length of the storage is the amount of items in it. 

41 """ 

42 return self.count 

43 

44 def __bool__(self) -> bool: 

45 """ 

46 The storage is truthy if it has any items. 

47 """ 

48 return self.count > 0 

49 

50 def add(self, entries: list[T_TwoFactorAuthDetails]) -> None: 

51 """ 

52 Extend the storage with new items. 

53 """ 

54 for entry in entries: 

55 name = (entry.name or "").lower() 

56 self._multidict[name].append(entry) 

57 

58 self.count += len(entries) 

59 

60 def __getitem__(self, item: str) -> "list[T_TwoFactorAuthDetails]": 

61 """ 

62 Get a service via the class[property] syntax. 

63 """ 

64 # 

65 return self._multidict[item.lower()] 

66 

67 def keys(self) -> list[str]: 

68 """ 

69 Return a list of services in this storage. 

70 

71 Usage: 

72 storage.keys() 

73 """ 

74 return list(self._multidict.keys()) 

75 

76 def items(self) -> typing.Generator[tuple[str, list[T_TwoFactorAuthDetails]], None, None]: 

77 """ 

78 Loop through tuples of key and values. 

79 

80 Usage: 

81 for key, value in storage.items(): ... 

82 # (like dict.items()) 

83 """ 

84 yield from self._multidict.items() 

85 

86 def _fuzzy_find(self, find: typing.Optional[str], fuzz_threshold: int) -> list[T_TwoFactorAuthDetails]: 

87 if not find: 

88 # don't loop 

89 return list(self) 

90 

91 all_items = self._multidict.items() 

92 

93 find = find.lower() 

94 # if nothing found exactly, try again but fuzzy (could be slower) 

95 # search in key: 

96 fuzzy = [ 

97 # search in key 

98 v 

99 for k, v in all_items 

100 if fuzzy_match(k.lower(), find) > fuzz_threshold 

101 ] 

102 if fuzzy and (flat := flatten(fuzzy)): 

103 return flat 

104 

105 # search in value: 

106 # str is short, repr is json 

107 return [ 

108 # search in value instead 

109 v 

110 for v in list(self) 

111 if fuzzy_match(repr(v).lower(), find) > fuzz_threshold 

112 ] 

113 

114 def generate(self) -> list[tuple[str, str]]: 

115 """ 

116 Create TOTP codes for all services in this storage. 

117 """ 

118 return [(_.name, _.generate()) for _ in self] 

119 

120 def find( 

121 self, target: Optional[str] = None, fuzz_threshold: int = 75 

122 ) -> "TwoFactorStorage[T_TwoFactorAuthDetails]": 

123 """ 

124 Create a new storage object with a subset of items in this storage, filtered by the search query in 'target'. 

125 

126 First, an exact search is tried and if that fails, fuzzy matching is applied. 

127 """ 

128 target = (target or "").lower() 

129 # first try exact match: 

130 if items := self._multidict.get(target): 

131 return new_auth_storage(items) 

132 # else: fuzzy match: 

133 return new_auth_storage(self._fuzzy_find(target, fuzz_threshold)) 

134 

135 def all(self) -> list[T_TwoFactorAuthDetails]: 

136 """ 

137 Return a list of services. 

138 """ 

139 return list(self) 

140 

141 def __iter__(self) -> typing.Generator[T_TwoFactorAuthDetails, None, None]: 

142 """ 

143 Allows for-looping through this storage. 

144 """ 

145 for entries in self._multidict.values(): 

146 yield from entries 

147 

148 def __repr__(self) -> str: 

149 """ 

150 Representation for repr(). 

151 """ 

152 return f"<TwoFactorStorage with {len(self._multidict)} keys and {self.count} entries>" 

153 

154 

155def new_auth_storage(initial_items: list[T_TwoFactorAuthDetails] = None) -> TwoFactorStorage[T_TwoFactorAuthDetails]: 

156 """ 

157 Create an instance of TwoFactorStorage and maybe load some items into it. 

158 """ 

159 storage: TwoFactorStorage[T_TwoFactorAuthDetails] = TwoFactorStorage() 

160 

161 if initial_items: 

162 storage.add(initial_items) 

163 

164 return storage 

165 

166 

167def load_services( 

168 filename: str | Path, _max_retries: int = 0, passphrase: Optional[str] = None 

169) -> TwoFactorStorage[TwoFactorAuthDetails] | None: 

170 """ 

171 Given a 2fas file, try to decrypt it (via stored password in keyring or by querying user) \ 

172 and load into a TwoFactorStorage object. 

173 

174 Args: 

175 filename: Path to a .2fas file 

176 _max_retries: how many password guesses are allowed? (default = unlimited) 

177 passphrase: password for the supplied 2fas file; leave empty to query the user. 

178 Note: when using the passphrase option, _max_retries is ignored and the keyring is not used. 

179 

180 Returns: 

181 A TwoFactorStorage instance, or None if e.g. the requested .2fas file does not exist. 

182 

183 Raises: 

184 PermissionError on invalid password. 

185 """ 

186 filepath = Path(filename).expanduser() 

187 

188 if not filepath.exists(): 

189 return None 

190 

191 with filepath.open() as f: 

192 data_raw = f.read() 

193 data = pyjson5.loads(data_raw) 

194 

195 storage: TwoFactorStorage[TwoFactorAuthDetails] = new_auth_storage() 

196 

197 if decrypted := data["services"]: 

198 services = into_class(decrypted, TwoFactorAuthDetails) 

199 storage.add(services) 

200 return storage 

201 

202 encrypted = data["servicesEncrypted"] 

203 

204 if passphrase is not None: 

205 # could raise PermissionError 

206 entries = decrypt(encrypted, passphrase) 

207 storage.add(entries) 

208 return storage 

209 else: 

210 retries = 0 

211 while True: 

212 # fmt: off 

213 password = ( 

214 keyring_manager.retrieve_credentials(str(filename)) 

215 or keyring_manager.save_credentials(str(filename)) 

216 ) 

217 # fmt: on 

218 

219 try: 

220 entries = decrypt(encrypted, password) 

221 storage.add(entries) 

222 return storage 

223 except PermissionError as e: 

224 retries += 1 # only really useful for pytest 

225 print(e, file=sys.stderr) 

226 keyring_manager.delete_credentials(str(filename)) 

227 

228 if _max_retries and retries > _max_retries: 

229 raise e