offload-sender/main.py

#!/usr/bin/env python
import base64
import datetime
import json
import logging
import os
import pathlib
import shutil
import time
from typing import Optional
import urllib.parse
import re

import certifi
import click
import minio
import requests
import urllib3

from progress import Progress

logging.basicConfig(level=logging.INFO)
BACKFEED_DELIM = "\n"
SKIP_BACKFEED = "SKIPBF"
SKIP_DISPATCHER = "SKIPDISPATCH://"


# TODO: Add rsync support
# TODO: Add rsync+ssh support
# TODO: Add webdav support.
# TODO: Fix the "ctrl-c handling" logic so it actually cleans up in the s3 bucket.

def retry_failures(fn, msg, *args, **kwargs):
    tries = 0
    while True:
        try:
            return fn(*args, **kwargs)
        except Exception:
            logging.exception(msg)
            delay = min(2 ** tries, 64)
            tries = tries + 1
            logging.info(f"Sleeping {delay} seconds...")
            time.sleep(delay)


@click.group()
def sender():
    pass


def watch_pass(input_directory: pathlib.Path, work_directory: pathlib.Path, ia_collection: str, ia_item_title: str,
               ia_item_prefix: str, ia_item_date: str, project: str, dispatcher: str, delete: bool, backfeed_key: str):
    logging.info("Checking for new items...")
    for original_directory in input_directory.iterdir():
        if original_directory.is_dir():
            original_name = original_directory.name
            new_directory = work_directory.joinpath(original_name)
            try:
                original_directory.rename(new_directory)
            except FileNotFoundError:
                logging.warning(f"Unable to move item {original_directory}")
                continue
            single_impl(new_directory, ia_collection, ia_item_title, ia_item_prefix, ia_item_date, project,
                        dispatcher, delete, backfeed_key)
            return True
    return False


@sender.command()
@click.option('--input-directory', envvar='UPLOAD_QUEUE_DIR', default="/data/upload-queue",
              type=click.Path(exists=True))
@click.option('--work-directory', envvar='UPLOADER_WORKING_DIR', default="/data/uploader-work",
              type=click.Path(exists=True))
@click.option('--ia-collection', envvar='IA_COLLECTION', required=True)
@click.option('--ia-item-title', envvar='IA_ITEM_TITLE', required=True)
@click.option('--ia-item-prefix', envvar='IA_ITEM_PREFIX', required=True)
@click.option('--ia-item-date', envvar='IA_ITEM_DATE', required=False)
@click.option('--project', envvar='PROJECT', required=True)
@click.option('--dispatcher', envvar='DISPATCHER', required=True)
@click.option('--delete/--no-delete', envvar='DELETE', default=False)
@click.option('--backfeed-key', envvar='BACKFEED_KEY', required=True)
def watch(input_directory: pathlib.Path, work_directory: pathlib.Path, ia_collection: str, ia_item_title: str,
          ia_item_prefix: str, ia_item_date: str, project: str, dispatcher: str, delete: bool, backfeed_key: str):
    if not isinstance(input_directory, pathlib.Path):
        input_directory = pathlib.Path(input_directory)
    if not isinstance(work_directory, pathlib.Path):
        work_directory = pathlib.Path(work_directory)

    while True:
        if not watch_pass(input_directory, work_directory, ia_collection, ia_item_title, ia_item_prefix, ia_item_date,
                          project, dispatcher, delete, backfeed_key):
            logging.info("No item found, sleeping...")
            time.sleep(10)


@sender.command()
@click.option('--item-directory', type=click.Path(exists=True), required=True)
@click.option('--ia-collection', envvar='IA_COLLECTION', required=True)
@click.option('--ia-item-title', envvar='IA_ITEM_TITLE', required=True)
@click.option('--ia-item-prefix', envvar='IA_ITEM_PREFIX', required=True)
@click.option('--ia-item-date', envvar='IA_ITEM_DATE', required=False)
@click.option('--project', envvar='PROJECT', required=True)
@click.option('--dispatcher', envvar='DISPATCHER', required=True)
@click.option('--delete/--no-delete', envvar='DELETE', default=False)
@click.option('--backfeed-key', envvar='BACKFEED_KEY', required=True)
def single(item_directory: pathlib.Path, ia_collection: str, ia_item_title: str, ia_item_prefix: str,
           ia_item_date: Optional[str], project: str, dispatcher: str, delete: bool, backfeed_key: str):
    single_impl(item_directory, ia_collection, ia_item_title, ia_item_prefix, ia_item_date, project, dispatcher, delete,
                backfeed_key)


