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