#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: MirOS OR LGPL-2.0-or-later
#
# Copyright © Thorsten Glaser <tglaser@b1-systems.de>
#
# Install to the monitored hosts as:
# /usr/lib/check_mk_agent/plugins/kea_dhcpd4.py
#
# check_mk_agent must run as root, or at least with sufficient
# privilegue to access SOCKFN and test for CONFIG existence.
#
# The "libdhcp_subnet_cmds.so" hook library must be loaded into
# kea-dhcp4-server for this agent to work. This library has only
# become part of ISC KEA core under FOSS licences with KEA 3.x.
#
# See corresponding check plugin for data documentation.

import json
import os
import re
import socket

# default buffer size for socket I/O, should not need adjusting
BUFSIZ = 8192

# default path for KEA socket, make sure it matches the config
SOCKFN = '/run/kea/kea4-ctrl-socket'

# configuration file path; this agent plugin is inert if absent
CONFIG = '/etc/kea/kea-dhcp4.conf'

# workaround for py3k JSON implementation
def _parse_constant(s):
    raise ValueError(f'Invalid JSON: "{s}"')

# an exception for the _m() function, see below
class MistypeException(Exception):
    def __init__(self, msg, *args, **kwargs):
        self.msg = msg
        super().__init__(msg, *args, **kwargs)

# extract member from dict and verify presence and type
def _m(obj, name, cls, path=None):
    if type(cls) != type:
        cls = type(cls)
    if path is None:
        fullname = name
    else:
        fullname = '%s.%s' % (path, name)
    if name not in obj:
        raise MistypeException('missing %s' % fullname)
    memb = obj[name]
    if isinstance(memb, cls):
        return memb
    if cls == dict:
        t = 'object'
    elif cls == list:
        t = 'array'
    elif cls == str:
        t = 'string'
    elif cls == int:
        t = 'integer'
    elif cls == float:
        t = 'numeric'
    elif cls == bool:
        t = 'boolean'
    elif cls == type(None):
        t = 'null'
    else:
        t = repr(cls)
    raise MistypeException('non-%s %s' % (t, fullname))

# on KEA communication errors
class KEAClientException(Exception):
    def __init__(self, msg, *args, **kwargs):
        self.msg = msg
        super().__init__(msg, *args, **kwargs)

    def __str__(self):
        r = super().__str__()
        if self.__cause__ is not None:
            r += ' ← ' + repr(self.__cause__)
        return r

    def __repr__(self):
        r = super().__repr__()
        if self.__cause__ is not None:
            r += ' ← ' + repr(self.__cause__)
        return r.expandtabs()

# general response codes from its documentation
_kea_codes = {
    0: 'the command has been processed successfully',
    1: 'a general error or failure has occurred during the command processing',
    2: 'the specified command is unsupported by the server receiving it',
    3: 'the requested operation has been completed but the requested resource was not found',
    4: 'the well-formed command has been processed but the requested changes could not be applied, because they were in conflict with the server state or its notion of the configuration',
}

