91 lines
2.5 KiB
Python
91 lines
2.5 KiB
Python
import struct
|
|
from asyncio import IncompleteReadError
|
|
|
|
from async_generator import async_generator, yield_
|
|
from nacl.secret import SecretBox
|
|
|
|
from .util import inc_nonce, split_chunks
|
|
|
|
HEADER_LENGTH = 2 + 16 + 16
|
|
MAX_SEGMENT_SIZE = 4 * 1024
|
|
TERMINATION_HEADER = (b'\x00' * 18)
|
|
|
|
|
|
def get_stream_pair(reader, writer, **kwargs):
|
|
"""Return a tuple with `(unbox_stream, box_stream)` (reader/writer).
|
|
|
|
:return: (:class:`secret_handshake.boxstream.UnboxStream`,
|
|
:class:`secret_handshake.boxstream.BoxStream`) """
|
|
box_args = {
|
|
'key': kwargs['encrypt_key'],
|
|
'nonce': kwargs['encrypt_nonce'],
|
|
}
|
|
unbox_args = {
|
|
'key': kwargs['decrypt_key'],
|
|
'nonce': kwargs['decrypt_nonce'],
|
|
}
|
|
return UnboxStream(reader, **unbox_args), BoxStream(writer, **box_args)
|
|
|
|
|
|
class UnboxStream(object):
|
|
def __init__(self, reader, key, nonce):
|
|
self.reader = reader
|
|
self.key = key
|
|
self.nonce = nonce
|
|
self.closed = False
|
|
|
|
async def read(self):
|
|
try:
|
|
data = await self.reader.readexactly(HEADER_LENGTH)
|
|
except IncompleteReadError:
|
|
self.closed = True
|
|
return None
|
|
|
|
box = SecretBox(self.key)
|
|
|
|
header = box.decrypt(data, self.nonce)
|
|
|
|
if header == TERMINATION_HEADER:
|
|
self.closed = True
|
|
return None
|
|
|
|
length = struct.unpack('>H', header[:2])[0]
|
|
mac = header[2:]
|
|
|
|
data = await self.reader.readexactly(length)
|
|
|
|
body = box.decrypt(mac + data, inc_nonce(self.nonce))
|
|
|
|
self.nonce = inc_nonce(inc_nonce(self.nonce))
|
|
return body
|
|
|
|
@async_generator
|
|
async def __aiter__(self):
|
|
while True:
|
|
data = await self.read()
|
|
if data is None:
|
|
return
|
|
await yield_(data)
|
|
|
|
|
|
class BoxStream(object):
|
|
def __init__(self, writer, key, nonce):
|
|
self.writer = writer
|
|
self.key = key
|
|
self.box = SecretBox(self.key)
|
|
self.nonce = nonce
|
|
|
|
def write(self, data):
|
|
for chunk in split_chunks(data, MAX_SEGMENT_SIZE):
|
|
body = self.box.encrypt(chunk, inc_nonce(self.nonce))[24:]
|
|
header = struct.pack('>H', len(body) - 16) + body[:16]
|
|
|
|
hdrbox = self.box.encrypt(header, self.nonce)[24:]
|
|
self.writer.write(hdrbox)
|
|
|
|
self.nonce = inc_nonce(inc_nonce(self.nonce))
|
|
self.writer.write(body[16:])
|
|
|
|
def close(self):
|
|
self.writer.write(self.box.encrypt(b'\x00' * 18, self.nonce)[24:])
|