check_polkadot_validator_active

#!/usr/bin/env python3

"""
Nagios/Icinga2 style check in an attempt to verify a validator's state in a more idempotent manner
    (or at least as much as possible)

# Thanks to @aperture-exe:matrix.org for providing some guidance on what to check
#   https://matrix.to/#/
#       !LhjZccBOqFNYKLdmbb:polkadot.builders/
#       $168184239112460RQLKz:matrix.org?via=parity.io&via=matrix.org&via=corepaper.org
# Also inspired by:
#   https://github.com/paritytech/substrate/blob/master/scripts/ci/monitoring/
#       alerting-rules/alerting-rules.yaml
"""

import argparse
import asyncio
import random
import sys
import time
import traceback
from typing import Optional

import requests

STATUS_OK = 0
STATUS_WARNING = 1
STATUS_CRITICAL = 2
STATUS_UNKNOWN = 3

PREFIXES = {
    STATUS_OK: 'OK',
    STATUS_WARNING: 'WARNING',
    STATUS_CRITICAL: 'CRITICAL',
    STATUS_UNKNOWN: 'UNKNOWN'
}

DEFAULT_WS_SERVER = 'localhost'
DEFAULT_WS_PORT = 9944

# NumberOfPeersLow (substrate)
DEFAULT_MIN_PEERS = 3
# NoIncomingConnection (substrate)
# Enhance: substrate_sub_libp2p_incoming_connections_total
DEFAULT_MAX_HEIGHT_DISTANCE_WARN = 5
DEFAULT_MAX_HEIGHT_DISTANCE_CRIT = 10
# BlockProductionSlow (substrate)
DEFAULT_BEST_INCREASE_TIMEOUT = 15
# BlockFinalizationSlow (substrate)
DEFAULT_FINALIZED_INCREASE_TIMEOUT = 15

BLOCK_TYPE_BEST = 'best'
BLOCK_TYPE_FINALIZED = 'finalized'

RETURN_VALUE = -1
RETURN_MESSAGE = ''
RETURN_PERF_STATS = {}

PERF_IS_SYNCING_LABEL = 'is_syncing'
PERF_NUM_PEERS_LABEL = 'num_peers'
PERF_BLOCK_HEIGHT_DIFFERENCE_LABEL = 'height_diff'
PERF_BEST_COMPARE_TO_DIFFERENCE = 'best_compare_diff'
PERF_FINALIZED_COMPARE_TO_DIFFERENCE = 'finalized_compare_diff'
PERF_BEST_INCREASE_TIME = 'best_increased'
PERF_FINALIZED_INCREASE_TIME = 'finalized_increased'

WS_MSG_METHOD_KEY = 'method'
WS_MSG_PARAMS_KEY = 'params'
RPC_HEADERS = {
    'Content-Type': 'application/json',
}

REQUEST_TIMEOUT = 30

MAX_BACKOFF_MS = 1500
BASE_BACKOFF_MS = 200


def _set_return(this_return: int):
    # pylint: disable=global-statement
    global RETURN_VALUE
    if this_return > RETURN_VALUE:
        RETURN_VALUE = this_return


def _set_message(this_msg: str):
    # pylint: disable=global-statement
    global RETURN_MESSAGE
    if RETURN_MESSAGE:
        RETURN_MESSAGE += ' '
    RETURN_MESSAGE += this_msg


def _update_status(this_msg: str, this_return: int):
    """Helper function to update message and return
    """
    _set_message(this_msg)
    _set_return(this_return)


def _get_perf_is_syncing_limits(args):
    if args.warn_is_syncing:
        return ';1;;0;1'
    return ';;1;0;1'


def _get_perf_num_peers_limits(args):
    return f';;{args.num_peers};0;'


def _get_perf_block_height_limits(args):
    return f';{args.max_height_distance_warn};{args.max_height_distance_crit};0;'


def _get_perf_best_block_increase_limits(args):
    return f's;;{args.best_block_timeout};0;'


def _get_perf_finalized_block_increase_limits(args):
    return f's;;{args.finalized_block_timeout};0;'


def _update_perf_stats(perf_label: str, perf_value: str):
    RETURN_PERF_STATS[perf_label] = perf_value


def _get_output_message() -> str:
    return (f'{PREFIXES[RETURN_VALUE]}: {RETURN_MESSAGE}'
            f' | {" ".join(f"{k}={v}" for k, v in RETURN_PERF_STATS.items())}')


def _build_post_data(update_dict) -> dict:
    empty_msg = {
        'id': 1,
        'jsonrpc': '2.0',
        WS_MSG_METHOD_KEY: '',
        WS_MSG_PARAMS_KEY: [],
    }
    empty_msg.update(update_dict)
    return empty_msg