def single_impl(item_directory: pathlib.Path, ia_collection: str, ia_item_title: str, ia_item_prefix: str,
                ia_item_date: Optional[str], project: str, dispatcher: str, delete: bool, backfeed_key: str):
    if not isinstance(item_directory, pathlib.Path):
        item_directory = pathlib.Path(item_directory)

    logging.info(f"Processing item {item_directory}...")

    if ia_item_date is None:
        s = item_directory.name.split("_")
        if len(s) > 0:
            ds = s[0]
            try:
                d = datetime.datetime.strptime(ds, "%Y%m%d%H%M%S")
                ia_item_date = d.strftime("%Y-%m")
            except ValueError:
                pass

    meta_json_loc = item_directory.joinpath('__upload_meta.json')
    if meta_json_loc.exists():
        raise Exception("META JSON EXISTS WTF")
    meta_json = {
        "IA_COLLECTION": ia_collection,
        "IA_ITEM_TITLE": f"{ia_item_title} {item_directory.name}",
        "IA_ITEM_DATE": ia_item_date,
        "IA_ITEM_NAME": f"{ia_item_prefix}{item_directory.name}",
        "PROJECT": project,
    }
    with open(meta_json_loc, 'w') as f:
        f.write(json.dumps(meta_json))
        logging.info("Wrote metadata json.")
    total_size = 0
    files = list(item_directory.glob("**/*"))
    for item in files:
        total_size = total_size + os.path.getsize(item)
    logging.info(f"Item size is {total_size} bytes across {len(files)} files.")
    meta_json["SIZE_HINT"] = str(total_size)

    def assign_target():
        if dispatcher.startswith(SKIP_DISPATCHER):
            target = dispatcher[len(SKIP_DISPATCHER):]
            logging.info(f"Skipping dispatcher, using {target} as target.")
            return target
        logging.info("Attempting to assign target...")
        r = requests.get(f"{dispatcher}/offload_target", params=meta_json, timeout=60)
        if r.status_code == 200:
            data = r.json()
            return data["url"]
        else:
            raise Exception(f"Invalid status code {r.status_code}: {r.text}")

    url = retry_failures(assign_target, "Failed to fetch target")
    logging.info(f"Assigned target {url}")

    parsed_url = urllib.parse.urlparse(url)
    parsed_qs = urllib.parse.parse_qs(parsed_url.query)

    def get_q(key, default):
        return parsed_qs.get(key, [str(default)])[0]

    bf_item = None
    if parsed_url.scheme == "minio+http" or parsed_url.scheme == "minio+https":
        secure = (parsed_url.scheme == "minio+https")
        ep = parsed_url.hostname
        if parsed_url.port is not None:
            ep = f"{ep}:{parsed_url.port}"
        client = None

        def create_client():
            logging.info("Connecting to minio...")
            cert_check = True
            timeout = datetime.timedelta(seconds=int(get_q("timeout", 60))).seconds
            total_timeout = datetime.timedelta(seconds=int(get_q("total_timeout", timeout*2))).seconds
            hclient = urllib3.PoolManager(
                timeout=urllib3.util.Timeout(connect=timeout, read=timeout, total=total_timeout),
                maxsize=10,
                cert_reqs='CERT_REQUIRED' if cert_check else 'CERT_NONE',
                ca_certs=os.environ.get('SSL_CERT_FILE') or certifi.where(),
                retries=urllib3.Retry(
                    total=5,
                    backoff_factor=0.2,
                    backoff_max=30,
                    status_forcelist=[500, 502, 503, 504]
                )
            )
            return minio.Minio(endpoint=ep, access_key=parsed_url.username, secret_key=parsed_url.password,
                               secure=secure, http_client=hclient)

        client = retry_failures(create_client, "Failed to connect to minio")

        bucket_name = item_directory.name.replace("_", "-")

        def make_bucket():
            logging.info("Attempting to make bucket...")
            if client.bucket_exists(bucket_name=bucket_name):
                # If bucket already exists a previous attempt was aborted.
                logging.warning("Bucket already exists!")
                return
            client.make_bucket(bucket_name=bucket_name)

        retry_failures(make_bucket, "Failed to make bucket")

        logging.info("Starting uploads...")
        for file in files:
            rel_file = file.relative_to(item_directory)

            def upload_file():
                logging.info(f"Uploading file {rel_file}...")
                client.fput_object(bucket_name=bucket_name, object_name=str(rel_file), file_path=file, num_parallel_uploads=8)

            retry_failures(upload_file, f"Failed to upload {rel_file}")

        item_data = {"url": url, "item_name": item_directory.name, "bucket_name": bucket_name}
        bf_item_part = base64.urlsafe_b64encode(str(json.dumps(item_data)).encode("UTF-8")).decode("UTF-8")
        bf_item = f"{project}:{parsed_url.hostname}:{bf_item_part}"
    elif parsed_url.scheme == "s3+cred":
        secure = True
        ep = parsed_url.hostname
        if parsed_url.port is not None:
            ep = f"{ep}:{parsed_url.port}"
        if not parsed_url.path:
            raise Exception("Invalid s3+cred url, no bucket specified. vhost style s3 urls are not supported.")
        bucket_name = parsed_url.path[1:].strip("/")
        if "/" in bucket_name:
            raise Exception("Invalid s3+cred url, path does not support sub-directories - only a main bucket in the url is supported.")
        client = None
        # ensure the cred name is reasonable
        if not re.match("^[a-zA-Z0-9_.-]+$", parsed_url.username):
            raise Exception(f"Invalid cred name in url: {parsed_url.username}")
        creds_file = f'creds/{parsed_url.username}.json'
        if not os.path.exists(creds_file):
            raise Exception(f"Unable to find credentials file {creds_file}")
        with open(creds_file) as f:
            creds = json.load(f)
        if "access_key" not in creds or "secret_key" not in creds:
            raise Exception(f"Unable to find access_key or secret_key in credentials file {creds_file}")
        def create_client():
            logging.info("Connecting to minio...")
            cert_check = True
            timeout = datetime.timedelta(seconds=int(get_q("timeout", 60))).seconds
            total_timeout = datetime.timedelta(seconds=int(get_q("total_timeout", timeout * 2))).seconds
            hclient = urllib3.PoolManager(
                timeout=urllib3.util.Timeout(connect=timeout, read=timeout, total=total_timeout),
                maxsize=10,
                cert_reqs='CERT_REQUIRED' if cert_check else 'CERT_NONE',
                ca_certs=os.environ.get('SSL_CERT_FILE') or certifi.where(),
                retries=urllib3.Retry(
                    total=5,
                    backoff_factor=0.2,
                    backoff_max=30,
                    status_forcelist=[500, 502, 503, 504]
                )
            )
            return minio.Minio(endpoint=ep, access_key=creds['access_key'], secret_key=creds['secret_key'],
                               secure=secure, http_client=hclient)

        client = retry_failures(create_client, "Failed to connect to minio")

        folder_name = item_directory.name.replace("_", "-")

        logging.info("Starting uploads...")
        for file in files:
            rel_file = file.relative_to(item_directory)

            def upload_file():
                logging.info(f"Uploading file {rel_file}...")
                client.fput_object(bucket_name = bucket_name, object_name=f"{folder_name}/{str(rel_file)}", file_path=file, num_parallel_uploads=8)

            retry_failures(upload_file, f"Failed to upload {rel_file}")

        item_data = {"url": url, "item_name": item_directory.name, "folder_name": folder_name}
        bf_item_part = base64.urlsafe_b64encode(str(json.dumps(item_data)).encode("UTF-8")).decode("UTF-8")
        bf_item = f"{project}:{parsed_url.hostname}:{bf_item_part}"
    else:
        raise Exception("Unable to upload, don't understand url: {url}")

    if bf_item is None:
        raise Exception("Unable to create backfeed item")

    if backfeed_key == SKIP_BACKFEED:
        logging.warning(f"Skipping backfeed! Would have submitted: {bf_item}")
    else:
        def submit_item():
            u = f"https://legacy-api.arpa.li/backfeed/legacy/{backfeed_key}"
            logging.info(f"Attempting to submit bf item {bf_item} to {u}...")
            resp = requests.post(u, params={"skipbloom": "1", "delimiter": BACKFEED_DELIM},
                                 data=f"{bf_item}{BACKFEED_DELIM}".encode("UTF-8"), timeout=60)
            if resp.status_code != 200:
                raise Exception(f"Failed to submit to backfeed {resp.status_code}: {resp.text}")

        retry_failures(submit_item, "Failed to submit to backfeed")
        logging.info("Backfeed submit complete!")

    if delete:
        logging.info("Removing item...")
        shutil.rmtree(item_directory)
    logging.info("Upload complete!")


if __name__ == '__main__':
    sender()