#!/usr/bin/python3
"""Collect the metrics of the router, keep them in PostgreSQL, serve them for scraping and push them with OTLP.

Run by vpn-router-metrics.service while metrics.enabled is true; vpn-router-setup enables and
disables the unit.

Every SAMPLE_SECONDS, on the minute, a sample is taken with psutil: the CPU utilization since the
previous sample, the memory used and in total, the free and the total space of /, the bytes
received and sent per second on the external and the internal interface, and the load averages;
and the VPN clients of each enabled VPN system (vpnrouter_status): IKEv2 (p2s) and
OpenVPN clients and WireGuard peers with a recent handshake; and of each site-to-site connection
that is set up, whether it is up, the bytes and packets per second through its child SAs since
the previous sample, and how long its IKE SA has been up. The interfaces, the enabled systems and
the connections are those vpn-router-setup resolved, read from /run/vpn-router/web.json.

The sample is written to the detail tables of vpnrouter_metrics. When the clock has passed into
a new hour, day, month or year, the complete periods are summarized, and once a day, after the
daily summary, the rows older than their retention are deleted; both also run on start. A
failure of the database is logged and tried again with the next sample.

GET /metrics answers the last sample in the Prometheus text format, written by
prometheus_client, with OpenTelemetry semantic convention names in their Prometheus form, for an
OpenTelemetry Collector's prometheus receiver. It needs HTTP Basic authentication as the service
account metrics.service_account names, which vpn-router-setup keeps with the password in
/etc/vpn-router/metrics.secret, checked with a bind to the router's directory, so failed attempts
lock the account as for any user. The service listens on the interfaces metrics.listen
names and on the loopback address, where Caddy serves it over HTTPS at metrics.https_port; the
address of a call through Caddy is the last entry of its X-Forwarded-For header.

When metrics.otlp_endpoint is set, every sample is also pushed there as an OTLP/HTTP JSON
ExportMetricsServiceRequest, with HTTP Basic authentication when metrics.otlp_username is set.
"""

import base64
import datetime
import json
import os
import signal
import socket
import sys
import threading
import time
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler

sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)))

import ldap
import psutil
from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
from prometheus_client.core import CollectorRegistry, CounterMetricFamily, GaugeMetricFamily

import vpnrouter
import vpnrouter_metrics
import vpnrouter_net
import vpnrouter_serve
import vpnrouter_status
import vpnrouter_users

STATE_FILE = '/run/vpn-router/web.json'
CONFIG_FILE = '/etc/vpn-router/vpn-router.json'

SAMPLE_SECONDS = 60

#: Seconds between two discoveries of where the service should listen.
DISCOVERY_SECONDS = 30

#: Seconds an OTLP push may take.
OTLP_TIMEOUT = 10

#: The mount point whose space is measured.
MOUNTPOINT = '/'

#: The last sample, as sample() makes it, or None before the first.
_latest = {'sample': None}
_latest_lock = threading.Lock()


def log(message):
    print(f'vpn-router-metrics: {message}', file=sys.stderr, flush=True)


def read_json(path):
    """The object in a JSON file, or {} when the file is missing or holds no object."""
    try:
        with open(path, encoding='utf-8') as handle:
            data = json.load(handle)
    except (OSError, ValueError):
        return {}
    return data if isinstance(data, dict) else {}


def section(configuration, name):
    """A section of the configuration, or {} when it is missing or not an object."""
    value = configuration.get(name)
    return value if isinstance(value, dict) else {}


def clients(state):
    """The VPN clients of each enabled VPN system, None where it cannot be asked."""
    counts = {}
    if state.get('p2s'):
        counts['p2s'] = vpnrouter_status.established('road-warrior')
    if state.get('openvpn'):
        counts['openvpn'] = vpnrouter_status.openvpn_clients()
    if state.get('wireguard'):
        counts['wireguard'] = vpnrouter_status.wireguard_connected()
    return counts


