Automatic Reboot Manager script

autoreboot.py

#!/usr/local/bin/python3.7

# Reboot linux box when connection to router lost 
# Reboot router when connection to internet lost
# Run via cron as root every minute 
# After 5 (ticks) failures a reboot will be initiated
# After 3 reboots wait for 30 reboot ticks before rebooting again

# # sudo crontab -e
# Add line
# * * * * * /usr/local/bin/rebootmanager.py >> /tmp/heartbeat-debug.log 2>&1
# 

import logging
import os
import socket
import configparser

# config keys
import subprocess

KEY_ACTUAL_REBOOTS_COUNT = "key_actual_reboots"
KEY_CONFIG_SECTION = 'DEFAULT'
KEY_REBOOT_TICKER_ROUTER = 'reboot_ticker_router'
KEY_REBOOT_TICKER_PI = 'reboot_ticker_pi'
KEY_FAIL_TICKER_ROUTER = 'fail_ticker_router'
KEY_FAIL_TICKER_PI = 'fail_ticker_pi'

CONF_FILE = 'hb.ini'
CONF_DIR = '/var/log/hb'

log = None


def get_logger(log_file_dir=None,
               log_file_name=None,
               log_to_stdout=True,
               log_to_file=False,
               log_level=logging.DEBUG):
    logger = logging.getLogger(log_file_name)
    logger.setLevel(log_level)

    if not (log_to_stdout or log_to_file):
        return

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

    if log_to_file:
        if not log_file_dir:
            log_file_dir = '.'

        log_file_path = os.path.join(log_file_dir, log_file_name)
        print('Logging to file', log_file_path)
        file_handler = logging.FileHandler(log_file_path)
        file_handler.setLevel(log_level)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

    if log_to_stdout:
        stdout_handler = logging.StreamHandler()
        stdout_handler.setLevel(log_level)
        stdout_handler.setFormatter(formatter)
        logger.addHandler(stdout_handler)

    return logger


def ping(host):
    try:
        socket.gethostbyaddr(host)
        return True
    except Exception:
        return False


def check_host(host=None,
               desc=None,
               conf=None,
               key_fail_ticker=None,
               key_reboot_ticker=None,
               key_total_reboots=None,
               max_failures_before_reboot=None,
               max_reboots_before_wait=None,
               long_wait_duration_in_ticks=None,
               logger=None,
               command=None):
    """
    @param host: host to ping
    @param desc: readable description of host
    @param conf: ConfigParser object
    @param key_fail_ticker: config key for fail ticker
    @param key_reboot_ticker: config key for reboot ticker
    @param max_failures_before_reboot: action will be initiated after N failures
    @param max_reboots_before_wait: after N reboots wait for LONG_WAIT_DURATION_IN_TICKS ticks (minutes)
    @param long_wait_duration_in_ticks: 30 min given that job is called every minute
    @param logger: logger object
    @param command: command to run if conditions are met
    """

    if not ping(host):
        do_reboot = False
        fail_count = conf.getint(KEY_CONFIG_SECTION, key_fail_ticker, fallback=0)
        reboot_ticker = conf.getint(KEY_CONFIG_SECTION, key_reboot_ticker, fallback=0)
        logger.error(f'{desc} [{host}]:  {fail_count + 1}/{max_failures_before_reboot} failures; {reboot_ticker}/{max_reboots_before_wait} reboots')

        fail_count += 1

        if fail_count > max_failures_before_reboot - 1:
            do_reboot = True
            fail_count = 0
            # need to increment the count even when not rebooting in order to keep track when long timeout is reached
            reboot_ticker += 1

        # too many reboots; supress reboots for a while...
        if long_wait_duration_in_ticks > reboot_ticker > max_reboots_before_wait:
            logger.info(f'Max reboot number exceeded. skipping {reboot_ticker}/{long_wait_duration_in_ticks}')
            do_reboot = False

        elif reboot_ticker > max_reboots_before_wait:
            reboot_ticker = 0

        if do_reboot:
            reboot_ticker += 1
            total_reboots = conf.getint(KEY_CONFIG_SECTION, key_total_reboots, fallback=0)
            total_reboots += 1
            conf.set(KEY_CONFIG_SECTION, key_total_reboots, str(total_reboots))
            result = run_command(command.split())
            logger.info(f'Running command {command}:\n{result}')

        # persist settings
        conf.set(KEY_CONFIG_SECTION, key_fail_ticker, str(fail_count))
        conf.set(KEY_CONFIG_SECTION, key_reboot_ticker, str(reboot_ticker))
    else:
        logger.info(f'{desc} ({host}) OK')


def run_command(command):
    p = subprocess.run(command,
                       stdout=subprocess.PIPE,
                       stderr=subprocess.PIPE)
    return p.stdout + b'\n' + p.stderr


def write_config(file_name, parser_obj):
    with open(file_name, 'w') as fn:
        parser_obj.write(fn)


# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
# '''''''''''''''''''''''''''''''''''''''' MAIN ''''''''''''''''''''''''''''''''''''''''
# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
if __name__ == '__main__':

    if not os.path.exists(CONF_DIR):
        print(f'Path {CONF_DIR} does not exist! Creating...')
        os.makedirs(CONF_DIR)

    log = get_logger(log_file_dir=CONF_DIR, log_file_name='hb-log.log', log_to_file=True, log_to_stdout=True)

    parser = configparser.ConfigParser()

    ini_path = os.path.join(CONF_DIR, CONF_FILE)

    if not os.path.exists(ini_path):
        write_config(ini_path, parser)

    parser.read(ini_path, encoding='utf-8')

    check_host('192.168.1.1',
               desc='Router',
               conf=parser,
               key_fail_ticker=KEY_FAIL_TICKER_PI,
               key_reboot_ticker=KEY_REBOOT_TICKER_PI,
               key_total_reboots='reboots_pi',
               max_failures_before_reboot=5,
               max_reboots_before_wait=3,
               long_wait_duration_in_ticks=30,
               logger=log,
               command="/sbin/reboot")

    check_host('1.1.1.1',
               desc='Internet',
               conf=parser,
               key_fail_ticker=KEY_FAIL_TICKER_ROUTER,
               key_reboot_ticker=KEY_REBOOT_TICKER_ROUTER,
               key_total_reboots='reboots_router',
               max_failures_before_reboot=10,
               max_reboots_before_wait=4,
               long_wait_duration_in_ticks=30,
               logger=log,
               command="/usr/local/bin/relay 180")

    write_config(ini_path, parser)