# does a KEA API request and returns the response or its .arguments subkey
def _kea_request(cmd, args=None, rvok=None, need_args=True):
    if not isinstance(cmd, str):
        raise KEAClientException('invalid command (not a string): %s' % repr(cmd))
    if args is None:
        args = {}
    if rvok is None:
        rvok = [0]  # default “ok” response code(s)

    try:
        cmdstr = json.dumps({
            'command': cmd,
            'arguments': args,
          }, ensure_ascii=True, allow_nan=False, indent=None,
          separators=(",", ":"), sort_keys=True)
    except Exception as ex:
        raise KEAClientException('cannot format JSON request') from ex

    if not os.path.exists(SOCKFN):
        raise KEAClientException('KEA control socket does not exist')
    rsp = b''
    phase = "opening"
    try:
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
            phase = "connecting"
            sock.connect(SOCKFN)
            phase = "writing"
            sock.sendall(cmdstr.encode("ASCII"))
            phase = "reading"
            while True:
                rspnew = sock.recv(BUFSIZ)
                if len(rspnew) == 0:
                    break
                rsp += rspnew
            phase = "closing"
    except OSError as ex:
        raise KEAClientException('error %s KEA control socket for %s' % \
          (phase, repr(cmd))) from ex

    # 15 is the length of the minimum response {"response":0} as KEA formats
    if len(rsp) < 15:
        raise KEAClientException('unrealistically short response for %s' % \
          repr(cmd))
    try:
        rsp = json.loads(rsp, parse_constant=_parse_constant)
    except ValueError as ex:
        raise KEAClientException('error parsing JSON response for %s' % \
          repr(cmd)) from ex
    if not isinstance(rsp, dict):
        raise KEAClientException('bogus JSON response (not an object) for %s' % \
          repr(cmd))
    try:
        res = _m(rsp, 'result', int)
    except MistypeException as ex:
        raise KEAClientException('bogus JSON response (%s) for %s' % \
          (ex, repr(cmd)))

    # is the exit status listed as “ok return value”?
    if res in rvok:
        if not need_args:
            return rsp
        # caller requests the .arguments object only
        try:
            args = _m(rsp, 'arguments', dict)
        except MistypeException as ex:
            raise KEAClientException('bogus JSON response (%s) for %s: %s' % \
              (ex, repr(cmd), repr(rsp)))
        return args
    # the return value isn’t “okay”, determine error message
    if res in _kea_codes:
        raise KEAClientException('error response (code %d %s) for %s: %s' % \
          (res, repr(_kea_codes[res]), repr(cmd), repr(rsp)))
    raise KEAClientException('error response (code %d) for %s: %s' % \
      (res, repr(cmd), repr(rsp)))

# basically a list of lists of strings, but stringifies and removes tab/nl on add
class ResultList(object):
    def __init__(self):
        self.arr = []

    def row(self, *args):
        self.arr.append([str(col).expandtabs().replace('\n', '\\n') \
          for col in args])
        return self

    def out(self):
        s = '\n'.join(['\t'.join(row) for row in self.arr])
        if s:
            print(s)

# gather general KEA info (about the dæmon itself)
def _get_general():
    l = ResultList()
    try:
        r = _kea_request('status-get')
    except KEAClientException as ex:
        l.row('g', 'err', repr(ex))
        return l

    try:
        pid = _m(r, 'pid', int)
        l.row('g', 'pid', pid)
    except MistypeException as ex:
        l.row('g', 'err', str(ex))

    try:
        uptime = _m(r, 'uptime', int)
        l.row('g', 'uptime', uptime)
    except MistypeException as ex:
        l.row('g', 'err', str(ex))

    try:
        is_disabled = []
        ds = _m(r, 'dhcp-state', dict)
        for k, v in ds.items():
            if v:
                is_disabled.append(k)
        if is_disabled:
            l.row('g', 'err', 'DHCP state is disabled: %s' % ','.join(is_disabled))
    except MistypeException as ex:
        pass

    try:
        ss = _m(r, 'sockets', dict)
    except MistypeException as ex:
        l.row('g', 'err', str(ex))
        return l
    try:
        st = _m(ss, 'status', str, 'sockets')
    except MistypeException as ex:
        l.row('g', 'err', str(ex))
        st = None
    if 'errors' in ss:
        se = ', errors=' + repr(ss['errors'])
    else:
        se = ''
    l.row('g', 'socket', repr(st) + se)

    return l

# from 'statistic-get-all', we are interested in only these stats
_stats_re = re.compile('^subnet\\[([0-9]+)\\]\\.pool\\[([0-9])+\\]\\.(assigned|total)-addresses$')

# gather KEA statistics for pools
def _get_stats():
    l = ResultList()
    try:
        r = _kea_request('statistic-get-all')
    except KEAClientException as ex:
        l.row('g', 'err', repr(ex))
        # avoid extra error about missing stats
        l.row('g', '/s')
        return l

    pools = set()
    stats = {'assigned':{},'total':{}}
    for k, v in r.items():
        g = _stats_re.fullmatch(k)
        if not g:
            continue
        subnet = int(g.group(1))
        pool = int(g.group(2))
        type = g.group(3)
        key = (subnet, pool)
        if not isinstance(v, list):
            l.row('s', subnet, pool, 'err', \
              'wrong value type for stat %s' % k)
            continue
        if len(v) == 0:
            # no statistics yet
            continue
        if (not isinstance(v[0], list)) or (len(v[0]) == 0):
            l.row('s', subnet, pool, 'err', \
              'wrong encompassing type for latest stat %s' % k)
            continue
        if not isinstance(v[0][0], int):
            l.row('s', subnet, pool, 'err', \
              'wrong value type for latest stat %s' % k)
            continue
        value = v[0][0]
        stats[type][key] = value
        pools.add(key)

    for k in pools:
        if (k not in stats['assigned']) or (k not in stats['total']):
            l.row('s', subnet, pool, 'err', 'only partial stats received')
            continue
        l.row('s', subnet, pool, 'ok', stats['assigned'][k], stats['total'][k])

    l.row('g', '/s')
    return l

