PutIO needs tus lib

libs/tus/__init__.py

@@ -0,0 +1,190 @@
+import os
+import base64
+import logging
+import argparse
+
+import requests
+
+LOG_LEVEL = logging.INFO
+DEFAULT_CHUNK_SIZE = 4 * 1024 * 1024
+TUS_VERSION = '1.0.0'
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+logger.addHandler(logging.NullHandler())
+
+
+class TusError(Exception):
+    pass
+
+
+def _init():
+    fmt = "[%(asctime)s] %(levelname)s %(message)s"
+    h = logging.StreamHandler()
+    h.setLevel(LOG_LEVEL)
+    h.setFormatter(logging.Formatter(fmt))
+    logger.addHandler(h)
+
+
+def _create_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('file', type=argparse.FileType('rb'))
+    parser.add_argument('--chunk-size', type=int, default=DEFAULT_CHUNK_SIZE)
+    parser.add_argument(
+        '--header',
+        action='append',
+        help="A single key/value pair"
+        " to be sent with all requests as HTTP header."
+        " Can be specified multiple times to send more then one header."
+        " Key and value must be separated with \":\".")
+    return parser
+
+
+def _cmd_upload():
+    _init()
+
+    parser = _create_parser()
+    parser.add_argument('tus_endpoint')
+    parser.add_argument('--file_name')
+    parser.add_argument(
+        '--metadata',
+        action='append',
+        help="A single key/value pair to be sent in Upload-Metadata header."
+        " Can be specified multiple times to send more than one pair."
+        " Key and value must be separated with space.")
+    args = parser.parse_args()
+
+    headers = dict([x.split(':') for x in args.header])
+    metadata = dict([x.split(' ') for x in args.metadata])
+
+    upload(
+        args.file,
+        args.tus_endpoint,
+        chunk_size=args.chunk_size,
+        file_name=args.file_name,
+        headers=headers,
+        metadata=metadata)
+
+
+def _cmd_resume():
+    _init()
+
+    parser = _create_parser()
+    parser.add_argument('file_endpoint')
+    args = parser.parse_args()
+
+    headers = dict([x.split(':') for x in args.header])
+
+    resume(
+        args.file,
+        args.file_endpoint,
+        chunk_size=args.chunk_size,
+        headers=headers)
+
+
+def upload(file_obj,
+           tus_endpoint,
+           chunk_size=DEFAULT_CHUNK_SIZE,
+           file_name=None,
+           headers=None,
+           metadata=None):
+    file_name = os.path.basename(file_obj.name)
+    file_size = _get_file_size(file_obj)
+    location = _create_file(
+        tus_endpoint,
+        file_name,
+        file_size,
+        extra_headers=headers,
+        metadata=metadata)
+    resume(
+        file_obj, location, chunk_size=chunk_size, headers=headers, offset=0)
+
+
+def _get_file_size(f):
+    pos = f.tell()
+    f.seek(0, 2)
+    size = f.tell()
+    f.seek(pos)
+    return size
+
+
+def _create_file(tus_endpoint,
+                 file_name,
+                 file_size,
+                 extra_headers=None,
+                 metadata=None):
+    logger.info("Creating file endpoint")
+
+    headers = {
+        "Tus-Resumable": TUS_VERSION,
+        "Upload-Length": str(file_size),
+    }
+
+    if extra_headers:
+        headers.update(extra_headers)
+
+    if metadata:
+        l = [k + ' ' + base64.b64encode(v) for k, v in metadata.items()]
+        headers["Upload-Metadata"] = ','.join(l)
+
+    response = requests.post(tus_endpoint, headers=headers)
+    if response.status_code != 201:
+        raise TusError("Create failed: %s" % response)
+
+    location = response.headers["Location"]
+    logger.info("Created: %s", location)
+    return location
+
+
+def resume(file_obj,
+           file_endpoint,
+           chunk_size=DEFAULT_CHUNK_SIZE,
+           headers=None,
+           offset=None):
+    if offset is None:
+        offset = _get_offset(file_endpoint, extra_headers=headers)
+
+    total_sent = 0
+    file_size = _get_file_size(file_obj)
+    while offset < file_size:
+        file_obj.seek(offset)
+        data = file_obj.read(chunk_size)
+        offset = _upload_chunk(
+            data, offset, file_endpoint, extra_headers=headers)
+        total_sent += len(data)
+        logger.info("Total bytes sent: %i", total_sent)
+
+
+def _get_offset(file_endpoint, extra_headers=None):
+    logger.info("Getting offset")
+
+    headers = {"Tus-Resumable": TUS_VERSION}
+
+    if extra_headers:
+        headers.update(extra_headers)
+
+    response = requests.head(file_endpoint, headers=headers)
+    response.raise_for_status()
+
+    offset = int(response.headers["Upload-Offset"])
+    logger.info("offset=%i", offset)
+    return offset
+
+
+def _upload_chunk(data, offset, file_endpoint, extra_headers=None):
+    logger.info("Uploading chunk from offset: %i", offset)
+
+    headers = {
+        'Content-Type': 'application/offset+octet-stream',
+        'Upload-Offset': str(offset),
+        'Tus-Resumable': TUS_VERSION,
+    }
+
+    if extra_headers:
+        headers.update(extra_headers)
+
+    response = requests.patch(file_endpoint, headers=headers, data=data)
+    if response.status_code != 204:
+        raise TusError("Upload chunk failed: %s" % response)
+
+    return int(response.headers["Upload-Offset"])