#!/usr/bin/env python import asyncio, websockets import sys, os, base64, argparse, json, pickle from xkcdpass import xkcd_password as xp from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from progress.bar import Bar VERSION_NUMBER="0.1.1" 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" 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") #arg_group.add_argument('--relay', action='store_true', help='Run as a Relay server (mutually exclusive with --receive and --send)') 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'{sys.argv[0]} --receive ' if args.receive: role = 'receive' password = args.password if args.send and args.password is None: wordlist = xp.generate_wordlist(wordfile = xp.locate_wordfile(), min_length = 5, max_length = 9) password = xp.generate_xkcdpassword(wordlist, numwords=4, delimiter = "-", case='capitalize') 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"{sys.argv[0]} 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, server_part, passwd_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 = round(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('Transferring', max=number_of_chunks, 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) await asyncio.sleep(0.05) chunk_id += 1 if chunk_id > number_of_chunks: break proceed = await ws.recv() bar.next() bar.suffix = '%(percent).1f%% done' print("\n") 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 i = 1 while True: message = await ws.recv() message = json.loads(message) payload = message["payload"] msg = pickle.loads(decrypt_chunk(k, payload)) if bar is None: filename = msg["filename"] number_of_chunks = msg["number_of_chunks"] bar = Bar('Receiving', max=number_of_chunks, suffix='%(percent).1f%% complete - %(eta_td)s remaining') if f is None: f = open(msg["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 i += 1 bar.next() if i > msg["number_of_chunks"]: bar.suffix = '%(percent).1f%% complete' bar.update() f.close() print("\n") break asyncio.run(main())