Commit aa0b2b02 authored by setop's avatar setop 🚀
Browse files

add progress bar

parent 556dd490
#!/usr/bin/env python3
import baker
import sys
import os
from tqdm import tqdm
from websocket import create_connection
from json import loads, dumps
from sjcl import SJCL, sjcl
from base64 import b64decode, b64encode
MAGIC = "XXMOJOXX"
CHUNK_SIZE = 2*1000*1000
CHUNK_SIZE = 1*1024*1024 # only used for upload
@baker.command(
......@@ -31,7 +31,7 @@ def upload(filepath, url, delay=None, del_at_first_view=False):
f = open(filepath, "rb")
wsurl = url.replace("http", "ws")+"upload/"
ws = create_connection(wsurl)
for (i, data) in enumerate(iter(lambda: f.read(CHUNK_SIZE), b"")):
for (i, data) in tqdm(enumerate(iter(lambda: f.read(CHUNK_SIZE), b"")), total=totalparts):
meta["part"] = i
data = SJCL().encrypt(b64encode(data), key) # b64 is unnecessary be kept for compatibility
for k in data.keys():
......@@ -39,7 +39,6 @@ def upload(filepath, url, delay=None, del_at_first_view=False):
data[k] = data[k].decode("ascii")
ws.send(payload=dumps(meta)+MAGIC+dumps(dumps(data))) # second dumps is unnecessary be kept for compatibility
result = ws.recv()
print(result, file=sys.stderr)
result = loads(result)
meta["id"] = result["short"]
ws.close()
......@@ -56,16 +55,18 @@ def download(url):
ws = create_connection(url)
p = 0
t = 1
progress = tqdm()
while p < t:
ws.send('{"part":%d}' % p)
result = ws.recv()
idx = result.find(MAGIC)
meta = loads(result[0:idx])
print(meta, file=sys.stderr)
data = loads(loads(result[idx+len(MAGIC):])) # second loads is unnecessary be kept for compatibility
t = meta["total"]
progress.total = t
f = open(meta["name"].split("/")[-1], "ab")
f.write(b64decode(SJCL().decrypt(data, key))) # b64 is unnecessary be kept for compatibility
progress.update()
p += 1
ws.close()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment