diff --git a/.gitignore b/.gitignore index 2879dc6..fa715f8 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ /result flamegraph.svg perf.data* -.env \ No newline at end of file +.env +/cli/env \ No newline at end of file diff --git a/cli/requirements.txt b/cli/requirements.txt new file mode 100644 index 0000000..8c73512 --- /dev/null +++ b/cli/requirements.txt @@ -0,0 +1,2 @@ +tqdm +websockets diff --git a/cli/transbeam.py b/cli/transbeam.py new file mode 100755 index 0000000..c21451e --- /dev/null +++ b/cli/transbeam.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 + +import argparse +import asyncio +import getpass +import json +import math +import pathlib +import sys +from tqdm import tqdm +from websockets import connect + +def file_loader(files): + with tqdm(desc="Total", total=sum(size for (path, size) in files), unit='B', unit_scale=True, leave=True, position=1) as total_progress: + for (path, size) in files: + with tqdm(desc=path.name, total=size, unit='B', unit_scale=True, leave=True, position=0) as file_progress: + with path.open("rb") as f: + while f.tell() < size: + data = f.read(min(16384, size - f.tell())) + if data == "": + tqdm.write("file ended early!") + exit(1) + total_progress.update(len(data)) + file_progress.update(len(data)) + yield data + +async def send(paths, uri, password, lifetime): + paths = [path for path in paths if path.is_file()] + fileMetadata = [ + { + "name": path.name, + "size": path.stat().st_size, + "modtime": math.floor(path.stat().st_mtime * 1000), + } for path in paths + ] + manifest = { + "files": fileMetadata, + "lifetime": lifetime, + "password": password, + } + async with connect(uri) as ws: + await ws.send(json.dumps(manifest)) + resp = json.loads(await ws.recv()) + if resp["type"] != "ready": + print("unexpected response: {}".format(resp)) + exit(1) + print("Download code: {}".format(resp["code"])) + loader = file_loader([(paths[i], fileMetadata[i]["size"]) for i in range(len(paths))]) + for data in loader: + await ws.send(data) + resp = await ws.recv() + if resp != "ack": + tqdm.write("unexpected response: {}".format(resp)) + exit(1) + +parser = argparse.ArgumentParser(description="Upload files to transbeam") +parser.add_argument("-l", "--lifetime", type=int, default=1, help="Lifetime in days for files (default 10)") +parser.add_argument("-u", "--uri", type=str, default="wss://transbeam.link/upload", help="Websocket URI for transbeam") +parser.add_argument("files", type=pathlib.Path, nargs="+", help="Files to upload") + +async def main(): + args = parser.parse_args() + password = getpass.getpass() + await send(args.files, args.uri, password, args.lifetime) + +asyncio.run(main())