63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
|
from automat import MethodicalMachine
|
||
|
|
from spake2 import SPAKE2_Symmetric
|
||
|
|
from hkdf import Hkdf
|
||
|
|
from nacl.secret import SecretBox
|
||
|
|
from .util import (to_bytes, bytes_to_hexstr, hexstr_to_bytes)
|
||
|
|
|
||
|
|
def HKDF(skm, outlen, salt=None, CTXinfo=b""):
|
||
|
|
return Hkdf(salt, skm).expand(CTXinfo, outlen)
|
||
|
|
|
||
|
|
def derive_key(key, purpose, length=SecretBox.KEY_SIZE):
|
||
|
|
if not isinstance(key, type(b"")): raise TypeError(type(key))
|
||
|
|
if not isinstance(purpose, type(b"")): raise TypeError(type(purpose))
|
||
|
|
if not isinstance(length, int): raise TypeError(type(length))
|
||
|
|
return HKDF(key, length, CTXinfo=purpose)
|
||
|
|
|
||
|
|
class SendMachine(object):
|
||
|
|
m = MethodicalMachine()
|
||
|
|
def __init__(self, timing):
|
||
|
|
self._timing = timing
|
||
|
|
def set_mailbox(self, mailbox):
|
||
|
|
self._mailbox = mailbox
|
||
|
|
|
||
|
|
@m.state(initial=True)
|
||
|
|
def S0_no_key(self): pass
|
||
|
|
@m.state()
|
||
|
|
def S1_verified_key(self): pass
|
||
|
|
|
||
|
|
def got_pake(self, payload):
|
||
|
|
if "pake_v1" in payload:
|
||
|
|
self.got_pake_good(hexstr_to_bytes(payload["pake_v1"]))
|
||
|
|
else:
|
||
|
|
self.got_pake_bad()
|
||
|
|
|
||
|
|
@m.input()
|
||
|
|
def got_verified_key(self, key): pass
|
||
|
|
@m.input()
|
||
|
|
def send(self, phase, payload): pass
|
||
|
|
|
||
|
|
@m.output()
|
||
|
|
def queue(self, phase, payload):
|
||
|
|
self._queue.append((phase, payload))
|
||
|
|
@m.output()
|
||
|
|
def record_key(self, key):
|
||
|
|
self._key = key
|
||
|
|
@m.output()
|
||
|
|
def drain(self, key):
|
||
|
|
del key
|
||
|
|
for (phase, payload) in self._queue:
|
||
|
|
self._encrypt_and_send(phase, payload)
|
||
|
|
@m.output()
|
||
|
|
def deliver(self, phase, payload):
|
||
|
|
self._encrypt_and_send(phase, payload)
|
||
|
|
|
||
|
|
def _encrypt_and_send(self, phase, payload):
|
||
|
|
data_key = self._derive_phase_key(self._side, phase)
|
||
|
|
encrypted = self._encrypt_data(data_key, plaintext)
|
||
|
|
self._mailbox.add_message(phase, encrypted)
|
||
|
|
|
||
|
|
S0_no_key.upon(send, enter=S0_no_key, outputs=[queue])
|
||
|
|
S0_no_key.upon(got_verified_key, enter=S1_verified_key,
|
||
|
|
outputs=[record_key, drain])
|
||
|
|
S1_verified_key.upon(send, enter=S1_verified_key, outputs=[deliver])
|