#!/usr/bin/env python3

# note that each client should only be running a single instance of the
# uploader, so there shouldn't be any race conditions caused by the same
# file being uploaded multiple times simultaneously

# accept a POST from an authenticated ark client
#   POST https://collector.ark.caida.org/upload/team-probing/
# save the file into the staging directory

# run this isolated in some way, so the staging directory being written to
# is separate from any real hosts (in a container or virtual machine?). The
# directory can be shared with ra so that archiver scripts can move files
# into correct locations.

# nginx needs to pass us the common name from the certificate so we can
# name the file appropriately
#
#   proxy_set_header X-SSL-Client-Cert-Subject-DN $ssl_client_s_dn;

# test with curl using something like:
# curl -X POST -H "Content-Type: application/octet-stream" -H "Content-Disposition: file; filename=\"foo.warts.gz\"" --data-binary @foo.warts --cert /etc/ark/ssl/aee9-zz.ark.caida.org.crt --key /etc/ark/ssl/aee9-zz.ark.caida.org.key --cacert /etc/ssl/certs/ca.ark.caida.org.crt https://collector.ark.caida.org/uploads/prefix-probing

import argparse
from datetime import datetime, timedelta, timezone
from http import HTTPStatus
import json
from pathlib import Path
import random
import re
import ssl
import threading

from ark_http_handler import ArkHTTPRequestHandler, ArkHTTPServer
from kafka import KafkaProducer

# run a proxy in front of it to check ssl client certificates
HOSTNAME = "127.0.0.1"
PORT = 8088
FILENAME_REGEX = re.compile(r"filename=\"([0-9a-z.-]+.warts.gz)\"")
DEFAULT_STAGING_DIR = "/staging"
CHUNK_SIZE = 1024 * 1024
MAX_UPLOADS = 3 # XXX test it works, then increase to a reasonable number
KAFKA_SERVER = "collector.ark.caida.org:9095"

ACTIVITIES = {
    "team-probing": {
        "name": "IPv4 team probing",
        "maxsize": 1024 * 1024,
        # hlz3-nz.team-probing.1783059791.c015310.20260703.warts.gz
        "fileregex": re.compile(r"^(?P<name>[a-z][a-z0-9-]{0,31})\.team-probing\.(?P<timestamp>[0-9]{10})\.c[0-9]+\.(?P<date>[0-9]{8})\.warts\.gz$"),
    },
    "topo-v6": {
        "name": "IPv6 probing",
        "maxsize": 1024 * 1024 * 200,
        # 1763047356.hlz2-nz.l8.20251113.warts.gz
        "fileregex": re.compile(r"^(?P<timestamp>[0-9]{10})\.(?P<name>[a-z][a-z0-9-]{0,31})\.l8\.(?P<date>[0-9]{8})\.warts\.gz$"),
    },
    "prefix-probing": {
        "name": "IPv4 prefix probing",
        "maxsize": 1024 * 1024 * 400,
        # eug-us.20260702.1782961709.warts.gz
        "fileregex": re.compile(r"^(?P<name>[a-z][a-z0-9-]{0,31})\.(?P<date>[0-9]{8})\.(?P<timestamp>[0-9]{10})\.warts\.gz$"),
    }
}


