1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
|
import enum
from typing import Dict, FrozenSet, List, Optional, Tuple, Type
from typing_extensions import NamedTuple, assert_never
import omemo
__all__ = [
"TrustLevel",
"BundleStorageKey",
"DeviceListStorageKey",
"BundleStorage",
"DeviceListStorage",
"MessageQueue",
"make_session_manager_impl"
]
@enum.unique
class TrustLevel(enum.Enum):
"""
Trust levels modeling simple manual trust.
"""
TRUSTED: str = "TRUSTED"
UNDECIDED: str = "UNDECIDED"
DISTRUSTED: str = "DISTRUSTED"
class BundleStorageKey(NamedTuple):
# pylint: disable=invalid-name
"""
The key identifying a bundle in the tests.
"""
namespace: str
bare_jid: str
device_id: int
class DeviceListStorageKey(NamedTuple):
# pylint: disable=invalid-name
"""
The key identifying a device list in the tests.
"""
namespace: str
bare_jid: str
BundleStorage = Dict[BundleStorageKey, omemo.Bundle]
DeviceListStorage = Dict[DeviceListStorageKey, Dict[int, Optional[str]]]
MessageQueue = List[Tuple[str, omemo.Message]]
def make_session_manager_impl(
own_bare_jid: str,
bundle_storage: BundleStorage,
device_list_storage: DeviceListStorage,
message_queue: MessageQueue
) -> Type[omemo.SessionManager]:
"""
Args:
own_bare_jid: The bare JID of the account that will be used with the session manager instances created
from this implementation.
bundle_storage: The dictionary to "upload", "download" and delete bundles to/from.
device_list_storage: The dictionary to "upload" and "download" device lists to/from.
message_queue: The list to "send" automated messages to. The first entry of each tuple is the bare JID
of the recipient. The second entry is the message itself.
Returns:
A session manager implementation which sends/uploads/downloads/deletes data to/from the collections
given as parameters.
"""
class SessionManagerImpl(omemo.SessionManager):
@staticmethod
async def _upload_bundle(bundle: omemo.Bundle) -> None:
bundle_storage[BundleStorageKey(
namespace=bundle.namespace,
bare_jid=bundle.bare_jid,
device_id=bundle.device_id
)] = bundle
@staticmethod
async def _download_bundle(namespace: str, bare_jid: str, device_id: int) -> omemo.Bundle:
try:
return bundle_storage[BundleStorageKey(
namespace=namespace,
bare_jid=bare_jid,
device_id=device_id
)]
except KeyError as e:
raise omemo.BundleDownloadFailed() from e
@staticmethod
async def _delete_bundle(namespace: str, device_id: int) -> None:
try:
bundle_storage.pop(BundleStorageKey(
namespace=namespace,
bare_jid=own_bare_jid,
device_id=device_id
))
except KeyError as e:
raise omemo.BundleDeletionFailed() from e
@staticmethod
async def _upload_device_list(namespace: str, device_list: Dict[int, Optional[str]]) -> None:
device_list_storage[DeviceListStorageKey(
namespace=namespace,
bare_jid=own_bare_jid
)] = device_list
@staticmethod
async def _download_device_list(namespace: str, bare_jid: str) -> Dict[int, Optional[str]]:
try:
return device_list_storage[DeviceListStorageKey(
namespace=namespace,
bare_jid=bare_jid
)]
except KeyError:
return {}
async def _evaluate_custom_trust_level(self, device: omemo.DeviceInformation) -> omemo.TrustLevel:
try:
trust_level = TrustLevel(device.trust_level_name)
except ValueError as e:
raise omemo.UnknownTrustLevel() from e
if trust_level is TrustLevel.TRUSTED:
return omemo.TrustLevel.TRUSTED
if trust_level is TrustLevel.UNDECIDED:
return omemo.TrustLevel.UNDECIDED
if trust_level is TrustLevel.DISTRUSTED:
return omemo.TrustLevel.DISTRUSTED
assert_never(trust_level)
async def _make_trust_decision(
self,
undecided: FrozenSet[omemo.DeviceInformation],
identifier: Optional[str]
) -> None:
for device in undecided:
await self.set_trust(device.bare_jid, device.identity_key, TrustLevel.TRUSTED.name)
@staticmethod
async def _send_message(message: omemo.Message, bare_jid: str) -> None:
message_queue.append((bare_jid, message))
return SessionManagerImpl
|
