parent
30d2ec4209
commit
3831485ff3
@ -1,6 +1,8 @@
|
||||
import struct
|
||||
from asyncio import IncompleteReadError
|
||||
|
||||
from async_generator import async_generator, yield_
|
||||
|
||||
from nacl.secret import SecretBox
|
||||
|
||||
from .util import split_chunks, inc_nonce
|
||||
@ -58,12 +60,13 @@ class UnboxStream(object):
|
||||
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
|
||||
yield data
|
||||
await yield_(data)
|
||||
|
||||
|
||||
class BoxStream(object):
|
||||
|
@ -21,6 +21,8 @@
|
||||
|
||||
from asyncio import open_connection, start_server, ensure_future
|
||||
|
||||
from async_generator import async_generator, yield_
|
||||
|
||||
from .boxstream import get_stream_pair
|
||||
from .crypto import SHSClientCrypto, SHSServerCrypto
|
||||
|
||||
@ -43,9 +45,10 @@ class SHSSocket(object):
|
||||
def disconnect(self):
|
||||
self.writer.close()
|
||||
|
||||
@async_generator
|
||||
async def __aiter__(self):
|
||||
async for msg in self.read_stream:
|
||||
yield msg
|
||||
await yield_(msg)
|
||||
|
||||
def on_connect(self, cb):
|
||||
self._on_connect = cb
|
||||
|
@ -24,6 +24,7 @@ from io import BytesIO
|
||||
|
||||
from .test_crypto import (CLIENT_ENCRYPT_KEY, CLIENT_ENCRYPT_NONCE)
|
||||
from secret_handshake.boxstream import BoxStream, UnboxStream, HEADER_LENGTH
|
||||
from secret_handshake.util import async_comprehend
|
||||
|
||||
MESSAGE_1 = (b'\xcev\xedE\x06l\x02\x13\xc8\x17V\xfa\x8bZ?\x88B%O\xb0L\x9f\x8e\x8c0y\x1dv\xc0\xc9\xf6\x9d\xc2\xdf\xdb'
|
||||
b'\xee\x9d')
|
||||
@ -71,7 +72,7 @@ async def test_unboxstream():
|
||||
|
||||
unbox_stream = UnboxStream(buffer, CLIENT_ENCRYPT_KEY, CLIENT_ENCRYPT_NONCE)
|
||||
assert not unbox_stream.closed
|
||||
assert [msg async for msg in unbox_stream] == [b'foo', b'foo', b'bar']
|
||||
assert (await async_comprehend(unbox_stream)) == [b'foo', b'foo', b'bar']
|
||||
assert unbox_stream.closed
|
||||
|
||||
|
||||
|
@ -4,6 +4,14 @@ NONCE_SIZE = 24
|
||||
MAX_NONCE = (8 * NONCE_SIZE)
|
||||
|
||||
|
||||
async def async_comprehend(generator):
|
||||
"""Emulate ``[elem async for elem in generator]``."""
|
||||
results = []
|
||||
async for msg in generator:
|
||||
results.append(msg)
|
||||
return results
|
||||
|
||||
|
||||
def inc_nonce(nonce):
|
||||
num = bytes_to_long(nonce) + 1
|
||||
if num > 2 ** MAX_NONCE:
|
||||
|
Loading…
Reference in New Issue
Block a user