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
« prev ^ index » next coverage.py v7.4.1, created at 2025-01-12 17:34 +0100
1"""
2This file contains the core functionality.
3"""
5import sys
6import typing
7from collections import defaultdict
8from pathlib import Path
9from typing import Optional
11import pyjson5
13from ._security import decrypt, keyring_manager
14from ._types import TwoFactorAuthDetails, into_class
15from .utils import flatten, fuzzy_match
17T_TwoFactorAuthDetails = typing.TypeVar("T_TwoFactorAuthDetails", bound=TwoFactorAuthDetails)
20class TwoFactorStorage(typing.Generic[T_TwoFactorAuthDetails]):
21 """
22 Container to make working with a collection of 2fas services easier.
23 """
25 _multidict: defaultdict[str, list[T_TwoFactorAuthDetails]]
26 count: int
28 def __init__(self, _klass: typing.Type[T_TwoFactorAuthDetails] = None) -> None:
29 """
30 Create a new instance, usually done by `new_auth_storage()`.
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
38 def __len__(self) -> int:
39 """
40 The length of the storage is the amount of items in it.
41 """
42 return self.count
44 def __bool__(self) -> bool:
45 """
46 The storage is truthy if it has any items.
47 """
48 return self.count > 0
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)
58 self.count += len(entries)
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()]
67 def keys(self) -> list[str]:
68 """
69 Return a list of services in this storage.
71 Usage:
72 storage.keys()
73 """
74 return list(self._multidict.keys())
76 def items(self) -> typing.Generator[tuple[str, list[T_TwoFactorAuthDetails]], None, None]:
77 """
78 Loop through tuples of key and values.
80 Usage:
81 for key, value in storage.items(): ...
82 # (like dict.items())
83 """
84 yield from self._multidict.items()
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)
91 all_items = self._multidict.items()
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
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 ]
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]
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'.
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))
135 def all(self) -> list[T_TwoFactorAuthDetails]:
136 """
137 Return a list of services.
138 """
139 return list(self)
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
148 def __repr__(self) -> str:
149 """
150 Representation for repr().
151 """
152 return f"<TwoFactorStorage with {len(self._multidict)} keys and {self.count} entries>"
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()
161 if initial_items:
162 storage.add(initial_items)
164 return storage
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.
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.
180 Returns:
181 A TwoFactorStorage instance, or None if e.g. the requested .2fas file does not exist.
183 Raises:
184 PermissionError on invalid password.
185 """
186 filepath = Path(filename).expanduser()
188 if not filepath.exists():
189 return None
191 with filepath.open() as f:
192 data_raw = f.read()
193 data = pyjson5.loads(data_raw)
195 storage: TwoFactorStorage[TwoFactorAuthDetails] = new_auth_storage()
197 if decrypted := data["services"]:
198 services = into_class(decrypted, TwoFactorAuthDetails)
199 storage.add(services)
200 return storage
202 encrypted = data["servicesEncrypted"]
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
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))
228 if _max_retries and retries > _max_retries:
229 raise e