def sample(previous):
    """A sample of the router, with the rates since the previous sample, None for the first."""
    state = read_json(STATE_FILE)
    network = section(state, 'network')
    names = [name for name in dict.fromkeys((network.get('external_interface'), network.get('internal_interface')))
             if isinstance(name, str) and name]
    counters = psutil.net_io_counters(pernic=True)
    memory = psutil.virtual_memory()
    disk = psutil.disk_usage(MOUNTPOINT)
    load1, load5, load15 = os.getloadavg()
    connections = state.get('site_to_site')
    tunnels = vpnrouter_status.site_to_site_tunnels(connections) if isinstance(connections, list) and connections else {}
    now = time.time()
    current = {
        'time': now,
        'monotonic': time.monotonic(),
        'system': {
            'cpu': psutil.cpu_percent(interval=None),
            'memory_used': memory.total - memory.available,
            'memory_total': memory.total,
            'disk_free': disk.free,
            'disk_total': disk.total,
            'load1': load1,
            'load5': load5,
            'load15': load15,
        },
        'counters': {name: (counters[name].bytes_recv, counters[name].bytes_sent) for name in names if name in counters},
        'clients': clients(state),
        'rates': {},
        'tunnels': tunnels or {},
        'site_to_site': {},
    }
    elapsed = current['monotonic'] - previous['monotonic'] if previous is not None else 0
    if elapsed > 0:
        for name, (received, sent) in current['counters'].items():
            if name in previous['counters']:
                before_received, before_sent = previous['counters'][name]
                current['rates'][name] = (max(received - before_received, 0) / elapsed, max(sent - before_sent, 0) / elapsed)
    # A connection's rates need its counters of the previous sample; a counter that went down was rekeyed and counts from zero.
    for name, tunnel in current['tunnels'].items():
        before = previous['tunnels'].get(name) if elapsed > 0 else None
        values = {'up': tunnel['up'], 'uptime': tunnel['uptime']}
        for value, counter in (('rx', 'bytes_in'), ('tx', 'bytes_out'), ('packets_rx', 'packets_in'), ('packets_tx', 'packets_out')):
            if before is None:
                values[value] = None
            else:
                count = tunnel[counter]
                values[value] = (count - before[counter] if count >= before[counter] else count) / elapsed
        current['site_to_site'][name] = values
    return current


def periods(moment):
    """The hour, the day, the month and the year of a UTC moment, keyed by the levels of vpnrouter_metrics.LEVELS."""
    return {'hourly': moment.strftime('%Y-%m-%dT%H'), 'daily': moment.strftime('%Y-%m-%d'),
            'monthly': moment.strftime('%Y-%m'), 'yearly': moment.strftime('%Y')}


def store(cfg, current, done):
    """Write the sample, summarize the periods the clock has passed and collect the garbage once a day.

    done holds the periods summarized so far, and garbage the day of the last garbage collection;
    it is updated only for what succeeded.
    """
    now = periods(datetime.datetime.now(datetime.timezone.utc))
    with vpnrouter_metrics.database(cfg) as connection:
        vpnrouter_metrics.write(connection, current['system'], current['rates'], current['clients'], current['site_to_site'])
        for level, _, _ in vpnrouter_metrics.LEVELS:
            if done.get(level) != now[level]:
                vpnrouter_metrics.summarize(connection, level)
                done[level] = now[level]
        if done.get('garbage') != now['daily']:
            vpnrouter_metrics.collect_garbage(connection, cfg)
            done['garbage'] = now['daily']


def attributes(**values):
    """OTLP attributes of string values."""
    return [{'key': key, 'value': {'stringValue': value}} for key, value in values.items()]


