Upgraded transmat protocol to ver. 0.2
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user