#!/usr/bin/env python import asyncio, websockets import sys, os, base64, argparse, json, pickle, math from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from progress.bar import Bar import diceware 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 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): salt_part = password.split('-')[0] key_part = "-".join(password.split('-')[1:]) digest = hashes.Hash(hashes.SHA256()) digest.update(salt_part.encode('utf-8')) salt = digest.finalize() kdf_key_length = 32 kdf_iterations = 1000000 kdf = PBKDF2HMAC( algorithm = hashes.SHA256(), length = kdf_key_length, salt = salt, iterations = kdf_iterations ) key = kdf.derive(key_part.encode('utf-8')) return base64.urlsafe_b64encode(key) def get_peer_group_id(password): peer_group_id_part = password.split('-')[0] digest = hashes.Hash(hashes.SHA256()) digest.update(peer_group_id_part.encode('utf-8')) peer_group_id = digest.finalize() return peer_group_id.hex() def encrypt_chunk(key, chunk): f = Fernet(key) encrypted_chunk = f.encrypt(chunk).decode('utf-8') return encrypted_chunk def decrypt_chunk(key, chunk): f = Fernet(key) return f.decrypt(chunk.encode('utf-8')) async def send_message(websocket, msg): await websocket.send(msg) async def send_msg(ws, msg): await ws.send(msg) async def send_encrypted_msg(ws, k, data): ( peer_group_id, filename, chunk_size, chunk_id, number_of_chunks, chunk ) = data payload = { "filename": filename, "chunk_size": chunk_size, "chunk_id": chunk_id, "number_of_chunks": number_of_chunks, "chunk": chunk } msg = { "peer_group_id": peer_group_id, "payload": encrypt_chunk(k, pickle.dumps(payload)) } await send_msg(ws, json.dumps(msg)) async def main(): WS_RELAY_SERVER = "wss://transmat.exocortex.ru/ws" parser = argparse.ArgumentParser() arg_group = parser.add_mutually_exclusive_group(required=True) arg_group.add_argument('--receive', '--recv', action='store_true', help='Receive a file from the remote party (mutually exclusive with --send and --relay)') arg_group.add_argument('--send', type=str, help='Send a file to the remote party (mutually exclusive with --receive and --relay)') arg_group.add_argument('--version', action='store_true', help="Show version") parser.add_argument('--server', type=str, help="Specify the Relay server URL (ignored with --relay)") parser.add_argument('--password', type=str, help="Specify the shared password (for --receive)") args = parser.parse_args() passwd_part = server_part = "" send_part = f'{SCRIPT_BASENAME} --receive' if args.receive: role = 'receive' password = args.password if args.send and args.password is None: password = diceware.generate_passphrase(4) passwd_part = f"--password {password}" if args.receive and args.password is None: print("Error: --password required when receiving files.") sys.exit(1) if args.password: password = args.password passwd_part = f"--password {password}" if args.send: role = 'send' file_path = args.send if args.server: WS_RELAY_SERVER = args.server server_part = f'--server {WS_RELAY_SERVER}' if args.version: print(f"{SCRIPT_BASENAME} ver. {VERSION_NUMBER}, protocol ver. {PROTOCOL_VERSION}") sys.exit(0) k = derive_key_from_password(password) peer_group_id = get_peer_group_id(password) if role == 'send': 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) 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: # 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 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 bar.suffix = '%(percent).1f%% complete' bar.update() print("") if role =='receive': async with websockets.connect(WS_RELAY_SERVER) as ws: 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') if f is None: f = open(filename, "wb") f.write(msg["chunk"]) else: f.write(msg["chunk"]) chunk_id += 1 bar.suffix = "%(percent).1f%% complete - %(eta_td)s remaining" bar.next() 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() bar.finish() break asyncio.run(main())