#!/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.1.8" 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 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): ( msgtype, peer_group_id, role, 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 = { "msgtype": msgtype, "peer_group_id": peer_group_id, "role": role, "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}") 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) chunk_id = 0 chunk_size = 1024 * 512 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, "")) while True: 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) await send_encrypted_msg(ws, k, msg) chunk_id += 1 if chunk_id > number_of_chunks: break proceed = await ws.recv() bar.next() 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, "", "", "", "", "")) bar = None f = None while True: message = await ws.recv() message = json.loads(message) 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"]) msgtype = "proceed" 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.next() if chunk_id >= number_of_chunks-1: # This was the last chunk, exit f.close() bar.suffix = '%(percent).1f%% complete' bar.update() bar.finish() break asyncio.run(main())