gergely/PySecretHandshake
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ci: Add mypy as a dev dependency and configure it

Browse Source
This commit is contained in:
Gergely Polonkai
2023-10-30 05:47:28 +01:00
parent 95039914ba
commit 4996931b54
14 changed files with 350 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
secret_handshake/boxstream.py
View File

@@ -1,8 +1,8 @@
"""Box stream utilities"""
from asyncio import IncompleteReadError
from asyncio import IncompleteReadError, StreamReader, StreamWriter
import struct
from typing import Tuple
from typing import Any, AsyncIterator, Optional, Tuple, TypedDict
from nacl.secret import SecretBox
@@ -13,36 +13,51 @@ MAX_SEGMENT_SIZE = 4 * 1024
TERMINATION_HEADER = b"\x00" * 18
def get_stream_pair(reader, writer, **kwargs) -> Tuple["UnboxStream", "BoxStream"]:
class BoxStreamKeys(TypedDict):
"""Dictionary to hold all box stream keys"""
decrypt_key: bytes
decrypt_nonce: bytes
encrypt_key: bytes
encrypt_nonce: bytes
shared_secret: bytes
def get_stream_pair( # pylint: disable=too-many-arguments
reader: StreamReader, # pylint: disable=unused-argument
writer: StreamWriter,
*,
decrypt_key: bytes,
decrypt_nonce: bytes,
encrypt_key: bytes,
encrypt_nonce: bytes,
**kwargs: Any,
) -> Tuple["UnboxStream", "BoxStream"]:
"""Create a new duplex box stream"""
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)
read_stream = UnboxStream(reader, key=decrypt_key, nonce=decrypt_nonce)
write_stream = BoxStream(writer, key=encrypt_key, nonce=encrypt_nonce)
return read_stream, write_stream
class UnboxStream:
"""Unboxing stream"""
def __init__(self, reader, key, nonce):
def __init__(self, reader: StreamReader, key: bytes, nonce: bytes):
self.reader = reader
self.key = key
self.nonce = nonce
self.closed = False
async def read(self):
async def read(self) -> Optional[bytes]:
"""Read data from the stream"""
try:
data = await self.reader.readexactly(HEADER_LENGTH)
except IncompleteReadError:
self.closed = True
return None
box = SecretBox(self.key)
@@ -51,6 +66,7 @@ class UnboxStream:
if header == TERMINATION_HEADER:
self.closed = True
return None
length = struct.unpack(">H", header[:2])[0]
@@ -61,12 +77,13 @@ class UnboxStream:
body = box.decrypt(mac + data, inc_nonce(self.nonce))
self.nonce = inc_nonce(inc_nonce(self.nonce))
return body
def __aiter__(self):
def __aiter__(self) -> AsyncIterator[bytes]:
return self
async def __anext__(self):
async def __anext__(self) -> bytes:
data = await self.read()
if data is None:
@@ -78,17 +95,17 @@ class UnboxStream:
class BoxStream:
"""Box stream"""
def __init__(self, writer, key, nonce):
def __init__(self, writer: StreamWriter, key: bytes, nonce: bytes):
self.writer = writer
self.key = key
self.box = SecretBox(self.key)
self.nonce = nonce
def write(self, data):
def write(self, data: bytes) -> None:
"""Write data to the box stream"""
for chunk in split_chunks(data, MAX_SEGMENT_SIZE):
body = self.box.encrypt(chunk, inc_nonce(self.nonce))[24:]
body = self.box.encrypt(bytes(chunk), inc_nonce(self.nonce))[24:]
header = struct.pack(">H", len(body) - 16) + body[:16]
hdrbox = self.box.encrypt(header, self.nonce)[24:]
@@ -97,7 +114,7 @@ class BoxStream:
self.nonce = inc_nonce(inc_nonce(self.nonce))
self.writer.write(body[16:])
def close(self):
def close(self) -> None:
"""Close the box stream"""
self.writer.write(self.box.encrypt(b"\x00" * 18, self.nonce)[24:])

108
secret_handshake/crypto.py
View File