async def _connect_rpc(uri: str,
                       message_dict: dict,
                       args,
                       return_result_dict: bool = True) -> Optional[dict]:
    response = requests.post(uri,
                             json=message_dict,
                             headers=RPC_HEADERS,
                             timeout=REQUEST_TIMEOUT)
    if args.verbose:
        print(f'Request: {uri}. json: {message_dict}.'
              f' Response: {response}. Text: {response.text}')
    if response.status_code == 200:
        if return_result_dict:
            return response.json()['result']
        return response.json()
    _update_status(f'RPC request to uri: {uri} and message: {message_dict}'
                   f' returned {response.status_code}: {response}',
                   STATUS_UNKNOWN)
    return _exit()


async def _is_syncing(args) -> bool:
    return (await _connect_rpc(args.node_uri,
                               _build_post_data({WS_MSG_METHOD_KEY: 'system_health'}),
                               args))['isSyncing']  # type: ignore


async def _get_num_peers(args) -> int:
    return (await _connect_rpc(args.node_uri,
                               _build_post_data({WS_MSG_METHOD_KEY: 'system_health'}),
                               args))['peers']  # type: ignore


async def _get_best_head_number(args, connect_node: bool = True) -> str:
    return (await _connect_rpc(args.node_uri if connect_node else args.compare_to_uri,
                               _build_post_data({WS_MSG_METHOD_KEY: 'chain_getHeader'}),
                               args))['number']  # type: ignore


async def _get_finalized_head(args, connect_node: bool = True) -> str:
    return (await _connect_rpc(args.node_uri if connect_node else args.compare_to_uri,
                               _build_post_data({WS_MSG_METHOD_KEY: 'chain_getFinalizedHead'}),
                               args,
                               return_result_dict=False))['result']  # type: ignore


async def _get_finalized_head_number(args, connect_node: bool = True) -> str:
    finalized_head = await _get_finalized_head(args, connect_node=connect_node)
    return (await _connect_rpc(args.node_uri if connect_node else args.compare_to_uri,
                               _build_post_data({WS_MSG_METHOD_KEY: 'chain_getHeader',
                                                 WS_MSG_PARAMS_KEY: [finalized_head]}),
                               args))['number']  # type: ignore


def _calculate_block_number_distance(best_block_num: str, finalized_block_num: str) -> int:
    return abs(int(best_block_num, 16) - int(finalized_block_num, 16))


async def _check_is_syncing(args):
    # Check if we're syncing
    is_syncing = await _is_syncing(args)
    if is_syncing:
        _update_status('isSyncing: True.',
                       STATUS_WARNING if args.warn_is_syncing else STATUS_CRITICAL)
    else:
        _set_return(STATUS_OK)
    _update_perf_stats(PERF_IS_SYNCING_LABEL,
                       f'{1 if is_syncing else 0}{_get_perf_is_syncing_limits(args)}')


async def _check_num_peers(args):
    # Check how many peers we have
    num_peers = await _get_num_peers(args)
    if num_peers < args.num_peers:
        _update_status(f'Peers: {num_peers}.', STATUS_CRITICAL)
    else:
        _set_return(STATUS_OK)
    _update_perf_stats(PERF_NUM_PEERS_LABEL,
                       f'{num_peers}{_get_perf_num_peers_limits(args)}')


async def _check_block_height_difference(args):
    # Check block height distance
    best_head_block_number = await _get_best_head_number(args)
    finalized_head_block_number = await _get_finalized_head_number(args)
    block_height_difference = _calculate_block_number_distance(best_head_block_number,
                                                               finalized_head_block_number)
    if block_height_difference >= args.max_height_distance_crit:
        _update_status(__get_block_height_difference_message(best_head_block_number,
                                                             finalized_head_block_number),
                       STATUS_CRITICAL)
    elif block_height_difference >= args.max_height_distance_warn:
        _update_status(__get_block_height_difference_message(best_head_block_number,
                                                             finalized_head_block_number),
                       STATUS_WARNING)
    else:
        _set_return(STATUS_OK)
    _update_perf_stats(PERF_BLOCK_HEIGHT_DIFFERENCE_LABEL,
                       f'{block_height_difference}{_get_perf_block_height_limits(args)}')


def __get_block_height_difference_message(best_number: str,
                                          finalized_number: str) -> str:
    # Generate detailed message used when block height difference is not ok
    return (f'Node Block Heights:'
            f' Best: {best_number} ({int(best_number, 16)}).'
            f' Finalized: {finalized_number} ({int(finalized_number, 16)}).'
            f' Diff: {_calculate_block_number_distance(best_number, finalized_number)}.')


