add quick and dirty CLI client

This commit is contained in:
xenofem 2022-05-24 21:45:37 -04:00
parent 7de5997d1b
commit c60b066d49
3 changed files with 70 additions and 1 deletions

3
.gitignore vendored
View file

@ -3,4 +3,5 @@
/result
flamegraph.svg
perf.data*
.env
.env
/cli/env

2
cli/requirements.txt Normal file
View file

@ -0,0 +1,2 @@
tqdm
websockets

66
cli/transbeam.py Executable file
View file

@ -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())