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

83 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-01-29 11:26 +0100

1""" 

2This file contains the core functionality. 

3""" 

4 

5import json 

6import sys 

7import typing 

8from collections import defaultdict 

9from pathlib import Path 

10from typing import Optional 

11 

12from ._security import decrypt, keyring_manager 

13from ._types import TwoFactorAuthDetails, into_class 

14from .utils import flatten, fuzzy_match 

15 

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

17 

18 

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

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

21 count: int 

22 

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

24 # _klass is purely for annotation atm 

25 

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

27 self.count = 0 

28 

29 def __len__(self) -> int: 

30 return self.count 

31 

32 def __bool__(self) -> bool: 

33 return self.count > 0 

34 

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

36 for entry in entries: 

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

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

39 

40 self.count += len(entries) 

41 

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

43 # class[property] syntax 

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

45 

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

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

48 

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

50 yield from self._multidict.items() 

51 

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

53 if not find: 

54 # don't loop 

55 return list(self) 

56 

57 all_items = self._multidict.items() 

58 

59 find = find.lower() 

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

61 # search in key: 

62 fuzzy = [ 

63 # search in key 

64 v 

65 for k, v in all_items 

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

67 ] 

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

69 return flat 

70 

71 # search in value: 

72 # str is short, repr is json 

73 return [ 

74 # search in value instead 

75 v 

76 for v in list(self) 

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

78 ] 

79 

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

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

82 

83 def find( 

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

85 ) -> "TwoFactorStorage[T_TwoFactorAuthDetails]": 

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

87 # first try exact match: 

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

89 return new_auth_storage(items) 

90 # else: fuzzy match: 

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

92 

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

94 return list(self) 

95 

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

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

98 yield from entries 

99 

100 def __repr__(self) -> str: 

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

102 

103 

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

105 storage: TwoFactorStorage[T_TwoFactorAuthDetails] = TwoFactorStorage() 

106 

107 if initial_items: 

108 storage.add(initial_items) 

109 

110 return storage 

111 

112 

113def load_services(filename: str, _max_retries: int = 0) -> TwoFactorStorage[TwoFactorAuthDetails]: 

114 filepath = Path(filename).expanduser() 

115 with filepath.open() as f: 

116 data_raw = f.read() 

117 data = json.loads(data_raw) 

118 

119 storage: TwoFactorStorage[TwoFactorAuthDetails] = new_auth_storage() 

120 

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

122 services = into_class(decrypted, TwoFactorAuthDetails) 

123 storage.add(services) 

124 return storage 

125 

126 encrypted = data["servicesEncrypted"] 

127 

128 retries = 0 

129 while True: 

130 password = keyring_manager.retrieve_credentials(filename) or keyring_manager.save_credentials(filename) 

131 try: 

132 entries = decrypt(encrypted, password) 

133 storage.add(entries) 

134 return storage 

135 except PermissionError as e: 

136 retries += 1 # only really useful for pytest 

137 print(e, file=sys.stderr) 

138 keyring_manager.delete_credentials(filename) 

139 

140 if _max_retries and retries > _max_retries: 

141 raise e