199 lines
6.9 KiB
Python
Executable File
199 lines
6.9 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
import asyncio, websockets
|
|
import sys, os, base64, argparse, json, pickle
|
|
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.3"
|
|
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 = 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%% 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
|
|
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"] + 1:
|
|
bar.index = 100
|
|
bar.suffix = '%(percent).1f%% complete'
|
|
bar.update()
|
|
f.close()
|
|
print("")
|
|
break
|
|
|
|
|
|
asyncio.run(main()) |