#! /usr/bin/python3

# Monitors the filesystem for completed warts files generated by ark
# activities, and uploads them to the ark-collector server.
#
# Currently encodes the data in the body of the POST similar to:
#
# curl --data-binary @foo.warts.gz \
#      -H "Content-Disposition: file; filename=\"foo.warts.gz\"" \
#      http://collector/uploads/prefix-probing
#
# For testing or manual uploads with curl be aware that --data-binary loads
# the entire file into memory so maybe an alternative like this is required:
#
# curl -T foo.warts.gz \
#      -X POST \
#      -H "Expect:" \
#      -H "Content-Disposition: file; filename=\"foo.warts.gz\""
#      http://collector/uploads/prefix-probing

import argparse
import contextlib
from http import HTTPStatus
from pathlib import Path
import random
import re
import signal
import time

import requests

try:
    from inotify.adapters import Inotify
    from inotify.calls import InotifyError
    from inotify.constants import IN_CLOSE_WRITE, IN_MOVED_TO
    HAS_INOTIFY = True
except ImportError:
    HAS_INOTIFY = False

from prometheus_client import (
    GC_COLLECTOR,
    PLATFORM_COLLECTOR,
    PROCESS_COLLECTOR,
    start_http_server,
)
from prometheus_client.core import REGISTRY, Counter, Gauge

try:
    # only added in python3-prometheus-client 0.18.0, but makes things tidy
    from prometheus_client import disable_created_metrics
    disable_created_metrics()
except ImportError:
    pass


FILENAME_REGEX = re.compile(r"/.*/[0-9a-z.-]+.warts.gz")
ACTIVITIES = ["team-probing", "topo-v6", "prefix-probing"]
MAX_RETRY_DELAY = 3600

TEMPORARY_FAILURES = [
    HTTPStatus.TOO_MANY_REQUESTS,
    HTTPStatus.INTERNAL_SERVER_ERROR,
    HTTPStatus.BAD_GATEWAY,
    HTTPStatus.SERVICE_UNAVAILABLE,
    HTTPStatus.GATEWAY_TIMEOUT,
]


