from twistednoise import genkeypair
from noise.connection import NoiseConnection, Keypair

server_key_pair = genkeypair()
client_key_pair = genkeypair()

init = NoiseConnection.from_name(b'Noise_XK_448_ChaChaPoly_SHA256')
init.set_as_initiator()

init.set_keypair_from_private_bytes(Keypair.STATIC, client_key_pair[1])
init.set_keypair_from_public_bytes(Keypair.REMOTE_STATIC, server_key_pair[0])

resp = NoiseConnection.from_name(b'Noise_XK_448_ChaChaPoly_SHA256')
resp.set_as_responder()

resp.set_keypair_from_private_bytes(Keypair.STATIC, server_key_pair[1])

init.start_handshake()
resp.start_handshake()

messages = []

cursend, currecv = init, resp
while not init.handshake_finished and not resp.handshake_finished:
	# send a message
	msg = cursend.write_message()
	currecv.read_message(msg)

	# record it
	messages.append(msg)

	# swap roles
	cursend, currecv = currecv, cursend

print(list(map(len, messages)))