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
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
|
from __future__ import annotations
from abc import abstractmethod
from typing import Optional, Tuple, Type, TypeVar
from .base_state import BaseState
from .crypto_provider import HashFunction
from .identity_key_pair import IdentityKeyPair
from .migrations import parse_base_state_model
from .models import BaseStateModel
from .signed_pre_key_pair import SignedPreKeyPair
from .types import Bundle, IdentityKeyFormat, Header, JSONObject
__all__ = [
"State"
]
StateTypeT = TypeVar("StateTypeT", bound="State")
class State(BaseState):
"""
This class is the core of this X3DH implementation. It manages the own :class:`~x3dh.types.Bundle` and
offers methods to perform key agreements with other parties. Use :class:`~x3dh.base_state.BaseState`
directly if manual state management is needed. Note that you can still use the methods available for
manual state management, but doing so shouldn't be required.
Warning:
:meth:`rotate_signed_pre_key` should be called periodically to check whether the signed pre key needs
to be rotated and to perform the rotation if necessary.
"""
def __init__(self) -> None:
super().__init__()
# Just the type definitions here
self.__signed_pre_key_rotation_period: int
self.__pre_key_refill_threshold: int
self.__pre_key_refill_target: int
@classmethod
def create(
cls: Type[StateTypeT],
identity_key_format: IdentityKeyFormat,
hash_function: HashFunction,
info: bytes,
identity_key_pair: Optional[IdentityKeyPair] = None,
signed_pre_key_rotation_period: int = 7 * 24 * 60 * 60,
pre_key_refill_threshold: int = 99,
pre_key_refill_target: int = 100
) -> StateTypeT:
"""
Args:
identity_key_format: The format in which the identity public key is included in bundles/headers.
hash_function: A 256 or 512-bit hash function.
info: A (byte) string identifying the application.
signed_pre_key_rotation_period: Rotate the signed pre key after this amount of time in seconds.
pre_key_refill_threshold: Threshold for refilling the pre keys.
pre_key_refill_target: When less then ``pre_key_refill_threshold`` pre keys are available,
generate new ones until there are ``pre_key_refill_target`` pre keys again.
identity_key_pair: If set, use the given identity key pair instead of generating a new one.
Returns:
A configured instance of :class:`~x3dh.state.State`.
"""
# pylint: disable=protected-access
if signed_pre_key_rotation_period < 1:
raise ValueError(
"Invalid value passed for the `signed_pre_key_rotation_period` parameter. The signed pre key"
" rotation period must be at least one day."
)
if not 1 <= pre_key_refill_threshold <= pre_key_refill_target:
raise ValueError(
"Invalid value(s) passed for the `pre_key_refill_threshold` / `pre_key_refill_target`"
" parameter(s). `pre_key_refill_threshold` must be greater than or equal to '1' and lower"
" than or equal to `pre_key_refill_target`."
)
self = super().create(identity_key_format, hash_function, info, identity_key_pair)
self.__signed_pre_key_rotation_period = signed_pre_key_rotation_period
self.__pre_key_refill_threshold = pre_key_refill_threshold
self.__pre_key_refill_target = pre_key_refill_target
self.generate_pre_keys(pre_key_refill_target)
# I believe this is a false positive by pylint
self._publish_bundle(self.bundle) # pylint: disable=no-member
return self
####################
# abstract methods #
####################
@abstractmethod
def _publish_bundle(self, bundle: Bundle) -> None:
"""
Args:
bundle: The bundle to publish, overwriting previously published data.
Note:
In addition to publishing the bundle, this method can be used as a trigger to persist the state.
Persisting the state in this method guarantees always remaining up-to-date.
Note:
This method is called from :meth:`create`, before :meth:`create` has returned the instance. Thus,
modifications to the object (``self``, in case of subclasses) may not have happened when this
method is called.
Note:
Even though this method is expected to perform I/O, it is deliberately not marked as async, since
completion of the I/O operation is not a requirement for the program flow to continue, and making
this method async would complicate API design with regards to inheritance from
:class:`~x3dh.base_state.BaseState`.
"""
raise NotImplementedError("Create a subclass of State and implement `_publish_bundle`.")
#################
# serialization #
#################
@classmethod
def from_model(
cls: Type[StateTypeT],
model: BaseStateModel,
identity_key_format: IdentityKeyFormat,
hash_function: HashFunction,
info: bytes,
signed_pre_key_rotation_period: int = 7 * 24 * 60 * 60,
pre_key_refill_threshold: int = 99,
pre_key_refill_target: int = 100
) -> StateTypeT:
"""
Args:
model: The pydantic model holding the internal state of a :class:`State`, as produced by
:attr:`~x3dh.base_state.BaseState.model`.
identity_key_format: The format in which the identity public key is included in bundles/headers.
hash_function: A 256 or 512-bit hash function.
info: A (byte) string identifying the application.
signed_pre_key_rotation_period: Rotate the signed pre key after this amount of time in seconds.
pre_key_refill_threshold: Threshold for refilling the pre keys.
pre_key_refill_target: When less then ``pre_key_refill_threshold`` pre keys are available,
generate new ones until there are ``pre_key_refill_target`` pre keys again.
Returns:
A configured instance of :class:`State`, with internal state restored from the model.
Warning:
Migrations are not provided via the :attr:`~x3dh.base_state.BaseState.model`/:meth:`from_model`
API. Use :attr:`~x3dh.base_state.BaseState.json`/:meth:`from_json` instead. Refer to
:ref:`serialization_and_migration` in the documentation for details.
"""
# pylint: disable=protected-access
if signed_pre_key_rotation_period < 1:
raise ValueError(
"Invalid value passed for the `signed_pre_key_rotation_period` parameter. The signed pre key"
" rotation period must be at least one day."
)
if not 1 <= pre_key_refill_threshold <= pre_key_refill_target:
raise ValueError(
"Invalid value(s) passed for the `pre_key_refill_threshold` / `pre_key_refill_target`"
" parameter(s). `pre_key_refill_threshold` must be greater than or equal to '1' and lower"
" than or equal to `pre_key_refill_target`."
)
self = super().from_model(model, identity_key_format, hash_function, info)
self.__signed_pre_key_rotation_period = signed_pre_key_rotation_period
self.__pre_key_refill_threshold = pre_key_refill_threshold
self.__pre_key_refill_target = pre_key_refill_target
self.rotate_signed_pre_key()
return self
@classmethod
def from_json(
cls: Type[StateTypeT],
serialized: JSONObject,
identity_key_format: IdentityKeyFormat,
hash_function: HashFunction,
info: bytes,
signed_pre_key_rotation_period: int = 7 * 24 * 60 * 60,
pre_key_refill_threshold: int = 99,
pre_key_refill_target: int = 100
) -> Tuple[StateTypeT, bool]:
"""
Args:
serialized: A JSON-serializable Python object holding the internal state of a :class:`State`,
as produced by :attr:`~x3dh.base_state.BaseState.json`.
identity_key_format: The format in which the identity public key is included in bundles/headers.
hash_function: A 256 or 512-bit hash function.
info: A (byte) string identifying the application.
signed_pre_key_rotation_period: Rotate the signed pre key after this amount of time in seconds.
pre_key_refill_threshold: Threshold for refilling the pre keys.
pre_key_refill_target: When less then ``pre_key_refill_threshold`` pre keys are available,
generate new ones until there are ``pre_key_refill_target`` pre keys again.
Returns:
A configured instance of :class:`State`, with internal state restored from the serialized data,
and a flag that indicates whether the bundle needed to be published. The latter was part of the
pre-stable serialization format and is handled automatically by this :meth:`from_json`
implementation.
"""
# pylint: disable=protected-access
model, bundle_needs_publish = parse_base_state_model(serialized)
self = cls.from_model(
model,
identity_key_format,
hash_function,
info,
signed_pre_key_rotation_period,
pre_key_refill_threshold,
pre_key_refill_target
)
if bundle_needs_publish:
# I believe this is a false positive by pylint
self._publish_bundle(self.bundle) # pylint: disable=no-member
return self, False
#################################
# key generation and management #
#################################
def rotate_signed_pre_key(self, force: bool = False) -> None:
"""
Check whether the signed pre key is due for rotation, and rotate it if necessary. Call this method
periodically to make sure the signed pre key is always up to date.
Args:
force: Whether to force rotation regardless of the age of the current signed pre key.
"""
if force or self.signed_pre_key_age() > self.__signed_pre_key_rotation_period:
super().rotate_signed_pre_key()
self._publish_bundle(self.bundle)
#################
# key agreement #
#################
async def get_shared_secret_passive(
self,
header: Header,
associated_data_appendix: bytes = b"",
require_pre_key: bool = True
) -> Tuple[bytes, bytes, SignedPreKeyPair]:
"""
Perform an X3DH key agreement, passively.
Args:
header: The header received from the active party.
associated_data_appendix: Additional information to append to the associated data, like usernames,
certificates or other identifying information.
require_pre_key: Use this flag to abort the key agreement if the active party did not use a pre
key.
Returns:
The shared secret and the associated data shared between both parties, and the signed pre key pair
that was used during the key exchange, for use by follow-up protocols.
Raises:
KeyAgreementException: If an error occurs during the key agreement. The exception message will
contain (human-readable) details.
"""
shared_secret, associated_data, signed_pre_key_pair = await super().get_shared_secret_passive(
header,
associated_data_appendix,
require_pre_key
)
# If a pre key was used, remove it from the pool and refill the pool if necessary
if header.pre_key is not None:
self.delete_pre_key(header.pre_key)
if self.get_num_visible_pre_keys() < self.__pre_key_refill_threshold:
self.generate_pre_keys(self.__pre_key_refill_target - self.get_num_visible_pre_keys())
self._publish_bundle(self.bundle)
return shared_secret, associated_data, signed_pre_key_pair
|