async def _check_block_compare_to(args, block_type: str):
    if block_type == BLOCK_TYPE_BEST:
        block_num_func = _get_best_head_number
        perf_label = PERF_BEST_COMPARE_TO_DIFFERENCE
    elif block_type == BLOCK_TYPE_FINALIZED:
        block_num_func = _get_finalized_head_number
        perf_label = PERF_FINALIZED_COMPARE_TO_DIFFERENCE
    else:
        print(f'Code error, unknown block_type: {block_type}')
        sys.exit(STATUS_UNKNOWN)
    await _do_check_block_compare_to(args,
                                     block_type,
                                     block_num_func,
                                     perf_label)


async def _do_check_block_compare_to(args, block_type, block_num_func, perf_label):
    if args.verbose:
        print(f'Comparing {block_type} block height to {args.compare_to_uri}')
    compare_block_number = await block_num_func(args, connect_node=False)
    node_block_number = await block_num_func(args, connect_node=True)
    block_height_difference = _calculate_block_number_distance(compare_block_number,
                                                               node_block_number)
    if block_height_difference >= args.max_height_distance_crit:
        _update_status(f'Compare To Block Heights ({block_type}):'
                       f' Compare To: {compare_block_number}.'
                       f' Node: {node_block_number}.',
                       STATUS_CRITICAL)
    elif block_height_difference >= args.max_height_distance_warn:
        _update_status(f'Compare To Block Heights ({block_type}):'
                       f' Compare To: {compare_block_number}.'
                       f' Node: {node_block_number}.',
                       STATUS_WARNING)
    else:
        _set_return(STATUS_OK)
    _update_perf_stats(
        perf_label, f'{block_height_difference}{_get_perf_block_height_limits(args)}')


async def _check_block_increasing(args, block_type: str):
    if block_type == BLOCK_TYPE_BEST:
        block_num_func = _get_best_head_number
        increase_timeout = args.best_block_timeout
        perf_label = PERF_BEST_INCREASE_TIME
        perf_units = _get_perf_best_block_increase_limits(args)
    elif block_type == BLOCK_TYPE_FINALIZED:
        block_num_func = _get_finalized_head_number
        increase_timeout = args.finalized_block_timeout
        perf_label = PERF_FINALIZED_INCREASE_TIME
        perf_units = _get_perf_finalized_block_increase_limits(args)
    else:
        print(f'Code error, unknown block_type: {block_type}')
        sys.exit(STATUS_UNKNOWN)
    await _do_check_block_increasing(args,
                                     block_type,
                                     block_num_func,
                                     increase_timeout,
                                     perf_label,
                                     perf_units)


async def _do_check_block_increasing(args, block_type, block_num_func,
                                     increase_timeout, perf_label, perf_units):
    initial_block_num = await block_num_func(args)
    if args.verbose:
        print(f'Checking {block_type} blocks. Initial: {initial_block_num}.')
    initial_block_time = time.time()
    new_block_time = initial_block_time
    block_increase_elapsed = 0
    iteration_count = 0
    while block_increase_elapsed < increase_timeout:
        time.sleep(__get_backoff_time(iteration_count))
        new_block_num = await block_num_func(args)
        if args.verbose:
            print(f'Checking {block_type} blocks. Found: {new_block_num}.')
        new_block_time = time.time()
        if new_block_num > initial_block_num:
            # block number increased
            break
        # block number did not increase, update our elapsed time
        block_increase_elapsed = new_block_time - initial_block_time
    if block_increase_elapsed >= increase_timeout:
        if args.verbose:
            print(f'Checking {block_type} exceeded timeout before increasing from'
                  f' {initial_block_num}.')
        _update_status(f'{block_type} block number did not increase after'
                       f' {block_increase_elapsed}s',
                       STATUS_CRITICAL)
        block_increase_elapsed = ('inf' if not args.fake_increasing_perfdata
                                  else (increase_timeout * 2.0))
    _update_perf_stats(perf_label, f'{block_increase_elapsed}{perf_units}')


def __get_backoff_time(iteration_count: int) -> float:
    """Calculate a backoff time using Decorrelated Jitter.
    Unit of time is ms and then convert to seconds as a float
    """
    temp = min(MAX_BACKOFF_MS, BASE_BACKOFF_MS * 2 ** iteration_count)
    sleep = temp / 2 + random.randint(0, int(temp / 2))
    sleep = min(MAX_BACKOFF_MS, random.randint(BASE_BACKOFF_MS, int(sleep * 3)))
    return sleep / 1000


def _exit():
    if RETURN_VALUE == STATUS_OK:
        _set_message('All checks passed.')
    print(_get_output_message())
    sys.exit(RETURN_VALUE)