@@ -28,7 +28,9 @@ from typing import Optional
from nacl.bindings import crypto_box_afternm, crypto_box_open_afternm, crypto_scalarmult
from nacl.exceptions import CryptoError
from nacl.public import PrivateKey
from nacl.signing import VerifyKey
from nacl.signing import SigningKey, VerifyKey
from .boxstream import BoxStreamKeys
APPLICATION_KEY = b64decode("1KHLiKZvAvjbY1ziZEHMXawbCEIM6qwjCDm3VYRan/s=")
@@ -40,29 +42,38 @@ class SHSError(Exception):
class SHSCryptoBase: # pylint: disable=too-many-instance-attributes
"""Base functions for SHS cryptography"""
def __init__(self, local_key, ephemeral_key=None, application_key=None):
self.local_key = local_key
def __init__(
self,
local_key: SigningKey,
ephemeral_key: Optional[PrivateKey] = None,
application_key: Optional[bytes] = None,
):
self.application_key = application_key or APPLICATION_KEY
self.shared_hash = None
self.remote_ephemeral_key = None
self.shared_secret = None
self.remote_app_hmac = None
self.remote_pub_key = None
self.box_secret = None
self.b_alice: Optional[bytes] = None
self.box_secret: Optional[bytes] = None
self.hello: Optional[bytes] = None
self.local_key = local_key
self.remote_app_hmac: Optional[bytes] = None
self.remote_ephemeral_key: Optional[bytes] = None
self.remote_pub_key: Optional[VerifyKey] = None
self.shared_hash: Optional[bytes] = None
self.shared_secret: Optional[bytes] = None
self._reset_keys(ephemeral_key or PrivateKey.generate())
def _reset_keys(self, ephemeral_key):
def _reset_keys(self, ephemeral_key: PrivateKey) -> None:
self.local_ephemeral_key = ephemeral_key
self.local_app_hmac = hmac.new(
self.application_key, bytes(ephemeral_key.public_key), digestmod="sha512"
).digest()[:32]
def generate_challenge(self):
def generate_challenge(self) -> bytes:
"""Generate and return a challenge to be sent to the server."""
return self.local_app_hmac + bytes(self.local_ephemeral_key.public_key)
def verify_challenge(self, data):
def verify_challenge(self, data: bytes) -> bool:
"""Verify the correctness of challenge sent from the client."""
assert len(data) == 64
sent_hmac, remote_ephemeral_key = data[:32], data[32:]
@@ -76,9 +87,10 @@ class SHSCryptoBase: # pylint: disable=too-many-instance-attributes
self.remote_ephemeral_key = remote_ephemeral_key
# this is hash(a * b)
self.shared_hash = hashlib.sha256(self.shared_secret).digest()
return ok
def clean(self, new_ephemeral_key=None):
def clean(self, new_ephemeral_key: Optional[PrivateKey] = None) -> None:
"""Clean internal data"""
self._reset_keys(new_ephemeral_key or PrivateKey.generate())
@@ -86,10 +98,15 @@ class SHSCryptoBase: # pylint: disable=too-many-instance-attributes
self.shared_hash = None
self.remote_ephemeral_key = None
def get_box_keys(self):
def get_box_keys(self) -> BoxStreamKeys:
"""Get the box stream’s keys"""
assert self.box_secret
assert self.remote_app_hmac
assert self.remote_pub_key
shared_secret = hashlib.sha256(self.box_secret).digest()
return {
"shared_secret": shared_secret,
"encrypt_key": hashlib.sha256(shared_secret + bytes(self.remote_pub_key)).digest(),
@@ -102,17 +119,15 @@ class SHSCryptoBase: # pylint: disable=too-many-instance-attributes
class SHSServerCrypto(SHSCryptoBase):
"""SHS server crypto algorithm"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.b_alice = None
self.hello = None
self.box_secret = None
self.remote_pub_key = None
def verify_client_auth(self, data):
def verify_client_auth(self, data: bytes) -> bool:
"""Verify client authentication data"""
assert self.remote_ephemeral_key
assert self.shared_hash
assert self.shared_secret
assert len(data) == 112
a_bob = crypto_scalarmult(bytes(self.local_key.to_curve25519_private_key()), self.remote_ephemeral_key)
box_secret = hashlib.sha256(self.application_key + self.shared_secret + a_bob).digest()
self.hello = crypto_box_open_afternm(data, b"\x00" * 24, box_secret)
@@ -127,15 +142,21 @@ class SHSServerCrypto(SHSCryptoBase):
bytes(self.local_ephemeral_key), bytes(self.remote_pub_key.to_curve25519_public_key())
)
self.box_secret = hashlib.sha256(self.application_key + self.shared_secret + a_bob + b_alice).digest()[:32]
return True
def generate_accept(self):
def generate_accept(self) -> bytes:
"""Generate an accept message"""
assert self.box_secret
assert self.hello
assert self.shared_hash
okay = self.local_key.sign(self.application_key + self.hello + self.shared_hash).signature
return crypto_box_afternm(okay, b"\x00" * 24, self.box_secret)
def clean(self, new_ephemeral_key=None):
def clean(self, new_ephemeral_key: Optional[PrivateKey] = None) -> None:
super().clean(new_ephemeral_key=new_ephemeral_key)
self.hello = None
self.b_alice = None
@@ -152,21 +173,26 @@ class SHSClientCrypto(SHSCryptoBase):
def __init__(
self,
local_key: PrivateKey,
local_key: SigningKey,
server_pub_key: bytes,
ephemeral_key: PrivateKey,
application_key: Optional[bytes] = None,
):
super().__init__(local_key, ephemeral_key, application_key)
self.remote_pub_key = VerifyKey(server_pub_key)
self.b_alice = None
self.a_bob = None
self.hello = None
self.box_secret = None
def verify_server_challenge(self, data):
self.a_bob: Optional[bytes] = None
self.remote_pub_key = VerifyKey(server_pub_key)
def verify_server_challenge(self, data: bytes) -> bool:
"""Verify the correctness of challenge sent from the server."""
assert self.remote_pub_key
assert super().verify_challenge(data)
assert self.shared_hash
assert self.shared_secret
curve_pkey = self.remote_pub_key.to_curve25519_public_key()
# a_bob is (a * B)
@@ -179,17 +205,30 @@ class SHSClientCrypto(SHSCryptoBase):
signed_message = self.local_key.sign(self.application_key + bytes(self.remote_pub_key) + self.shared_hash)
message_to_box = signed_message.signature + bytes(self.local_key.verify_key)
self.hello = message_to_box
return True
def generate_client_auth(self):
def generate_client_auth(self) -> bytes:
"""Generate box[K|a*b|a*B](H)"""
assert self.box_secret
assert self.hello
nonce = b"\x00" * 24
# return box(K | a * b | a * B)[H]
return crypto_box_afternm(self.hello, nonce, self.box_secret)
def verify_server_accept(self, data):
def verify_server_accept(self, data: bytes) -> bool:
"""Verify that the server's accept message is sane"""
assert self.a_bob
assert self.hello
assert self.remote_ephemeral_key
assert self.remote_pub_key
assert self.shared_hash
assert self.shared_secret
curve_lkey = self.local_key.to_curve25519_private_key()
# b_alice is (A * b)
b_alice = crypto_scalarmult(bytes(curve_lkey), self.remote_ephemeral_key)
@@ -208,9 +247,10 @@ class SHSClientCrypto(SHSCryptoBase):
# we should have received sign(B)[K | H | hash(a * b)]
# let's see if that signature can verify the reconstructed data on our side
self.remote_pub_key.verify(self.application_key + self.hello + self.shared_hash, signature)
return True
def clean(self, new_ephemeral_key=None):
def clean(self, new_ephemeral_key: Optional[PrivateKey] = None) -> None:
super().clean(new_ephemeral_key=new_ephemeral_key)
self.a_bob = None
self.b_alice = None

111
secret_handshake/network.py
View File

@@ -20,10 +20,15 @@
"""Networking functionality"""
import asyncio
from asyncio import StreamReader, StreamWriter, ensure_future, open_connection, start_server
from typing import AsyncIterator, Awaitable, Callable, List, Optional
from .boxstream import get_stream_pair
from .crypto import SHSClientCrypto, SHSServerCrypto
from nacl.public import PrivateKey
from nacl.signing import SigningKey
from typing_extensions import Self
from .boxstream import BoxStream, UnboxStream, get_stream_pair
from .crypto import SHSClientCrypto, SHSCryptoBase, SHSServerCrypto
class SHSClientException(Exception):
@@ -33,32 +38,37 @@ class SHSClientException(Exception):
class SHSDuplexStream:
"""SHS duplex stream"""
def __init__(self):
self.write_stream = None
self.read_stream = None
def __init__(self) -> None:
self.write_stream: Optional[BoxStream] = None
self.read_stream: Optional[UnboxStream] = None
self.is_connected = False
def write(self, data):
def write(self, data: bytes) -> None:
"""Write data to the write stream"""
assert self.write_stream
self.write_stream.write(data)
async def read(self):
async def read(self) -> Optional[bytes]:
"""Read data from the read stream"""
assert self.read_stream
return await self.read_stream.read()
def close(self):
def close(self) -> None:
"""Close the duplex stream"""
self.write_stream.close()
self.read_stream.close()
if self.write_stream:
self.write_stream.close()
self.is_connected = False
def __aiter__(self):
def __aiter__(self) -> AsyncIterator[bytes]:
return self
async def __anext__(self):
async def __anext__(self) -> bytes:
msg = await self.read()
if msg is None:
@@ -70,45 +80,53 @@ class SHSDuplexStream:
class SHSEndpoint:
"""SHS endpoint"""
def __init__(self):
self._on_connect = None
self.crypto = None
def __init__(self) -> None:
self._on_connect: Optional[Callable[[SHSDuplexStream], Awaitable[None]]] = None
self.crypto: Optional[SHSCryptoBase] = None
def on_connect(self, cb):
def on_connect(self, cb: Callable[[SHSDuplexStream], Awaitable[None]]) -> None:
"""Set the function to be called when a new connection arrives"""
self._on_connect = cb
def disconnect(self):
def disconnect(self) -> None:
"""Disconnect the endpoint"""
raise NotImplementedError
class SHSServer(SHSEndpoint):
"""SHS server"""
def __init__(self, host, port, server_kp, application_key=None):
def __init__(self, host: str, port: int, server_kp: SigningKey, application_key: Optional[bytes] = None):
super().__init__()
self.host = host
self.port = port
self.crypto = SHSServerCrypto(server_kp, application_key=application_key)
self.connections = []
self.crypto: SHSServerCrypto = SHSServerCrypto(server_kp, application_key=application_key)
self.connections: List[SHSServerConnection] = []
async def _handshake(self, reader: StreamReader, writer: StreamWriter) -> None:
assert self.crypto
async def _handshake(self, reader, writer):
data = await reader.readexactly(64)
if not self.crypto.verify_challenge(data):
raise SHSClientException("Client challenge is not valid")
writer.write(self.crypto.generate_challenge())
data = await reader.readexactly(112)
if not self.crypto.verify_client_auth(data):
raise SHSClientException("Client auth is not valid")
writer.write(self.crypto.generate_accept())
async def handle_connection(self, reader, writer):
async def handle_connection(self, reader: StreamReader, writer: StreamWriter) -> None:
"""Handle incoming connections"""
assert self.crypto
self.crypto.clean()
await self._handshake(reader, writer)
keys = self.crypto.get_box_keys()
@@ -118,14 +136,14 @@ class SHSServer(SHSEndpoint):
self.connections.append(conn)
if self._on_connect:
asyncio.ensure_future(self._on_connect(conn))
ensure_future(self._on_connect(conn))
async def listen(self):
async def listen(self) -> None:
"""Listen for connections"""
await asyncio.start_server(self.handle_connection, self.host, self.port)
await start_server(self.handle_connection, self.host, self.port)
def disconnect(self):
def disconnect(self) -> None:
for connection in self.connections:
connection.close()
@@ -133,52 +151,63 @@ class SHSServer(SHSEndpoint):
class SHSServerConnection(SHSDuplexStream):
"""SHS server connection"""
def __init__(self, read_stream, write_stream):
def __init__(self, read_stream: UnboxStream, write_stream: BoxStream):
super().__init__()
self.read_stream = read_stream
self.write_stream = write_stream
@classmethod
def from_byte_streams(cls, reader, writer, **keys):
def from_byte_streams(cls, reader: StreamReader, writer: StreamWriter, **keys: bytes) -> Self:
"""Create a server connection from an existing byte stream"""
reader, writer = get_stream_pair(reader, writer, **keys)
box_reader, box_writer = get_stream_pair(reader, writer, **keys)
return cls(reader, writer)
return cls(box_reader, box_writer)
class SHSClient(SHSDuplexStream, SHSEndpoint):
"""SHS client"""
def __init__( # pylint: disable=too-many-arguments
self, host, port, client_kp, server_pub_key, ephemeral_key=None, application_key=None
self,
host: str,
port: int,
client_kp: SigningKey,
server_pub_key: bytes,
ephemeral_key: Optional[PrivateKey] = None,
application_key: Optional[bytes] = None,
):
SHSDuplexStream.__init__(self)
SHSEndpoint.__init__(self)
self.host = host
self.port = port
self.writer = None
self.crypto = SHSClientCrypto(
client_kp, server_pub_key, ephemeral_key=ephemeral_key, application_key=application_key
self.writer: Optional[StreamWriter] = None
self.crypto: SHSClientCrypto = SHSClientCrypto(
client_kp,
server_pub_key,
ephemeral_key=ephemeral_key or PrivateKey.generate(),
application_key=application_key,
)
async def _handshake(self, reader, writer):
async def _handshake(self, reader: StreamReader, writer: StreamWriter) -> None:
writer.write(self.crypto.generate_challenge())
data = await reader.readexactly(64)
if not self.crypto.verify_server_challenge(data):
raise SHSClientException("Server challenge is not valid")
writer.write(self.crypto.generate_client_auth())
data = await reader.readexactly(80)
if not self.crypto.verify_server_accept(data):
raise SHSClientException("Server accept is not valid")
async def open(self):
async def open(self) -> None:
"""Open the TCP connection"""
reader, writer = await asyncio.open_connection(self.host, self.port)
reader, writer = await open_connection(self.host, self.port)
await self._handshake(reader, writer)
keys = self.crypto.get_box_keys()
@@ -189,7 +218,7 @@ class SHSClient(SHSDuplexStream, SHSEndpoint):
self.is_connected = True
if self._on_connect:
await self._on_connect()
await self._on_connect(self)
def disconnect(self):
def disconnect(self) -> None:
self.close()

0
secret_handshake/py.typed Normal file
View File

39
secret_handshake/util.py
View File

@@ -21,12 +21,14 @@
"""Utility functions"""
import struct
from typing import Generator, Sequence, TypeVar
NONCE_SIZE = 24
MAX_NONCE = 8 * NONCE_SIZE
T = TypeVar("T")
def inc_nonce(nonce):
def inc_nonce(nonce: bytes) -> bytes:
"""Increment nonce"""
num = bytes_to_long(nonce) + 1
@@ -40,34 +42,38 @@ def inc_nonce(nonce):
return bnum
def split_chunks(seq, n):
def split_chunks(seq: Sequence[T], n: int) -> Generator[Sequence[T], None, None]:
"""Split sequence in equal-sized chunks.
The last chunk is not padded."""
while seq:
yield seq[:n]
seq = seq[n:]
# Stolen from PyCypto (Public Domain)
def b(s):
def b(s: str) -> bytes:
"""Shorthand for s.encode("latin-1")"""
return s.encode("latin-1") # utf-8 would cause some side-effects we don't want
def long_to_bytes(n, blocksize=0):
"""long_to_bytes(n:long, blocksize:int) : string
Convert a long integer to a byte string.
If optional blocksize is given and greater than zero, pad the front of the
byte string with binary zeros so that the length is a multiple of
blocksize.
def long_to_bytes(n: int, blocksize: int = 0) -> bytes:
"""Convert a long integer to a byte string.
If optional ``blocksize`` is given and greater than zero, pad the front of the byte string with binary zeros so
that the length is a multiple of blocksize.
"""
# after much testing, this algorithm was deemed to be the fastest
s = b("")
pack = struct.pack
while n > 0:
s = pack(">I", n & 0xFFFFFFFF) + s
n = n >> 32
# strip off leading zeros
for i, c in enumerate(s):
if c != b("\000")[0]:
@@ -76,26 +82,33 @@ def long_to_bytes(n, blocksize=0):
# only happens when n == 0
s = b("\000")
i = 0
s = s[i:]
# add back some pad bytes. this could be done more efficiently w.r.t. the
# de-padding being done above, but sigh...
if blocksize > 0 and len(s) % blocksize:
s = (blocksize - len(s) % blocksize) * b("\000") + s
return s
def bytes_to_long(s):
"""bytes_to_long(string) : long
Convert a byte string to a long integer.
This is (essentially) the inverse of long_to_bytes().
def bytes_to_long(s: bytes) -> int:
"""Convert a byte string to a long integer.
This is (essentially) the inverse of ``long_to_bytes()``.
"""
acc = 0
unpack = struct.unpack
length = len(s)
if length % 4:
extra = 4 - length % 4
s = b("\000") * extra + s
length = length + extra
for i in range(0, length, 4):
acc = (acc << 32) + unpack(">I", s[i : i + 4])[0]
return acc