def otlp_request(current, started):
    """The OTLP/HTTP JSON ExportMetricsServiceRequest of a sample; started is when the counters were first read, in Unix seconds."""
    stamp = str(int(current['time'] * 1e9))
    start = str(int(started * 1e9))
    system = current['system']

    def gauge(name, unit, points):
        return {'name': name, 'unit': unit, 'gauge': {'dataPoints': [
            {'timeUnixNano': stamp, 'asDouble': float(value), 'attributes': attributes(**labels)}
            for value, labels in points if value is not None]}}

    metrics = [
        gauge('system.cpu.utilization', '1', [(system['cpu'] / 100, {})]),
        gauge('system.memory.usage', 'By', [(system['memory_used'], {'state': 'used'})]),
        gauge('system.memory.limit', 'By', [(system['memory_total'], {})]),
        gauge('system.filesystem.usage', 'By', [(system['disk_free'], {'mountpoint': MOUNTPOINT, 'state': 'free'}),
                                               (system['disk_total'] - system['disk_free'], {'mountpoint': MOUNTPOINT, 'state': 'used'})]),
        gauge('system.filesystem.utilization', '1',
              [(1 - system['disk_free'] / system['disk_total'] if system['disk_total'] else None, {'mountpoint': MOUNTPOINT})]),
        gauge('system.cpu.load_average.1m', '{thread}', [(system['load1'], {})]),
        gauge('system.cpu.load_average.5m', '{thread}', [(system['load5'], {})]),
        gauge('system.cpu.load_average.15m', '{thread}', [(system['load15'], {})]),
        gauge('vpn_router.clients', '{client}', [(count, {'kind': kind}) for kind, count in current['clients'].items()]),
        gauge('vpn_router.site_to_site.up', '1', [(values['up'], {'connection': name})
                                                 for name, values in current['site_to_site'].items()]),
        gauge('vpn_router.site_to_site.uptime', 's', [(values['uptime'], {'connection': name})
                                                     for name, values in current['site_to_site'].items()]),
        gauge('vpn_router.site_to_site.io', 'By/s', [(values[value], {'connection': name, 'direction': direction})
                                                    for name, values in current['site_to_site'].items()
                                                    for value, direction in (('rx', 'receive'), ('tx', 'transmit'))]),
        gauge('vpn_router.site_to_site.packets', '{packet}/s',
              [(values[value], {'connection': name, 'direction': direction})
               for name, values in current['site_to_site'].items()
               for value, direction in (('packets_rx', 'receive'), ('packets_tx', 'transmit'))]),
        {'name': 'system.network.io', 'unit': 'By', 'sum': {
            'aggregationTemporality': 2, 'isMonotonic': True, 'dataPoints': [
                {'startTimeUnixNano': start, 'timeUnixNano': stamp, 'asInt': str(value),
                 'attributes': attributes(device=name, direction=direction)}
                for name, counters in current['counters'].items()
                for value, direction in zip(counters, ('receive', 'transmit'))]}},
    ]
    return {'resourceMetrics': [{
        'resource': {'attributes': attributes(**{'service.name': 'vpn-router', 'host.name': socket.gethostname()})},
        'scopeMetrics': [{'scope': {'name': 'vpn-router-metrics'}, 'metrics': metrics}],
    }]}


def push(metrics, current, started):
    """Push the sample to metrics.otlp_endpoint; raises OSError when the collector does not take it."""
    request = urllib.request.Request(metrics['otlp_endpoint'], data=json.dumps(otlp_request(current, started)).encode(),
                                     headers={'Content-Type': 'application/json'}, method='POST')
    if metrics['otlp_username']:
        credentials = f"{metrics['otlp_username']}:{metrics['otlp_password']}".encode()
        request.add_header('Authorization', 'Basic ' + base64.b64encode(credentials).decode())
    with urllib.request.urlopen(request, timeout=OTLP_TIMEOUT) as response:
        response.read()


