78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
import struct
|
|
from nacl.secret import SecretBox
|
|
|
|
from .util import bytes_to_long, long_to_bytes
|
|
|
|
NONCE_SIZE = 24
|
|
HEADER_LENGTH = 2 + 16 + 16
|
|
MAX_SEGMENT_SIZE = 4 * 1024
|
|
TERMINATION_HEADER = (b'\x00' * HEADER_LENGTH)
|
|
MAX_NONCE = (8 * NONCE_SIZE)
|
|
|
|
|
|
def inc_nonce(nonce):
|
|
num = bytes_to_long(nonce) + 1
|
|
if num > 2**MAX_NONCE:
|
|
num = 0
|
|
bnum = long_to_bytes(num)
|
|
bnum = b'\x00' * (NONCE_SIZE - len(bnum)) + bnum
|
|
return bnum
|
|
|
|
|
|
def get_stream_pair(reader, writer, **kwargs):
|
|
shared_secret = kwargs.pop('shared_secret')
|
|
return UnboxStream(reader, shared_secret, **kwargs), BoxStream(writer, shared_secret, **kwargs)
|
|
|
|
|
|
class UnboxStream(object):
|
|
def __init__(self, reader, shared_secret, **key_data):
|
|
self.reader = reader
|
|
self.decrypt_key = key_data.get('decrypt_key')
|
|
self.decrypt_nonce = key_data.get('decrypt_nonce')
|
|
|
|
async def process(self):
|
|
while True:
|
|
data = await self.reader.read(HEADER_LENGTH)
|
|
if not data:
|
|
break
|
|
|
|
box = SecretBox(self.decrypt_key)
|
|
header = box.decrypt(data, self.decrypt_nonce)
|
|
|
|
if header == TERMINATION_HEADER:
|
|
return
|
|
|
|
length = struct.unpack('>H', header[:2])[0]
|
|
mac = header[2:]
|
|
|
|
data = await self.reader.read(length)
|
|
|
|
self.decrypt_nonce = inc_nonce(self.decrypt_nonce)
|
|
body = box.decrypt(mac + data, self.decrypt_nonce)
|
|
|
|
self.decrypt_nonce = inc_nonce(self.decrypt_nonce)
|
|
yield body
|
|
print('Disconnect')
|
|
|
|
|
|
class BoxStream(object):
|
|
def __init__(self, writer, shared_secret, **key_data):
|
|
self.writer = writer
|
|
self.encrypt_key = key_data.get('decrypt_key')
|
|
self.encrypt_nonce = key_data.get('decrypt_nonce')
|
|
|
|
async def write(self, data):
|
|
box = SecretBox(self.encrypt_key)
|
|
|
|
# XXX: This nonce logic is almost for sure wrong
|
|
|
|
self.encrypt_nonce = inc_nonce(self.encrypt_nonce)
|
|
|
|
body = box.encrypt(data, self.encrypt_nonce)
|
|
header = struct.pack('>H', len(body)) + body[:16]
|
|
|
|
self.writer.write(box.encrypt(header, self.encrypt_nonce))
|
|
|
|
self.encrypt_nonce = inc_nonce(self.encrypt_nonce)
|
|
self.writer.write(body)
|