notify on startup and exit

clarify notifications
refactoring: comments, docstrings

1 files changed, 101 insertions, 38 deletions

diff --git a/src/server_status.py b/src/server_status.py
index d6a1989..a0dc6b3 100644
--- a/src/server_status.py
+++ b/src/server_status.py
@@ -1,14 +1,53 @@
 import telegram  # Notifications
 import asyncio  # Telegram library is async
-import requests  # HTTP
+import requests  # Check HTTP status code
 import os  # Unix socket file handling
-import subprocess  # Ping
+import subprocess  # Ping check if host online
 import socket  # Hostname, unix socket
 import yaml  # Config file
 import logging  # Log messages
+import signal  # Notify before exiting
 
-logger = logging.getLogger('server_status')
-logger.setLevel(level=logging.DEBUG)
+
+class ColorFormatter(logging.Formatter):
+    grey = "\x1b[90;20m"
+    cyan = "\x1b[96;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+    format = "%(name)s %(asctime)s,%(msecs)03d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s"
+
+    FORMATS = {
+        logging.DEBUG: grey + format + reset,
+        logging.INFO: cyan + format + reset,
+        logging.WARNING: yellow + format + reset,
+        logging.ERROR: red + format + reset,
+        logging.CRITICAL: bold_red + format + reset
+    }
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
+
+logger = logging.getLogger(__name__)  # Instantiate a logger to be used in this module
+
+# Display every message Change this to INFO to see INFO and above (filter out DEBUG)
+# See: https://docs.python.org/3/howto/logging.html#logging-levels
+logger.root.setLevel(logging.DEBUG)
+
+stream_handler = logging.StreamHandler()  # This catches and handles log messages on the root handler
+stream_handler.setFormatter(ColorFormatter())
+logger.root.addHandler(stream_handler)
+
+# Disable logging for sub libraries
+logging.getLogger("urllib3").setLevel(logging.WARNING)
+logging.getLogger("httpx").setLevel(logging.WARNING)
+logging.getLogger("httpcore").setLevel(logging.WARNING)
+logging.getLogger("asyncio").setLevel(logging.WARNING)
+logging.getLogger("telegram").setLevel(logging.WARNING)
 
 with open("config.yml", "r") as file:
     config = yaml.safe_load(file)
@@ -20,7 +59,7 @@ socket_path = "/tmp/server_status.sock"
 bot = telegram.Bot(config["telegram"]["token"])
 hostname = socket.gethostname()
 # Used to only notify on state change
-# Initialise data stucture based on targets in config file
+# Initialise data structure based on targets in config file
 state = {}
 for host, host_config in config["hosts"].items():
     state[host] = {"online": None, "open_ports": []}
@@ -47,7 +86,12 @@ def ping(host: str) -> bool:
     return response.returncode == 0
 
 
-def check_target(host: str, host_config: dict):
+def check_target(host: str, host_config: dict) -> bool:
+    """
+    :param host: Hostname or IP
+    :param host_config: Config from config file, specifying what to test and desired state
+    :return: Host is desired state
+    """
     # Ping
     online = ping(host)
     # Default to desire host online, if not specified in config file
@@ -62,49 +106,55 @@ def check_target(host: str, host_config: dict):
     ''')
     # notify if the host online state changed, or on fresh boot, if it isn't the desired online state
     if changed_state or (first_test and not host_is_desired_state):
-        message = f"`{host}` is {'online' if online else 'offline'}"
+        message = f"❌ `{host}` is {'online' if online else 'offline'}"
         logger.warning(message)
         send_message(message)
     state[host]['online'] = online  # Save current state
 
-    if not online:
-        return
-
-    messages = []
-    for url in host_config['urls']:
-        http_state = None
-        # HTTP Code
-        try:
-            # print(f'''
-            # ---=== Getting: {url} ===---
-            # ''')
-            r = requests.get(url)
-            if r.status_code != 200:
-                http_state = r.status_code
-        except requests.exceptions.SSLError as exception:
+    if online:
+        messages = []
+        for url in host_config['urls']:
+            http_state = None
+            # HTTP Code
             try:
-                reason = exception.args[0].reason.args[0].verify_message
+                # print(f'''
+                # ---=== Getting: {url} ===---
+                # ''')
+                r = requests.get(url)
+                if r.status_code != 200:
+                    http_state = r.status_code
+            except requests.exceptions.SSLError as exception:
+                try:
+                    reason = exception.args[0].reason.args[0].verify_message
+                except Exception as exception:
+                    logger.warning("Unable to get reason")
+                    raise
+                http_state = "SSL error"
             except Exception as exception:
-                logger.warning("Unable to get reason")
                 raise
-            http_state = "SSL error"
-        except Exception as exception:
-            raise
-        if http_state:
-            if state[host]['http_status'][url] is None or state[host]['http_status'][url] != http_state:
-                messages.append(f'[{url}]({url}) {http_state}')
-                state[host]['http_status'][url] = http_state
+            if http_state:
+                if state[host]['http_status'][url] is None or state[host]['http_status'][url] != http_state:
+                    messages.append(f'[{url}]({url}) {http_state}')
+                    state[host]['http_status'][url] = http_state
+
+        if len(messages) > 0:
+            send_message('❌ ' + '\n'.join(messages))
 
-    if len(messages) > 0:
-        send_message('\n'.join(
-            messages))
+    return host_is_desired_state
 
 
-def check_status():
+def check_status() -> bool:
+    """
+    :return: All hosts are in desired state
+    """
     try:
         logger.info("Checking status")
+        all_hosts_are_in_desired_state = True
         for host, host_config in config["hosts"].items():
-            check_target(host, host_config)
+            host_is_desired_state = check_target(host, host_config)
+            if host_is_desired_state == False:
+                all_hosts_are_in_desired_state = False
+        return all_hosts_are_in_desired_state
     except Exception as ex:
         send_message("Error getting status")
         raise ex
@@ -158,10 +208,23 @@ def listen():
 
 
 def main():
-    # send_message(f"`{hostname}` online")
-    check_status()
+    send_message(f"✅ Server status checker on `{hostname}` *online*")
+    alles_gut = check_status()
+    if alles_gut == True:
+        # First test, notify of outcome
+        # If a non-desired state was caught in first check a notification would have been triggered.
+        send_message("✅ All hosts are in desired state,\nand no discrepancies found")  # Vægter
     listen()
 
 
+def service_exiting(signum, frame):
+    # Ship going down
+    send_message(f"❌ Server status checker on `{hostname}` *going down*")
+    exit(0)
+
+
+# Intercept SIGTERM (15) and notify before exiting
+signal.signal(signal.SIGTERM, service_exiting)
+
 if __name__ == "__main__":
     asyncio.run(main())