def collector_loop():
    """Take a sample on every minute, store it and push it; every new failure is logged once."""
    previous, done, reported = None, {}, {}
    started = time.time()
    psutil.cpu_percent(interval=None)

    def report(kind, message):
        if message != reported.get(kind):
            if message:
                log(message)
            elif reported.get(kind):
                log(f'{kind} works again')
            reported[kind] = message

    while True:
        time.sleep(SAMPLE_SECONDS - time.time() % SAMPLE_SECONDS)
        try:
            cfg = vpnrouter.load_config()
        except SystemExit as exc:
            report('configuration', str(exc.code))
            continue
        report('configuration', '')
        current = sample(previous)
        previous = current
        with _latest_lock:
            _latest['sample'] = current
        try:
            store(cfg, current, done)
            report('database', '')
        except vpnrouter_metrics.DatabaseError as exc:
            report('database', str(exc))
        metrics = cfg['metrics']
        if metrics['otlp_endpoint']:
            try:
                push(metrics, current, started)
                report('OTLP push', '')
            except (OSError, ValueError) as exc:
                report('OTLP push', f"pushing to {metrics['otlp_endpoint']} failed: {exc}")


class Sample:
    """A prometheus_client collector of a sample."""

    def __init__(self, current):
        self.current = current

    def collect(self):
        system = self.current['system']
        yield GaugeMetricFamily('system_cpu_utilization', 'CPU utilization since the previous sample, 0 to 1',
                                value=system['cpu'] / 100)
        memory = GaugeMetricFamily('system_memory_usage_bytes', 'Memory in use', labels=['state'])
        memory.add_metric(['used'], system['memory_used'])
        yield memory
        yield GaugeMetricFamily('system_memory_limit_bytes', 'Total memory', value=system['memory_total'])
        usage = GaugeMetricFamily('system_filesystem_usage_bytes', 'Space of the file system', labels=['mountpoint', 'state'])
        usage.add_metric([MOUNTPOINT, 'free'], system['disk_free'])
        usage.add_metric([MOUNTPOINT, 'used'], system['disk_total'] - system['disk_free'])
        yield usage
        if system['disk_total']:
            utilization = GaugeMetricFamily('system_filesystem_utilization', 'Used part of the file system, 0 to 1',
                                            labels=['mountpoint'])
            utilization.add_metric([MOUNTPOINT], 1 - system['disk_free'] / system['disk_total'])
            yield utilization
        for period in ('1m', '5m', '15m'):
            yield GaugeMetricFamily(f'system_cpu_load_average_{period}', f'Load average over {period}',
                                    value=system[f'load{period[:-1]}'])
        network = CounterMetricFamily('system_network_io_bytes', 'Bytes received and sent', labels=['device', 'direction'])
        for name, (received, sent) in self.current['counters'].items():
            network.add_metric([name, 'receive'], received)
            network.add_metric([name, 'transmit'], sent)
        yield network
        vpn = GaugeMetricFamily('vpn_router_clients', 'Connected VPN clients', labels=['kind'])
        for kind, count in self.current['clients'].items():
            if count is not None:
                vpn.add_metric([kind], count)
        yield vpn
        tunnels = self.current['site_to_site']
        up = GaugeMetricFamily('vpn_router_site_to_site_up', 'Whether the site-to-site connection is established, 1 or 0',
                               labels=['connection'])
        uptime = GaugeMetricFamily('vpn_router_site_to_site_uptime_seconds', 'Seconds the IKE SA of the site-to-site connection has been up',
                                   labels=['connection'])
        io = GaugeMetricFamily('vpn_router_site_to_site_io_bytes_per_second',
                               'Bytes per second received and sent through the site-to-site connection', labels=['connection', 'direction'])
        packets = GaugeMetricFamily('vpn_router_site_to_site_packets_per_second',
                                    'Packets per second received and sent through the site-to-site connection',
                                    labels=['connection', 'direction'])
        for name, values in tunnels.items():
            up.add_metric([name], values['up'])
            if values['uptime'] is not None:
                uptime.add_metric([name], values['uptime'])
            for family, received, sent in ((io, 'rx', 'tx'), (packets, 'packets_rx', 'packets_tx')):
                for value, direction in ((received, 'receive'), (sent, 'transmit')):
                    if values[value] is not None:
                        family.add_metric([name, direction], values[value])
        yield from (up, uptime, io, packets)


