akvo-top-asn

#!/usr/bin/env python3

import clickhouse_connect
from collections import namedtuple
import math
import numpy as np
import argparse
import os
import datetime
import configparser
import requests
import yaml
import logging
from typing import Tuple, List, Dict


def setup_logging(args):
    logger = logging.getLogger("akvo-top-asn")

    if args.debug:
        logger.setLevel(logging.DEBUG)
    elif args.quiet:
        logger.setLevel(logging.WARNING)
    else:
        logger.setLevel(logging.INFO)

    # create console handler and set level to debug
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)

    # create formatter
    formatter = logging.Formatter('%(levelname)s - %(message)s')

    # add formatter to ch
    ch.setFormatter(formatter)

    # add ch to logger
    logger.addHandler(ch)

    return logger


def check_conf_filename(fn):
    """
    Some simple (Race condition!) config filename checks to create
    polite error messages.
    """
    if not os.path.exists(fn):
        raise argparse.ArgumentTypeError(f"{fn} does not exist")
    if not os.path.isfile(fn):
        raise argparse.ArgumentTypeError(f"{fn} is not a file")

    return fn


def check_time_stamp(ts):
    return datetime.datetime.fromisoformat(ts)


# Instantiate the parser
parser = argparse.ArgumentParser(description="Akvorado top ASN stats")
parser.add_argument(
    "-c", "--config", type=check_conf_filename, default="akvo-top-asn.conf", help="configuration filename"
)
parser.add_argument("--starttime", type=check_time_stamp)
parser.add_argument("--endtime", type=check_time_stamp)

parser.add_argument("--days", type=int, default=0)
parser.add_argument("--weeks", type=int, default=1)
parser.add_argument("--months", type=int, default=0)

parser.add_argument("--filename", default=argparse.SUPPRESS, help="filename strftime pattern for saving and uploading results" )
group = parser.add_argument_group("export actions")
group.add_argument("--save", action="store_true", help="save result to local file")
group.add_argument("--upload", action="store_true", help="upload result to remote http url")
group.add_argument("--print", action="store_true", help="print result to stdout")

parser.add_argument("--query-limit", type=int, default=50)
parser.add_argument("--value-limit", type=int, default=1000)

# logging settings
logging_group = parser.add_mutually_exclusive_group()
logging_group.add_argument("--debug", action="store_true", help="enable debug logging")
logging_group.add_argument("--quiet", action="store_true", help="suppress informational logging")

args = parser.parse_args()
logger = setup_logging(args)

if not (args.print or args.upload or args.save):
  parser.error("You need to specify at least one action!")

config = configparser.ConfigParser(interpolation=None)
config.read(args.config)

starttime = None
endtime = None


# if no starttime is given we will take a timestamp thats in the past by some default offset
if args.starttime is None:
    delta = datetime.timedelta(
        days=args.days + args.weeks * 7 + args.months * 30,
    )

    starttime = datetime.datetime.utcnow() - delta
else:
    startime = args.startime

# if no endtime is given we take the current time as endtime
if args.endtime is None:
    endtime = datetime.datetime.utcnow()
else:
    endtime = args.endtime

time_range = [starttime.strftime("%Y-%m-%d %H:%M:%S"), endtime.strftime("%Y-%m-%d %H:%M:%S")]
si_prefixes = ["", "K", "M", "G"]
directions = ["Out", "In"]


def format_bps(n):
    """
    Format value with bps unit and SI prefix.
    """
    idx = max(0, min(len(si_prefixes) - 1, int(math.floor(0 if n == 0 else math.log10(abs(n)) / 3))))

    return "{:.0f}{}bps".format(n / 10 ** (3 * idx), si_prefixes[idx])


def get_connection(config):
    """
    Create a clickhouse client instance based on the config
    parameters of the [clickhouse] section.
    """
    params = {"host": config["clickhouse"].get("host", "localhost")}

    # add optional parameters
    for param_name, param_type in [
        ("secure", bool),
        ("port", int),
        ("username", str),
        ("password", str),
    ]:
        if param_name in config["clickhouse"]:
            try:
                params[param_name] = param_type(config["clickhouse"][param_name])
            except ValueError as ex:
                print(f"Bad config value for clickhouse.{param_name}: {ex}")
                exit(1)

    logger.info("connecting to clickhouse at %s", params["host"])
    return clickhouse_connect.get_client(**params)