class ArkUploader:
    """Watch directories for results files and upload them to the collector."""

    def __init__(self, monitor, collector, directories, remove):
        self.monitor = monitor
        self.collector = collector
        self.directories = directories
        self.remove = remove
        self.metrics = self._setup_metrics()
        self._running = False
        signal.signal(signal.SIGINT, self._graceful_shutdown)
        signal.signal(signal.SIGTERM, self._graceful_shutdown)


    def _setup_metrics(self):
        """Create all the prometheus metrics."""
        metrics = {}
        #disable_created_metrics()

        # remove autogenerated python process metrics
        REGISTRY.unregister(PROCESS_COLLECTOR)
        REGISTRY.unregister(PLATFORM_COLLECTOR)
        REGISTRY.unregister(GC_COLLECTOR)

        metrics["requests_total"] = Counter(
                "requests_total",
                "Total number of result files that we have tried to upload",
                labelnames=["activity", "status"],
                namespace="uploader",
        )
        metrics["last_success_timestamp"] = Gauge(
                "last_success_timestamp",
                "Timestamp when the last file was successfully uploaded",
                namespace="uploader",
        )
        metrics["last_attempt_timestamp"] = Gauge(
                "last_attempt_timestamp",
                "Timestamp when the last file was attempted to be uploaded",
                namespace="uploader",
        )

        return metrics


    def _graceful_shutdown(self, _signum, _frame):
        # existing uploads will be allowed to finish, but the watched file
        # file generators will exit so that no more files are uploaded
        self._running = False
        # TODO if running is already false then just exit?


    # TODO would it be better to explicitly state on the command line which
    # activity a watched directory belongs to?
    # TODO or easier to look for known activity strings anywhere in the path?
    # e.g. --dir "team-probing:/var/lib/ark/activity/team-probing/finished"
    def _get_activity(self, filepath):
        match = re.match(r"^.*/([0-9a-z.-]+)/finished/.*$", str(filepath))
        if match and match.group(1) in ACTIVITIES:
            return match.group(1)
        return None


    def upload(self, filepath, dryrun=False):
        """Upload a results file to the collector."""

        activity = self._get_activity(filepath)
        if activity is None:
            print(f"Skipping upload due to unknown activity: {filepath}")
            return

        url = f"https://{self.collector}/upload/{activity}"
        print(f"Uploading {filepath} to {url}")

        if dryrun:
            return

        # if we don't explicitly set Content-Length when using the
        # `data` parameter then we end up in the chunked encoding
        # codepath, which can't do retries for connection failures?
        headers = {
            "Connection": "close",
            "Content-Disposition": f"file; filename=\"{filepath.name}\"",
            "Content-Type": "application/octet-stream",
            "Content-Length": str(filepath.stat().st_size),
        }

        # TODO can we send Expect:100-continue? I don't think the python
        # server handles that appropriately by default?
        errors = 0

        # try to upload the file until the process is told to stop, or we get
        # a permanent error and give up
        while self._running:
            response = None
            try:
                with filepath.open("rb") as data:
                    # passing a file-like object as body data will apparently
                    # stream the file without loading it entirely into memory.
                    self.metrics["last_attempt_timestamp"].set(time.time())
                    response = requests.post(
                        url=url,
                        verify=True,
                        cert=(f'/etc/ark/ssl/{self.monitor}.ark.caida.org.crt',
                            f'/etc/ark/ssl/{self.monitor}.ark.caida.org.key'),
                        headers=headers,
                        data=data,
                        timeout=(10, 20),
                        )
                response.raise_for_status()
            except FileNotFoundError:
                print(f"Skipping upload of missing file: {filepath}")
                return
            except (requests.exceptions.ConnectionError,
                    requests.exceptions.Timeout,
                    requests.exceptions.SSLError,
                    requests.HTTPError) as e:
                print(e)
                # these are hopefully transient errors that should be retried
                if response is None or response.status_code in TEMPORARY_FAILURES:
                    try:
                        delay = int(response.headers["Retry-After"])
                        errors = 0
                    except (KeyError, ValueError, AttributeError):
                        delay = min(2 ** errors, MAX_RETRY_DELAY)
                        errors += 1
                    self._responsive_sleep(delay + random.randrange(delay))
                    continue
            # stop retrying and process any response
            break

        # stopped running while between upload attempts, without a result
        if response is None or response.status_code in TEMPORARY_FAILURES:
            print("Stopping incomplete upload")
            return

        # by now we've received an http response of some kind, check it
        self.metrics["requests_total"].labels(activity=activity,
                    status=response.status_code).inc()

        if response.ok:
            print(f"File {filepath} uploaded successfully")
            self.metrics["last_success_timestamp"].set(time.time())
        elif response.status_code == HTTPStatus.CONFLICT:
            print(f"File {filepath} already uploaded")
        else:
            # any other error is a permanent failure, give up on this file
            print(f"Failed to upload file: {filepath}")

        # we've either succeeded in the upload, or the error isn't
        # something that we can fix by retrying, so delete the file
        if self.remove:
            print(f"Removing local file {filepath}")
            with contextlib.suppress(FileNotFoundError):
                filepath.unlink()


    def run(self, poll=False, dryrun=False):
        """Start watching and uploading results files."""

        self._running = True

        # decide which approach we take to watching for new files
        watcher = self._watch_poll() if poll else self._watch_inotify()

        # start prometheus metrics endpoint
        print("Starting metrics endpoint at http://127.0.0.1:8000/metrics")
        server, thread = start_http_server(8000, addr="127.0.0.1")

        print("Watching for results files in the following locations:")
        for directory in self.directories:
            print(f"  {directory}")

        # watch for files until signalled to stop
        for filepath in watcher:
            if FILENAME_REGEX.fullmatch(str(filepath)):
                self.upload(filepath, dryrun)

        # stop the prometheus metrics endpoint
        server.shutdown()
        server.server_close()
        thread.join()


    def _responsive_sleep(self, seconds):
        for _ in range(seconds):
            if not self._running:
                return
            time.sleep(1)


    def _check(self):
        """Return all results files in watched directories."""
        # TODO remove any files older than a week or so rather than uploading?
        for directory in self.directories:
            filepaths = sorted(Path(directory).glob("*.warts.gz"))
            yield from filepaths


    def _watch_poll(self):
        # if we're only polling then it's easy just to regularly check for new
        # files in the monitored directories until asked to stop
        while self._running:
            for filepath in self._check():
                if not self._running:
                    return
                yield filepath
            self._responsive_sleep(30)


    def _watch_inotify(self):
        # start monitoring for any new results files in the target directories
        monitor = Inotify()
        for directory in self.directories:
            try:
                monitor.add_watch(directory, mask=IN_MOVED_TO|IN_CLOSE_WRITE)
            except InotifyError:
                print(f"Ignoring missing directory: {directory}")

        # upload any existing files already in those directories
        for filepath in self._check():
            if not self._running:
                return
            yield filepath

        # wait for and process any new files that appear
        for event in monitor.event_gen(yield_nones=True):
            if not self._running:
                return
            if event:
                (_, _type_names, path, filename) = event
                yield Path(path, filename)


def main():
    # TODO is it worth trying to build this list dynamically?
    directories = [
        "/var/lib/ark/activity/prefix-probing/finished/",
        "/var/lib/ark/activity/team-probing/finished/",
        "/var/lib/ark/activity/topo-v6/finished/",
    ]

    # TODO do we need to be able to set the monitor config file, the monitor
    # name, certificate file locations, etc?
    parser = argparse.ArgumentParser(description='ark uploader')
    parser.add_argument("--collector",
            default="collector.ark.caida.org",
            help="Collector server where results should be uploaded")
    parser.add_argument("--dir", action="append", default=directories,
            help="Add directory to watch (may be specified multiple times)")
    parser.add_argument("--dryrun", default=False, action="store_true",
            help="Don't actually upload or delete files")
    parser.add_argument("--no-remove", default=True, action="store_false",
            dest="remove",
            help="Keep files after uploading, don't delete them")
    parser.add_argument("--monitor", required=True,
            help="Ark monitor name")
    parser.add_argument("--poll", action='store_true',
            default=(not HAS_INOTIFY),
            help="Poll for new files instead of using inotify")
    args = parser.parse_args()

    uploader = ArkUploader(args.monitor, args.collector, args.dir, args.remove)
    uploader.run(args.poll, args.dryrun)


if __name__ == "__main__":
    main()