def authorized(header, account):
    """Whether an Authorization header holds the password of the service account; LDAP errors propagate."""
    if not header.startswith('Basic '):
        return False
    try:
        name, _, password = base64.b64decode(header[6:].strip(), validate=True).decode().partition(':')
    except ValueError:
        return False
    if name != account:
        return False
    try:
        with vpnrouter_users.session() as connection:
            if not vpnrouter_users.service(vpnrouter_users.entry(connection, name)):
                return False
        vpnrouter_users.login(name, password, None)
    except vpnrouter_users.Failure:
        return False
    return True


class Handler(BaseHTTPRequestHandler):
    protocol_version = 'HTTP/1.1'

    def client(self):
        """The client address: the last X-Forwarded-For entry of a call through Caddy on the loopback address, else the peer."""
        peer = self.client_address[0]
        forwarded = self.headers.get('X-Forwarded-For', '')
        if peer == vpnrouter_net.LOOPBACK and forwarded.strip():
            return forwarded.split(',')[-1].strip()
        return peer

    def answer(self, send_body):
        self.principal = None
        extra = []
        path = self.path.split('?', 1)[0]
        account = str(section(read_json(CONFIG_FILE), 'metrics').get('service_account',
                                                                       vpnrouter.SCHEMA['metrics']['service_account'])).strip()
        with _latest_lock:
            current = _latest['sample']
        if path != '/metrics':
            body, kind, code = b'Not found\n', 'text/plain; charset=utf-8', 404
        elif not account:
            body, kind, code = b'metrics.service_account is not set\n', 'text/plain; charset=utf-8', 503
        else:
            try:
                allowed = authorized(self.headers.get('Authorization', ''), account)
            except (ldap.LDAPError, OSError) as exc:
                allowed = None
                body, kind, code = f'the user directory failed: {exc}\n'.encode(), 'text/plain; charset=utf-8', 500
            if allowed is False:
                body, kind, code = b'Unauthorized\n', 'text/plain; charset=utf-8', 401
                extra.append(('WWW-Authenticate', 'Basic realm="vpn-router metrics", charset="UTF-8"'))
            elif allowed:
                self.principal = account
                if current is None:
                    body, kind, code = b'no sample yet\n', 'text/plain; charset=utf-8', 503
                else:
                    registry = CollectorRegistry()
                    registry.register(Sample(current))
                    body, kind, code = generate_latest(registry), CONTENT_TYPE_LATEST, 200
        self.send_response(code)
        self.send_header('Content-Type', kind)
        self.send_header('Content-Length', str(len(body)))
        self.send_header('Cache-Control', 'no-store')
        for name, value in extra:
            self.send_header(name, value)
        self.end_headers()
        if send_body:
            self.wfile.write(body)

    def do_GET(self):
        self.answer(True)

    def do_HEAD(self):
        self.answer(False)

    def log_message(self, format, *args):
        principal = getattr(self, 'principal', None) or '-'
        sys.stderr.write('%s - %s [%s] %s\n' % (self.client(), principal, self.log_date_time_string(), format % args))


def wanted():
    """Where the service listens: the addresses of the interfaces metrics.listen names and LOOPBACK, at metrics.port."""
    configuration = read_json(CONFIG_FILE)
    metrics = section(configuration, 'metrics')
    defaults = vpnrouter.SCHEMA['metrics']
    if metrics.get('enabled', defaults['enabled']) is False:
        return [], 'metrics are disabled in the configuration'
    port = metrics.get('port', defaults['port'])
    if isinstance(port, bool) or not isinstance(port, int) or not 1 <= port <= 65535:
        port = defaults['port']
    addresses = vpnrouter_net.listen_addresses(configuration, metrics.get('listen', defaults['listen']))
    if not addresses:
        return [], 'the addresses to listen on are not known yet'
    return [(address, port) for address in [*addresses, vpnrouter_net.LOOPBACK]], ''


def main():
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
    threading.Thread(target=collector_loop, name='collector', daemon=True).start()
    try:
        vpnrouter_serve.serve(wanted, Handler, log, DISCOVERY_SECONDS)
    except KeyboardInterrupt:
        pass
    return 0


if __name__ == '__main__':
    sys.exit(main())