# gather pool configuration (specifically, which ranges they have)
def _get_pools():
    l = ResultList()
    try:
        r = _kea_request('subnet4-list')
    except KEAClientException as ex:
        l.row('g', 'err', repr(ex))
        # avoid extra error about missing pools
        l.row('g', '/p')
        return l

    try:
        snl = _m(r, 'subnets', list)
    except MistypeException as ex:
        l.row('g', 'err', str(ex))
        # avoid extra error about missing pools
        l.row('g', '/p')
        return l

    subnets = set()
    for sne in snl:
        if not isinstance(sne, dict):
            l.row('g', 'err', 'subnet4-list subentry is not an object')
            continue
        try:
            snid = _m(sne, 'id', int, 'subnet4-list[]')
        except MistypeException as ex:
            l.row('g', 'err', str(ex))
            continue
        subnets.add(snid)

    for sn in subnets:
        try:
            r = _kea_request('subnet4-get', {'id': sn})
        except KEAClientException as ex:
            l.row('p', sn, '*', 'err', repr(ex))
            continue
        try:
            snl = _m(r, 'subnet4', list)
        except MistypeException as ex:
            l.row('p', sn, '*', 'err', str(ex))
            continue
        snll = len(snl)
        if snll != 1:
            l.row('p', sn, '*', 'err', \
              '%d subnet4-get results, not just one' % snll)
            if snll == 0:
                continue
        snl = snl[0]
        if not isinstance(snl, dict):
            l.row('p', sn, '*', 'err', \
              'subnet4-get subentry is not an object')
            continue
        try:
            snid = _m(snl, 'id', int, 'subnet4[%d]' % sn)
        except MistypeException as ex:
            l.row('p', sn, '*', 'err', str(ex))
            continue
        if snid != sn:
            l.row('p', sn, '*', 'err', 'subnet4[%d] == %d' % (sn, snid))
            continue
        if 'pools' not in snl:
            continue
        try:
            snp = _m(snl, 'pools', list, 'subnet4[%d]' % sn)
        except MistypeException as ex:
            l.row('p', sn, '*', 'err', str(ex))
            continue
        for pid, po in enumerate(snp):
            if not isinstance(po, dict):
                l.row('p', sn, pid, 'err', 'pool info is not an object')
                continue
            try:
                pp = _m(po, 'pool', str, 'subnet4[%d].pools[%d]' % (sn, pid))
            except MistypeException as ex:
                l.row('p', sn, pid, 'err', str(ex))
                continue
            l.row('p', sn, pid, 'ok', pp)

    l.row('g', '/p')
    return l

# acquire and output lines
def _get_all():
    try:
        general = _get_general()
    except Exception as ex:
        general = ResultList().row('g', 'err', 'general: unexpected ' + repr(ex))
    general.out()

    try:
        stats = _get_stats()
    except Exception as ex:
        stats = ResultList().row('g', 'err', 'stats: unexpected ' + repr(ex))
    stats.out()

    try:
        pools = _get_pools()
    except Exception as ex:
        pools = ResultList().row('g', 'err', 'pools: unexpected ' + repr(ex))
    pools.out()

def _main():
    # only run if the dæmon is configured
    confdir = os.path.dirname(CONFIG)
    if not os.path.exists(confdir):
        return
    if os.access(confdir, os.R_OK | os.X_OK):
        if not os.path.exists(CONFIG):
            return
        # it is configured
        err = None
    else:
        err = 'cannot access configuration, are you running as root?'

    print('<<<kea_dhcpd4:sep(9)>>>')
    if err:
        ResultList().row('g', 'err', err).out()
    else:
        _get_all()
    # last line, to verify completeness
    ResultList().row('g', '/g').out()

if __name__ == "__main__":
    _main()