def query_clickhouse(time_range: Tuple[str, str], direction: str, offset:int = 0, limit: int = 50):
    query = f"""
        WITH
        source AS (SELECT * FROM flows_5m0s SETTINGS asterisk_include_alias_columns = 1),
        rows AS (SELECT SrcAS FROM source WHERE TimeReceived BETWEEN toDateTime('{time_range[0]}', 'UTC') AND toDateTime('{time_range[1]}', 'UTC') AND (InIfBoundary = 'external') AND WHERE SUM(Bytes) >= GROUP BY SrcAS ORDER BY SUM(Bytes) DESC LIMIT {limit} OFFSET {offset})
        SELECT 1 AS axis, * FROM (
        SELECT
        toStartOfInterval(TimeReceived + INTERVAL 900 second, INTERVAL 900 second) - INTERVAL 900 second AS time,
        SUM(Bytes*SamplingRate*8)/900 AS xps,
        if((SrcAS) IN rows, [concat(toString(SrcAS), ': ', dictGetOrDefault('asns', 'name', SrcAS, '???'))], ['0: Other']) AS dimensions
        FROM source
        WHERE TimeReceived BETWEEN toDateTime('{time_range[0]}', 'UTC') AND toDateTime('{time_range[1]}', 'UTC') AND ({direction}IfBoundary = 'external')
        GROUP BY time, dimensions
        ORDER BY time WITH FILL
        FROM toDateTime('{time_range[0]}', 'UTC')
        TO toDateTime('{time_range[1]}', 'UTC') + INTERVAL 1 second
        STEP 900
        INTERPOLATE (dimensions AS ['0: Other']))
        """

    logger.debug("clickhouse query: %s", query)
    return client.query(query).result_rows


# connect to clickhouse
try:
    client = get_connection(config)
except clickhouse_connect.driver.exceptions.DatabaseError as ex:
    logger.error("Failed to connect to clickhouse: %s", ex)
    exit(2)

logger.info("query time range: %s - %s", *time_range)


value_limit = config["query"].get("value_limit", 1000)
query_limit = config["query"].get("query_limit", 100)

# build list of bandwidths per ASN
asn_xps = {}
for direction in directions: # in bound and outbound
    i = 0
    while True:
        logger.info(f"query results from {i * query_limit} until {(i + 1) * query_limit} ... ")

        query_result = query_clickhouse(time_range, direction, limit=query_limit, offset= i * query_limit)
        for axis, ts, xps, _asn in query_result:
            asn = _asn[0]
            if asn in asn_xps:
                if direction in asn_xps[asn]:
                    asn_xps[asn][direction].append(xps)
                else:
                    asn_xps[asn][direction] = [xps]
            else:
                asn_xps[asn] = {direction: [xps]}

        if query_result < query_limit and (i * query_limit) < value_limit:
            i += 1
        else:
            break
            

logger.info("building statistics...")


# build stats for each ASN
asn_stats = {}
for _asn, xps in asn_xps.items():
    asn, org = map(str.strip, _asn.split(':', 1))
    asn = int(asn)

    # skip 'Other'
    if asn == 0:
        continue

    stats = {
        'org': org,
    }
    for direction in directions:
        d = direction.lower()

        if direction in xps:
            arr = np.array(xps[direction])

            stats[f"{d}_avg"] = np.mean(arr)
            stats[f"{d}_p95"] = np.percentile(arr, 95)
            stats[f"{d}_max"] = np.max(arr)
        else:
            stats[f"{d}_avg"] = 0
            stats[f"{d}_p95"] = 0
            stats[f"{d}_max"] = 0

    asn_stats[asn] = stats

if args.save or args.upload:
    fn = starttime.strftime(getattr(args, 'filename', config['upload'].get('filename', 'export-%Y-W%W.yml')))

    # dump result to string
    yaml.add_representer(np.float64, lambda dumper, data: dumper.represent_scalar('tag:yaml.org,2002:float', str(data)))
    result = yaml.dump({
        'from': time_range[0],
        'to': time_range[1],
        # 'local_asn': local_asn,
        'top_peers': asn_stats,
    })

    # save result to local file
    if args.save:
        logger.info("saving to file: %s", fn)
        with open(fn, 'w') as fh:
            fh.write(result)

    # upload result to http remote location
    if args.upload:
        params = {
            "headers": {
                "X-Requested-With": "XMLHttpRequest",
            },
            "data": result,
        }

        if "username" in config["upload"] or "password" in config["upload"]:
            params["auth"] = (config["upload"].get("username"), config["upload"].get("password"))

        url = config["upload"]["url"] + "/" + fn
        logger.info("uploading to: %s", url)
        response = requests.put(url, **params)
        if response.status_code in range(200, 299):
            logger.info("upload succeeded")
        else:
            logging.error("upload failed: %d - %s", response.status_code, response.text)


# print to stdout
if args.print:
    # poor-mans table headers
    print("{asn: <7} {org: <36}".format(asn="ASN", org="ORG"), end="")
    columns = []
    for direction in directions:
        for metric in ["avg", "p95", "max"]:
            columns.append(f"{{{direction.lower()}_{metric}: >9}}")
            print(" {metric: <9}".format(metric=f"{direction.lower()}_{metric}"), end="")

    # sort & print ASN stats
    for asn, stats in sorted(asn_stats.items(), key=lambda i: i[1]["out_p95"] + i[1]["in_p95"], reverse=True):
        print(f'{{asn: <7}} {{org: <36}} {" ".join(columns)}'.format(asn=asn, org=stats['org'], **{k: format_bps(v) for k, v in stats.items() if k != "org"}))