# the nginx proxy sets all the headers, so we should be able to trust:
#   X-SSL-Client-Cert-Subject-DN
#   X-Forwarded-For
class UploadHandler(ArkHTTPRequestHandler):

    # XXX should this send the error, or the calling location?
    # XXX use urllib to get the last part of the url?
    def get_activity(self):
        # check that the activity being reported is valid
        activity = self.path.strip("/ ")
        if activity not in ACTIVITIES:
            self.send_error(HTTPStatus.BAD_REQUEST,
                    message=f"Unknown activity: {activity}")
            return None
        return activity


    # XXX is it better to use multipart/form-data and try to find the
    # content-disposition header inside of that, use content-disposition
    # with an octect stream, or just set a global x-file-name header?
    def get_filename(self):
        header = self.headers.get("Content-Disposition", None)
        if header is None:
            self.send_error(HTTPStatus.BAD_REQUEST,
                    message="Missing Content-Disposition header")
            return None

        filename = FILENAME_REGEX.search(header)
        if filename is None:
            self.send_error(HTTPStatus.BAD_REQUEST,
                    message=f"Invalid Content-Disposition header: {header}")
            return None

        # we don't want any path that might somehow be attached to the name
        return Path(filename.group(1)).name


    def validate_filename(self, filename, activity, name):
        # check the filename is properly formed for the activity type
        match = ACTIVITIES[activity]["fileregex"].fullmatch(filename)
        if match is None:
            self.send_error(HTTPStatus.BAD_REQUEST,
                    message=f"Invalid file name for {activity}")
            return False

        # check the name in the filename matches the name in the certificate
        if match["name"] != name:
            self.send_error(HTTPStatus.BAD_REQUEST,
                    message="Ark name doesn't match file name")
            return False

        # check the date and timestamp agree on what the date is
        when = datetime.fromtimestamp(int(match["timestamp"]), timezone.utc)
        if match["date"] != when.strftime("%Y%m%d"):
            self.send_error(HTTPStatus.BAD_REQUEST,
                    message="File timestamp and date don't match")
            return False

        # check the file isn't too old, or from the future
        if when < datetime.now(timezone.utc) - timedelta(days=14):
            self.send_error(HTTPStatus.BAD_REQUEST,
                    message="File timestamp is too old")
            return False
        if when > datetime.now(timezone.utc):
            self.send_error(HTTPStatus.BAD_REQUEST,
                    message="File timestamp is from the future")
            return False

        # it's probably properly formatted
        return True


    def report_download_stats(self, stats):
        if self.server.producer:
            self.server.producer.send(topic="ark-download-stat",
                    value=json.dumps(stats).encode("utf-8"))


    def report_download_complete(self, name, activity, filename):
        if self.server.producer:
            metadata = f"{name}|{activity}|{filename}"
            self.server.producer.send(topic="ark-downloaded-data",
                    value=metadata.encode("utf-8"))


    def do_POST(self):
        # "/upload/team-probing" gives "/team-probing" after nginx proxy
        # strips the matching /upload portion
        #print(self.path)

        # check that there is a useful client name in the certificate
        # e.g. CN=aee9-zz.ark.caida.org
        name = self.get_name_from_certificate()
        if name is None:
            return

        # XXX temporarily only allow uploads from a limited set of nodes
        if name not in ["hlz2-nz", "hlz3-nz"]:
            return

        # check that the activity being reported is valid
        activity = self.get_activity()
        if activity is None:
            return

        #self.log_message(f"{name}: {activity}")

        # check that the node exists in the database
        monitor = self.server.monitors.get(name)
        if monitor is None:
            self.send_error(HTTPStatus.FORBIDDEN,
                    message=f"Unknown node {name}")
            return

        # check that the named node has that activity enabled
        if ACTIVITIES[activity]["name"] not in monitor.get("activities", []):
            self.send_error(HTTPStatus.FORBIDDEN,
                    message=f"Activity {activity} is disabled for node {name}")
            return

        # check that the file isn't too big for the activity type
        try:
            length = int(self.headers["Content-Length"])
        except (KeyError, ValueError):
            self.send_error(HTTPStatus.LENGTH_REQUIRED,
                    message="Invalid Content-Length header")
            return

        if length > ACTIVITIES[activity]["maxsize"]:
            self.send_error(HTTPStatus.REQUEST_ENTITY_TOO_LARGE,
                    message="Payload too large")
            return

        # check there is a filename attached to the request
        filename = self.get_filename()
        if filename is None:
            return

        if not self.validate_filename(filename, activity, name):
            return

        # create the staging directory for this test and host if not present
        directory = Path(self.server.staging_directory, activity, name)
        try:
            Path.mkdir(directory, parents=True, exist_ok=True)
        except OSError as e:
            self.log_error(f"Failed to create directory: {e}")
            self.send_error(HTTPStatus.INTERNAL_SERVER_ERROR,
                    message="Failed to save file")
            return

        # determine the final location of the file, and a temporary version
        # that will be initially created. This way the final version only
        # exists if everything was successful
        filepath = Path(directory, filename)
        tmpfilepath = filepath.with_suffix(filepath.suffix + ".part")

        self.log_message(f"{name}: {filepath}")

        # check if the final file already exists, we won't try to replace it
        if filepath.is_file():
            self.send_error(HTTPStatus.CONFLICT, message="File already present")
            return

        # protect just the file upload portion with a semaphore to limit how
        # many concurrent uploads we allow. If there are too many at once,
        # send a 503 error with a Retry-After header
        if not self.server.upload_semaphore.acquire(timeout=5):
            self.log_message(f"{name}: too busy, sending 503 with retry-after")
            self.send_response(HTTPStatus.SERVICE_UNAVAILABLE)
            self.send_header("Retry-After", str(random.randint(20, 60)))
            self.end_headers()
            return

        remaining = length
        try:
            with tmpfilepath.open("xb") as outfile:
                while remaining > 0:
                    # the underlying ArkHTTPServer has a socket timeout,
                    # so this won't block forever if a node disconnects
                    data = self.rfile.read(min(CHUNK_SIZE, remaining))
                    if not data:
                        break
                    outfile.write(data)
                    remaining -= len(data)
        except FileExistsError:
            self.log_message(f"{name}: {filepath} already exists")
            self.send_error(HTTPStatus.CONFLICT, message="File already present")
            return
        except (ConnectionError, TimeoutError, ssl.SSLError) as e:
            self.log_error(f"{name}: connection error: {e}")
        except OSError as e:
            self.log_error(f"{name}: failed to save file: {e}")
            self.send_error(HTTPStatus.INTERNAL_SERVER_ERROR,
                    message="Failed to save file")
        finally:
            self.server.upload_semaphore.release()

        if remaining == 0:
            # file upload was successful, rename the file to the final name
            tmpfilepath.rename(filepath)
            self.send_response(HTTPStatus.NO_CONTENT)
            self.end_headers()
            # push statistics to kafka topic for downstream tools
            self.report_download_complete(name, activity, filename)
            self.report_download_stats({
                "timestamp": int(datetime.now(timezone.utc).timestamp()),
                "monitor": name,
                "activity": activity,
                "file": filename,
                "size": length,
            })
        else:
            # file upload failed, remove the temporary file
            try:
                tmpfilepath.unlink()
            except FileNotFoundError:
                pass
            except OSError as e:
                self.log_error(f"Error removing temporary file: {e}")

            # remote end has probably disconnected and is unlikely to get this
            self.send_response(HTTPStatus.SERVICE_UNAVAILABLE)
            self.send_header("Retry-After", str(random.randint(30, 90)))
            self.end_headers()

        return


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ark collector")
    parser.add_argument("-k", "--disable-kafka",
            dest="kafka",
            action='store_false',
            default=True,
            help="Don't send statistics to the local kafka server")
    parser.add_argument("-c", "--concurrent-transfers",
            dest="concurrent",
            default=MAX_UPLOADS,
            help="Maximum number of simultaneous file transfers")
    parser.add_argument("-d", "--staging-directory",
            dest="staging_directory",
            default=DEFAULT_STAGING_DIR,
            help="Directory to upload files to")
    args = parser.parse_args()

    server = ArkHTTPServer((HOSTNAME, PORT), UploadHandler)

    server.upload_semaphore = threading.BoundedSemaphore(args.concurrent)
    server.request_queue_size = args.concurrent * 2
    server.staging_directory = args.staging_directory
    server.producer = None

    if args.kafka:
        server.producer = KafkaProducer(
                bootstrap_servers="127.0.0.1:9092",
                client_id="ark-collector",
                security_protocol="SASL_PLAINTEXT",
                sasl_mechanism="SCRAM-SHA-256",
                sasl_plain_username="ark-collector",
                sasl_plain_password="ark-collector",
                #enable_idempotence=True,
                )

    print(f"Starting server on {HOSTNAME}:{PORT}")

    try:
        server.serve_forever()
    finally:
        print("Closing open connections...")
        server.server_close()
        if server.producer:
            print("Flushing kafka producer queue...")
            server.producer.flush(timeout=10)
            server.producer.close()
