Transphase/transmat/transmat.py

#!/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())