copyparty/bin/up2k.py

724 lines
21 KiB
Python
Executable file

#!/usr/bin/env python3
from __future__ import print_function, unicode_literals
"""
up2k.py: upload to copyparty
2021-09-30, v0.6, ed <irc.rizon.net>, MIT-Licensed
https://github.com/9001/copyparty/blob/hovudstraum/bin/up2k.py
- dependencies: requests
- supports python 2.7 and 3.3 through 3.10
- almost zero error-handling
- but if something breaks just try again and it'll autoresume
"""
import os
import sys
import stat
import math
import time
import atexit
import signal
import base64
import hashlib
import argparse
import platform
import threading
import requests
import datetime
# from copyparty/__init__.py
PY2 = sys.version_info[0] == 2
if PY2:
from Queue import Queue
sys.dont_write_bytecode = True
bytes = str
else:
from queue import Queue
unicode = str
VT100 = platform.system() != "Windows"
req_ses = requests.Session()
class File(object):
"""an up2k upload task; represents a single file"""
def __init__(self, top, rel, size, lmod):
self.top = top # type: bytes
self.rel = rel.replace(b"\\", b"/") # type: bytes
self.size = size # type: int
self.lmod = lmod # type: float
self.abs = os.path.join(top, rel) # type: bytes
self.name = self.rel.split(b"/")[-1].decode("utf-8", "replace") # type: str
# set by get_hashlist
self.cids = [] # type: list[tuple[str, int, int]] # [ hash, ofs, sz ]
self.kchunks = {} # type: dict[str, tuple[int, int]] # hash: [ ofs, sz ]
# set by handshake
self.ucids = [] # type: list[str] # chunks which need to be uploaded
self.wark = None # type: str
self.url = None # type: str
# set by upload
self.up_b = 0 # type: int
self.up_c = 0 # type: int
# m = "size({}) lmod({}) top({}) rel({}) abs({}) name({})\n"
# eprint(m.format(self.size, self.lmod, self.top, self.rel, self.abs, self.name))
class FileSlice(object):
"""file-like object providing a fixed window into a file"""
def __init__(self, file, cid):
# type: (File, str) -> FileSlice
self.car, self.len = file.kchunks[cid]
self.cdr = self.car + self.len
self.ofs = 0 # type: int
self.f = open(file.abs, "rb", 512 * 1024)
self.f.seek(self.car)
# https://stackoverflow.com/questions/4359495/what-is-exactly-a-file-like-object-in-python
# IOBase, RawIOBase, BufferedIOBase
funs = "close closed __enter__ __exit__ __iter__ isatty __next__ readable seekable writable"
try:
for fun in funs.split():
setattr(self, fun, getattr(self.f, fun))
except:
pass # py27 probably
def tell(self):
return self.ofs
def seek(self, ofs, wh=0):
if wh == 1:
ofs = self.ofs + ofs
elif wh == 2:
ofs = self.len + ofs # provided ofs is negative
if ofs < 0:
ofs = 0
elif ofs >= self.len:
ofs = self.len - 1
self.ofs = ofs
self.f.seek(self.car + ofs)
def read(self, sz):
sz = min(sz, self.len - self.ofs)
ret = self.f.read(sz)
self.ofs += len(ret)
return ret
def eprint(*a, **ka):
ka["file"] = sys.stderr
ka["end"] = ""
if not PY2:
ka["flush"] = True
print(*a, **ka)
if PY2:
sys.stderr.flush()
def termsize():
import os
env = os.environ
def ioctl_GWINSZ(fd):
try:
import fcntl, termios, struct, os
cr = struct.unpack("hh", fcntl.ioctl(fd, termios.TIOCGWINSZ, "1234"))
except:
return
return cr
cr = ioctl_GWINSZ(0) or ioctl_GWINSZ(1) or ioctl_GWINSZ(2)
if not cr:
try:
fd = os.open(os.ctermid(), os.O_RDONLY)
cr = ioctl_GWINSZ(fd)
os.close(fd)
except:
pass
if not cr:
try:
cr = (env["LINES"], env["COLUMNS"])
except:
cr = (25, 80)
return int(cr[1]), int(cr[0])
class CTermsize(object):
def __init__(self):
self.ev = False
self.margin = None
self.g = None
self.w, self.h = termsize()
try:
signal.signal(signal.SIGWINCH, self.ev_sig)
except:
return
thr = threading.Thread(target=self.worker)
thr.daemon = True
thr.start()
def worker(self):
while True:
time.sleep(0.5)
if not self.ev:
continue
self.ev = False
self.w, self.h = termsize()
if self.margin is not None:
self.scroll_region(self.margin)
def ev_sig(self, *a, **ka):
self.ev = True
def scroll_region(self, margin):
self.margin = margin
if margin is None:
self.g = None
eprint("\033[s\033[r\033[u")
else:
self.g = 1 + self.h - margin
m = "{0}\033[{1}A".format("\n" * margin, margin)
eprint("{0}\033[s\033[1;{1}r\033[u".format(m, self.g - 1))
ss = CTermsize()
def statdir(top):
"""non-recursive listing of directory contents, along with stat() info"""
if hasattr(os, "scandir"):
with os.scandir(top) as dh:
for fh in dh:
yield [os.path.join(top, fh.name), fh.stat()]
else:
for name in os.listdir(top):
abspath = os.path.join(top, name)
yield [abspath, os.stat(abspath)]
def walkdir(top):
"""recursive statdir"""
for ap, inf in sorted(statdir(top)):
if stat.S_ISDIR(inf.st_mode):
for x in walkdir(ap):
yield x
else:
yield ap, inf
def walkdirs(tops):
"""recursive statdir for a list of tops, yields [top, relpath, stat]"""
for top in tops:
if os.path.isdir(top):
for ap, inf in walkdir(top):
yield top, ap[len(top) + 1 :], inf
else:
sep = "{0}".format(os.sep).encode("ascii")
d, n = top.rsplit(sep, 1)
yield d, n, os.stat(top)
# from copyparty/util.py
def humansize(sz, terse=False):
"""picks a sensible unit for the given extent"""
for unit in ["B", "KiB", "MiB", "GiB", "TiB"]:
if sz < 1024:
break
sz /= 1024.0
ret = " ".join([str(sz)[:4].rstrip("."), unit])
if not terse:
return ret
return ret.replace("iB", "").replace(" ", "")
# from copyparty/up2k.py
def up2k_chunksize(filesize):
"""gives The correct chunksize for up2k hashing"""
chunksize = 1024 * 1024
stepsize = 512 * 1024
while True:
for mul in [1, 2]:
nchunks = math.ceil(filesize * 1.0 / chunksize)
if nchunks <= 256 or chunksize >= 32 * 1024 * 1024:
return chunksize
chunksize += stepsize
stepsize *= mul
# mostly from copyparty/up2k.py
def get_hashlist(file, pcb):
# type: (File, any) -> None
"""generates the up2k hashlist from file contents, inserts it into `file`"""
chunk_sz = up2k_chunksize(file.size)
file_rem = file.size
file_ofs = 0
ret = []
with open(file.abs, "rb", 512 * 1024) as f:
while file_rem > 0:
hashobj = hashlib.sha512()
chunk_sz = chunk_rem = min(chunk_sz, file_rem)
while chunk_rem > 0:
buf = f.read(min(chunk_rem, 64 * 1024))
if not buf:
raise Exception("EOF at " + str(f.tell()))
hashobj.update(buf)
chunk_rem -= len(buf)
digest = hashobj.digest()[:33]
digest = base64.urlsafe_b64encode(digest).decode("utf-8")
ret.append([digest, file_ofs, chunk_sz])
file_ofs += chunk_sz
file_rem -= chunk_sz
if pcb:
pcb(file, file_ofs)
file.cids = ret
file.kchunks = {}
for k, v1, v2 in ret:
file.kchunks[k] = [v1, v2]
def handshake(req_ses, url, file, pw, search):
# type: (requests.Session, str, File, any, bool) -> List[str]
"""
performs a handshake with the server; reply is:
if search, a list of search results
otherwise, a list of chunks to upload
"""
req = {
"hash": [x[0] for x in file.cids],
"name": file.name,
"lmod": file.lmod,
"size": file.size,
}
if search:
req["srch"] = 1
headers = {"Content-Type": "text/plain"} # wtf ed
if pw:
headers["Cookie"] = "=".join(["cppwd", pw])
if file.url:
url = file.url
elif b"/" in file.rel:
url += file.rel.rsplit(b"/", 1)[0].decode("utf-8", "replace")
while True:
try:
r = req_ses.post(url, headers=headers, json=req)
break
except:
eprint("handshake failed, retry...\n")
time.sleep(1)
try:
r = r.json()
except:
raise Exception(r.text)
if search:
return r["hits"]
try:
pre, url = url.split("://")
pre += "://"
except:
pre = ""
file.url = pre + url.split("/")[0] + r["purl"]
file.name = r["name"]
file.wark = r["wark"]
return r["hash"]
def upload(req_ses, file, cid, pw):
# type: (requests.Session, File, str, any) -> None
"""upload one specific chunk, `cid` (a chunk-hash)"""
headers = {
"X-Up2k-Hash": cid,
"X-Up2k-Wark": file.wark,
"Content-Type": "application/octet-stream",
}
if pw:
headers["Cookie"] = "=".join(["cppwd", pw])
f = FileSlice(file, cid)
try:
r = req_ses.post(file.url, headers=headers, data=f)
if not r:
raise Exception(repr(r))
_ = r.content
finally:
f.f.close()
class Daemon(threading.Thread):
def __init__(self, *a, **ka):
threading.Thread.__init__(self, *a, **ka)
self.daemon = True
class Ctl(object):
"""
this will be the coordinator which runs everything in parallel
(hashing, handshakes, uploads) but right now it's p dumb
"""
def __init__(self, ar):
self.ar = ar
ar.files = [
os.path.abspath(os.path.realpath(x.encode("utf-8"))) for x in ar.files
]
ar.url = ar.url.rstrip("/") + "/"
if "://" not in ar.url:
ar.url = "http://" + ar.url
eprint("\nscanning {0} locations\n".format(len(ar.files)))
nfiles = 0
nbytes = 0
for _, _, inf in walkdirs(ar.files):
nfiles += 1
nbytes += inf.st_size
eprint("found {0} files, {1}\n\n".format(nfiles, humansize(nbytes)))
self.nfiles = nfiles
self.nbytes = nbytes
if ar.td:
req_ses.verify = False
if ar.te:
req_ses.verify = ar.te
self.filegen = walkdirs(ar.files)
if ar.safe:
self.safe()
else:
self.fancy()
def safe(self):
"""minimal basic slow boring fallback codepath"""
search = self.ar.s
for nf, (top, rel, inf) in enumerate(self.filegen):
file = File(top, rel, inf.st_size, inf.st_mtime)
upath = file.abs.decode("utf-8", "replace")
print("{0} {1}\n hash...".format(self.nfiles - nf, upath))
get_hashlist(file, None)
while True:
print(" hs...")
hs = handshake(req_ses, self.ar.url, file, self.ar.a, search)
if search:
if hs:
for hit in hs:
print(" found: {0}{1}".format(self.ar.url, hit["rp"]))
else:
print(" NOT found")
break
file.ucids = hs
if not hs:
break
print("{0} {1}".format(self.nfiles - nf, upath))
ncs = len(hs)
for nc, cid in enumerate(hs):
print(" {0} up {1}".format(ncs - nc, cid))
upload(req_ses, file, cid, self.ar.a)
print(" ok!")
def fancy(self):
self.hash_f = 0
self.hash_c = 0
self.hash_b = 0
self.up_f = 0
self.up_c = 0
self.up_b = 0
self.up_br = 0
self.hasher_busy = 1
self.handshaker_busy = 0
self.uploader_busy = 0
self.t0 = time.time()
self.t0_up = None
self.spd = None
self.mutex = threading.Lock()
self.q_handshake = Queue() # type: Queue[File]
self.q_recheck = Queue() # type: Queue[File] # partial upload exists [...]
self.q_upload = Queue() # type: Queue[tuple[File, str]]
self.st_hash = [None, "(idle, starting...)"] # type: tuple[File, int]
self.st_up = [None, "(idle, starting...)"] # type: tuple[File, int]
if VT100:
atexit.register(self.cleanup_vt100)
ss.scroll_region(3)
Daemon(target=self.hasher).start()
for _ in range(self.ar.j):
Daemon(target=self.handshaker).start()
Daemon(target=self.uploader).start()
idles = 0
while idles < 3:
time.sleep(0.07)
with self.mutex:
if (
self.q_handshake.empty()
and self.q_upload.empty()
and not self.hasher_busy
and not self.handshaker_busy
and not self.uploader_busy
):
idles += 1
else:
idles = 0
if VT100:
maxlen = ss.w - len(str(self.nfiles)) - 14
txt = "\033[s\033[{0}H".format(ss.g)
for y, k, st, f in [
[0, "hash", self.st_hash, self.hash_f],
[1, "send", self.st_up, self.up_f],
]:
txt += "\033[{0}H{1}:".format(ss.g + y, k)
file, arg = st
if not file:
txt += " {0}\033[K".format(arg)
else:
if y:
p = 100 * file.up_b / file.size
else:
p = 100 * arg / file.size
name = file.abs.decode("utf-8", "replace")[-maxlen:]
if "/" in name:
name = "\033[36m{0}\033[0m/{1}".format(*name.rsplit("/", 1))
m = "{0:6.1f}% {1} {2}\033[K"
txt += m.format(p, self.nfiles - f, name)
txt += "\033[{0}H ".format(ss.g + 2)
else:
txt = " "
if not self.up_br:
spd = self.hash_b / (time.time() - self.t0)
eta = (self.nbytes - self.hash_b) / (spd + 1)
else:
spd = self.up_br / (time.time() - self.t0_up)
spd = self.spd = (self.spd or spd) * 0.9 + spd * 0.1
eta = (self.nbytes - self.up_b) / (spd + 1)
spd = humansize(spd)
eta = str(datetime.timedelta(seconds=int(eta)))
left = humansize(self.nbytes - self.up_b)
tail = "\033[K\033[u" if VT100 else "\r"
m = "eta: {0} @ {1}/s, {2} left".format(eta, spd, left)
eprint(txt + "\033]0;{0}\033\\\r{1}{2}".format(m, m, tail))
def cleanup_vt100(self):
ss.scroll_region(None)
eprint("\033[J\033]0;\033\\")
def cb_hasher(self, file, ofs):
self.st_hash = [file, ofs]
def hasher(self):
for top, rel, inf in self.filegen:
file = File(top, rel, inf.st_size, inf.st_mtime)
while True:
with self.mutex:
if (
self.hash_b - self.up_b < 1024 * 1024 * 128
and self.hash_c - self.up_c < 64
and (
not self.ar.nh
or (
self.q_upload.empty()
and self.q_handshake.empty()
and not self.uploader_busy
)
)
):
break
time.sleep(0.05)
get_hashlist(file, self.cb_hasher)
with self.mutex:
self.hash_f += 1
self.hash_c += len(file.cids)
self.hash_b += file.size
self.q_handshake.put(file)
self.hasher_busy = 0
self.st_hash = [None, "(finished)"]
def handshaker(self):
search = self.ar.s
q = self.q_handshake
while True:
file = q.get()
if not file:
if q == self.q_handshake:
q = self.q_recheck
q.put(None)
continue
self.q_upload.put(None)
break
with self.mutex:
self.handshaker_busy += 1
upath = file.abs.decode("utf-8", "replace")
try:
hs = handshake(req_ses, self.ar.url, file, self.ar.a, search)
except Exception as ex:
if q == self.q_handshake and "<pre>partial upload exists" in str(ex):
self.q_recheck.put(file)
hs = []
else:
raise
if search:
if hs:
for hit in hs:
m = "found: {0}\n {1}{2}\n"
print(m.format(upath, self.ar.url, hit["rp"]), end="")
else:
print("NOT found: {0}\n".format(upath), end="")
with self.mutex:
self.up_f += 1
self.up_c += len(file.cids)
self.up_b += file.size
self.handshaker_busy -= 1
continue
with self.mutex:
if not hs:
# all chunks done
self.up_f += 1
self.up_c += len(file.cids) - file.up_c
self.up_b += file.size - file.up_b
if hs and file.up_c:
# some chunks failed
self.up_c -= len(hs)
file.up_c -= len(hs)
for cid in hs:
sz = file.kchunks[cid][1]
self.up_b -= sz
file.up_b -= sz
file.ucids = hs
self.handshaker_busy -= 1
if not hs:
print("uploaded {0}".format(upath))
for cid in hs:
self.q_upload.put([file, cid])
def uploader(self):
while True:
task = self.q_upload.get()
if not task:
self.st_up = [None, "(finished)"]
break
with self.mutex:
self.uploader_busy += 1
self.t0_up = self.t0_up or time.time()
file, cid = task
try:
upload(req_ses, file, cid, self.ar.a)
except:
eprint("upload failed, retry...\n")
pass # handshake will fix it
with self.mutex:
sz = file.kchunks[cid][1]
file.ucids = [x for x in file.ucids if x != cid]
if not file.ucids:
self.q_handshake.put(file)
self.st_up = [file, cid]
file.up_b += sz
self.up_b += sz
self.up_br += sz
file.up_c += 1
self.up_c += 1
self.uploader_busy -= 1
def main():
time.strptime("19970815", "%Y%m%d") # python#7980
if not VT100:
os.system("rem") # enables colors
# fmt: off
ap = app = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
ap.add_argument("url", type=unicode, help="server url, including destination folder")
ap.add_argument("files", type=unicode, nargs="+", help="files and/or folders to process")
ap.add_argument("-a", metavar="PASSWORD", help="password")
ap.add_argument("-s", action="store_true", help="file-search (disables upload)")
ap = app.add_argument_group("performance tweaks")
ap.add_argument("-j", type=int, metavar="THREADS", default=4, help="parallel connections")
ap.add_argument("-nh", action="store_true", help="disable hashing while uploading")
ap.add_argument("--safe", action="store_true", help="use simple fallback approach")
ap = app.add_argument_group("tls")
ap.add_argument("-te", metavar="PEM_FILE", help="certificate to expect/verify")
ap.add_argument("-td", action="store_true", help="disable certificate check")
# fmt: on
Ctl(app.parse_args())
if __name__ == "__main__":
main()