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
« prev ^ index » next coverage.py v7.4.1, created at 2024-01-29 11:26 +0100
1"""
2This file contains the core functionality.
3"""
5import json
6import sys
7import typing
8from collections import defaultdict
9from pathlib import Path
10from typing import Optional
12from ._security import decrypt, keyring_manager
13from ._types import TwoFactorAuthDetails, into_class
14from .utils import flatten, fuzzy_match
16T_TwoFactorAuthDetails = typing.TypeVar("T_TwoFactorAuthDetails", bound=TwoFactorAuthDetails)
19class TwoFactorStorage(typing.Generic[T_TwoFactorAuthDetails]):
20 _multidict: defaultdict[str, list[T_TwoFactorAuthDetails]]
21 count: int
23 def __init__(self, _klass: typing.Type[T_TwoFactorAuthDetails] = None) -> None:
24 # _klass is purely for annotation atm
26 self._multidict = defaultdict(list) # one name can map to multiple keys
27 self.count = 0
29 def __len__(self) -> int:
30 return self.count
32 def __bool__(self) -> bool:
33 return self.count > 0
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)
40 self.count += len(entries)
42 def __getitem__(self, item: str) -> "list[T_TwoFactorAuthDetails]":
43 # class[property] syntax
44 return self._multidict[item.lower()]
46 def keys(self) -> list[str]:
47 return list(self._multidict.keys())
49 def items(self) -> typing.Generator[tuple[str, list[T_TwoFactorAuthDetails]], None, None]:
50 yield from self._multidict.items()
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)
57 all_items = self._multidict.items()
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
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 ]
80 def generate(self) -> list[tuple[str, str]]:
81 return [(_.name, _.generate()) for _ in self]
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))
93 def all(self) -> list[T_TwoFactorAuthDetails]:
94 return list(self)
96 def __iter__(self) -> typing.Generator[T_TwoFactorAuthDetails, None, None]:
97 for entries in self._multidict.values():
98 yield from entries
100 def __repr__(self) -> str:
101 return f"<TwoFactorStorage with {len(self._multidict)} keys and {self.count} entries>"
104def new_auth_storage(initial_items: list[T_TwoFactorAuthDetails] = None) -> TwoFactorStorage[T_TwoFactorAuthDetails]:
105 storage: TwoFactorStorage[T_TwoFactorAuthDetails] = TwoFactorStorage()
107 if initial_items:
108 storage.add(initial_items)
110 return storage
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)
119 storage: TwoFactorStorage[TwoFactorAuthDetails] = new_auth_storage()
121 if decrypted := data["services"]:
122 services = into_class(decrypted, TwoFactorAuthDetails)
123 storage.add(services)
124 return storage
126 encrypted = data["servicesEncrypted"]
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)
140 if _max_retries and retries > _max_retries:
141 raise e