Reorganize files and folders
This commit is contained in:
201
transmat/transmat.py
Executable file
201
transmat/transmat.py
Executable file
@@ -0,0 +1,201 @@
|
||||
#!/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"
|
||||
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"
|
||||
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'{SCRIPT_BASENAME} --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"{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%% done'
|
||||
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"]:
|
||||
bar.index = 100
|
||||
bar.suffix = '%(percent).1f%% complete'
|
||||
bar.update()
|
||||
f.close()
|
||||
print("")
|
||||
break
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user