Upgraded transmat protocol to ver. 0.2

This commit is contained in:
SG
2023-05-25 15:24:35 +02:00
parent 2074c001c7
commit c88224009c

View File

@@ -9,16 +9,24 @@ from progress.bar import Bar
import diceware import diceware
VERSION_NUMBER="0.1.8" VERSION_NUMBER="0.2.0"
PROTOCOL_VERSION="0.2"
SCRIPT_BASENAME = os.path.basename(sys.argv[0]) SCRIPT_BASENAME = os.path.basename(sys.argv[0])
async def read_chunks(filename, chunk_size = 1024): #async def read_chunks(filename, chunk_size = 1024):
with open(filename, "rb") as f: # with open(filename, "rb") as f:
while True: # while True:
chunk = f.read(chunk_size) # chunk = f.read(chunk_size)
if not chunk: # if not chunk:
break # break
yield chunk # yield chunk
def chunk_from_file(file, chunk_id, chunk_size):
position = chunk_id * chunk_size
file.seek(position)
chunk = file.read(chunk_size)
return chunk
def derive_key_from_password(password): def derive_key_from_password(password):
@@ -68,9 +76,7 @@ async def send_msg(ws, msg):
async def send_encrypted_msg(ws, k, data): async def send_encrypted_msg(ws, k, data):
( (
msgtype,
peer_group_id, peer_group_id,
role,
filename, filename,
chunk_size, chunk_size,
chunk_id, chunk_id,
@@ -86,9 +92,7 @@ async def send_encrypted_msg(ws, k, data):
"chunk": chunk "chunk": chunk
} }
msg = { msg = {
"msgtype": msgtype,
"peer_group_id": peer_group_id, "peer_group_id": peer_group_id,
"role": role,
"payload": encrypt_chunk(k, pickle.dumps(payload)) "payload": encrypt_chunk(k, pickle.dumps(payload))
} }
await send_msg(ws, json.dumps(msg)) await send_msg(ws, json.dumps(msg))
@@ -125,7 +129,7 @@ async def main():
WS_RELAY_SERVER = args.server WS_RELAY_SERVER = args.server
server_part = f'--server {WS_RELAY_SERVER}' server_part = f'--server {WS_RELAY_SERVER}'
if args.version: if args.version:
print(f"{SCRIPT_BASENAME} ver {VERSION_NUMBER}") print(f"{SCRIPT_BASENAME} ver. {VERSION_NUMBER}, protocol ver. {PROTOCOL_VERSION}")
sys.exit(0) sys.exit(0)
k = derive_key_from_password(password) k = derive_key_from_password(password)
@@ -135,41 +139,57 @@ async def main():
print('Run the following command on the remote party:\n\n', send_part, passwd_part, server_part, "\n") print('Run the following command on the remote party:\n\n', send_part, passwd_part, server_part, "\n")
filename = os.path.basename(file_path) filename = os.path.basename(file_path)
file_size = os.path.getsize(file_path) file_size = os.path.getsize(file_path)
chunk_id = 0 file = open(file_path, 'rb')
chunk_size = 1024 * 512 chunk_size = 1024 * 512
chunk_id = ""
chunk = ""
number_of_chunks = math.ceil(file_size / chunk_size) number_of_chunks = math.ceil(file_size / chunk_size)
WS_RELAY_SERVER = WS_RELAY_SERVER.replace('http', 'ws', 1) WS_RELAY_SERVER = WS_RELAY_SERVER.replace('http', 'ws', 1)
async with websockets.connect(WS_RELAY_SERVER) as ws: async with websockets.connect(WS_RELAY_SERVER) as ws:
msgtype = "announce" # Announce ourself to the relay
await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, filename, "", "", number_of_chunks, "")) await send_encrypted_msg(ws, k, (peer_group_id, filename, chunk_size, chunk_id, number_of_chunks, chunk))
bar = Bar(filename, max=number_of_chunks, check_tty=False, suffix='%(percent).1f%% complete - %(eta_td)s remaining')
while True: while True:
# Wait for messages from receiver
message = await ws.recv() message = await ws.recv()
message = json.loads(message) message = json.loads(message)
if message["msgtype"] == "announce" and message["peer_group_id"] == peer_group_id: # If message is not meant for us - skip it
break if 'peer_group_id' not in message or message["peer_group_id"] != peer_group_id:
bar = Bar(filename, max=number_of_chunks, check_tty=False, suffix='%(percent).1f%% complete - %(eta_td)s remaining') pass
msgtype = "data" payload = pickle.loads(decrypt_chunk(k, message["payload"]))
async for chunk in read_chunks(file_path, chunk_size): chunk_id = payload["chunk_id"]
msg = (msgtype, peer_group_id, role, filename, chunk_size, chunk_id, number_of_chunks, chunk) chunk = chunk_from_file(file, chunk_id, chunk_size)
#bar = Bar(filename, max=number_of_chunks, check_tty=False, suffix='%(percent).1f%% complete - %(eta_td)s remaining')
msg = (peer_group_id, filename, chunk_size, chunk_id, number_of_chunks, chunk)
await send_encrypted_msg(ws, k, msg) await send_encrypted_msg(ws, k, msg)
#bar.index = chunk_id
bar.next()
chunk_id += 1 chunk_id += 1
if chunk_id > number_of_chunks: if chunk_id > number_of_chunks:
break break
proceed = await ws.recv()
bar.next()
bar.suffix = '%(percent).1f%% complete' bar.suffix = '%(percent).1f%% complete'
bar.update() bar.update()
print("") print("")
if role =='receive': if role =='receive':
async with websockets.connect(WS_RELAY_SERVER) as ws: async with websockets.connect(WS_RELAY_SERVER) as ws:
msgtype = "announce" filename = ""
await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", "")) chunk_size = ""
chunk_id = 0
number_of_chunks = ""
chunk = ""
bar = None bar = None
f = None f = None
while True: while True:
# Announce ourself to the relay
await send_encrypted_msg(ws, k, (peer_group_id, filename, chunk_size, chunk_id, number_of_chunks, chunk))
# Wait for message from sender
message = await ws.recv() message = await ws.recv()
message = json.loads(message) message = json.loads(message)
# If message is not meant for us - skip it
if 'peer_group_id' not in message or message["peer_group_id"] != peer_group_id:
pass
payload = message["payload"] payload = message["payload"]
msg = pickle.loads(decrypt_chunk(k, payload)) msg = pickle.loads(decrypt_chunk(k, payload))
chunk_id = msg["chunk_id"] chunk_id = msg["chunk_id"]
@@ -182,13 +202,12 @@ async def main():
f.write(msg["chunk"]) f.write(msg["chunk"])
else: else:
f.write(msg["chunk"]) f.write(msg["chunk"])
msgtype = "proceed" chunk_id += 1
await send_encrypted_msg(ws, k, (msgtype, peer_group_id, role, "", "", "", "", "")) # request the next chunk from sender
bar.suffix = "%(percent).1f%% complete - %(eta_td)s remaining" bar.suffix = "%(percent).1f%% complete - %(eta_td)s remaining"
bar.next() bar.next()
if chunk_id >= number_of_chunks:
if chunk_id >= number_of_chunks-1:
# This was the last chunk, exit # This was the last chunk, exit
await send_encrypted_msg(ws, k, (peer_group_id, "", "", chunk_id, "", ""))
f.close() f.close()
bar.suffix = '%(percent).1f%% complete' bar.suffix = '%(percent).1f%% complete'
bar.update() bar.update()