async def run(args):
    """
    Perform the checks in a reasonable order.
    If the height distance is too large, skip the checks regarding Production and Finalization speed
    """
    await _check_is_syncing(args)
    await _check_num_peers(args)
    await _check_block_height_difference(args)
    if args.compare_to_uri:
        await _check_block_compare_to(args, BLOCK_TYPE_BEST)
        await _check_block_compare_to(args, BLOCK_TYPE_FINALIZED)
    else:
        await _check_block_increasing(args, BLOCK_TYPE_BEST)
        await _check_block_increasing(args, BLOCK_TYPE_FINALIZED)
    _exit()


def parse_args():
    """Parse command line arguments
    """
    parser = argparse.ArgumentParser(
        "This plugin checks a Polkadot/Kusama validator's Active status.")
    parser.add_argument('-s', '--server', type=str,
                        required=False,
                        dest='server',
                        default=DEFAULT_WS_SERVER,
                        help=f'WS server.'
                             f' Default: {DEFAULT_WS_SERVER}')
    parser.add_argument('-p', '--port', type=int,
                        required=False,
                        dest='port',
                        default=DEFAULT_WS_PORT,
                        help=f'WS port to connect to.'
                             f' Default: {DEFAULT_WS_PORT}')
    parser.add_argument('--use-tls',
                        dest='use_tls',
                        action='store_true',
                        help='use TLS when connecting to node <server>.'
                             ' Default: false (use http://)')
    parser.add_argument('--warn-is-syncing',
                        dest='warn_is_syncing',
                        action='store_true',
                        help='isSyncing should be treated as a Warning.'
                             ' Default: Critical')
    parser.add_argument('--min-peers', type=int,
                        required=False,
                        dest='num_peers',
                        default=DEFAULT_MIN_PEERS,
                        help=f'minimum number of peers.'
                             f' Default: {DEFAULT_MIN_PEERS}')
    parser.add_argument('--max-distance-warn', type=int,
                        required=False,
                        dest='max_height_distance_warn',
                        default=DEFAULT_MAX_HEIGHT_DISTANCE_WARN,
                        help=f'max height distance for warning.'
                             f' Default: {DEFAULT_MAX_HEIGHT_DISTANCE_WARN}')
    parser.add_argument('--max-distance-crit', type=int,
                        required=False,
                        dest='max_height_distance_crit',
                        default=DEFAULT_MAX_HEIGHT_DISTANCE_CRIT,
                        help=f'max height distance for critical.'
                             f' Default: {DEFAULT_MAX_HEIGHT_DISTANCE_CRIT}')
    parser.add_argument('--best-timeout', type=int,
                        required=False,
                        dest='best_block_timeout',
                        default=DEFAULT_BEST_INCREASE_TIMEOUT,
                        help=f'best block increase timeout (in seconds).'
                             f' Default: {DEFAULT_BEST_INCREASE_TIMEOUT}')
    parser.add_argument('--finalized-timeout', type=int,
                        required=False,
                        dest='finalized_block_timeout',
                        default=DEFAULT_FINALIZED_INCREASE_TIMEOUT,
                        help=f'finalized block increase timeout (in seconds).'
                             f' Default: {DEFAULT_FINALIZED_INCREASE_TIMEOUT}')
    parser.add_argument('--compare-to-uri', type=str,
                        required=False,
                        dest='compare_to_uri',
                        default=None,
                        help='node to compare best and finalized block heights to.'
                             ' Default: Use --best-timeout and --finalized-timeout'
                             ' and connect to http://<server>:<port>')
    parser.add_argument('--fake-increasing-perfdata',
                        dest='fake_increasing_perfdata',
                        action='store_true',
                        help='use 2x timeout value instead of inf.'
                             ' Default: inf when timeout occurs')
    parser.add_argument('-v', '--verbose',
                        dest='verbose',
                        action='store_true',
                        help='verbose output')
    args = parser.parse_args()
    if args.max_height_distance_warn > args.max_height_distance_crit:
        print('Max Height Distance warning must be less than critical')
        sys.exit(STATUS_CRITICAL)
    # build a full URI to use for node <server>
    args.node_uri = f'http{"s" if args.use_tls else ""}://{args.server}:{args.port}'
    # check if ignored arguments are provided
    if args.compare_to_uri:
        if args.best_block_timeout != DEFAULT_BEST_INCREASE_TIMEOUT:
            print('--best-timeout was set but will be ignored'
                  ' since --compare-to-uri is also set')
        if args.finalized_block_timeout != DEFAULT_FINALIZED_INCREASE_TIMEOUT:
            print('--finalized-timeout was set but will be ignored'
                  ' since --compare-to-uri is also set')
    return args


if __name__ == '__main__':
    try:
        asyncio.run(run(parse_args()))
    # pylint: disable=broad-except
    except Exception as ex:
        print(f'{sys.argv[0]}: Unhandled exception {type(ex)}')
        print(ex)
        traceback.print_exc()
        sys.exit(STATUS_CRITICAL)