Handles short sessions

Can now handle sessions of duration < FETCH_INTERVAL

prometheus-ssh-exporter.py

@@ -2,11 +2,21 @@ import prometheus_client
 import time
 import argparse
 import utmp
+from watchdog.observers import Observer
+from watchdog.events import FileSystemEventHandler, FileModifiedEvent
+
 
 # These defaults can be overwritten by command line arguments
 SERVER_HOST = '0.0.0.0'
 SERVER_PORT = 9999
 FETCH_INTERVAL = 15
+WATCHFILE = '/var/run/utmp'
+
+
+class FileOpenedHandler(FileSystemEventHandler):
+    def on_modified(self, event):
+        if not event.is_directory and isinstance(event, FileModifiedEvent) and event.src_path == WATCHFILE:
+            handle_sessions_changed()
 
 
 class Session:
@@ -53,54 +63,13 @@ def get_utmp_data() -> list[Session]:
 
 
 
-def parse_arguments() -> None:
+def handle_sessions_changed() -> None:
-
-    global FETCH_INTERVAL, SERVER_PORT, SERVER_HOST
-
-    parser = argparse.ArgumentParser(
-        prog='python prometheus-ssh-exporter.py',
-        description='Prometheus exporter for info about SSH sessions')
-    parser.add_argument('-H', '--host', type=str,
-                        default='0.0.0.0', help='Hostname to bind to')
-    parser.add_argument('-p', '--port', type=int, default=9999,
-                        help='Port for the server to listen to')
-    parser.add_argument('-i', '--interval', type=int, default=15,
-                        help='Interval in seconds to fetch SSH sessions data')
-
-    args = parser.parse_args()
-    FETCH_INTERVAL = args.interval
-    SERVER_PORT = args.port
-    SERVER_HOST = args.host
-
-
-if __name__ == '__main__':
     """ 
-    This program exports the number of SSH sessions as a metric "ssh_num_sessions" for prometheus.
+    This function fetches the current list of SSH sessions and compares it to the previous list of SSH sessions.
-    It applies a label to each increment or decrement of that number, containing the remote IP address.
+    If the number of sessions has changed, it increments or decrements the gauge_num_sessions metric
-    That way we can filter by the remote IP in Grafana, getting the number of SSH sessions by IP address,
+    and updates the session_data and num_sessions variables.
-    or sum them up to get the total number of sessions.
     """
-
+    global session_data, num_sessions, gauge_num_sessions, old_session_data, old_num_sessions
-    parse_arguments()
-
-    # Start up the server to expose the metrics.
-    prometheus_client.start_http_server(SERVER_PORT)
-    print("Started metrics server bound to {}:{}".format(SERVER_HOST, SERVER_PORT))
-    gauge_num_sessions = prometheus_client.Gauge(
-        'ssh_num_sessions', 'Number of SSH sessions', ['remote_ip'])
-    
-    session_data = get_utmp_data()
-    num_sessions = len(session_data)
-    
-    # Initial metrics
-    print("Active sessions at startup:")
-    for session in session_data:
-        gauge_num_sessions.labels(remote_ip=session.from_).inc()
-        print("Initial connection: {}".format(session.from_))
-
-    # Generate some requests.
-    print("Looking for SSH connection changes at interval {}".format(FETCH_INTERVAL))
-    while True:
 
     old_session_data = session_data
     old_num_sessions = len(old_session_data)
@@ -120,4 +89,78 @@ if __name__ == '__main__':
             print("Session disconnected: %s" % maybe_old_session.from_)
             gauge_num_sessions.labels(remote_ip=maybe_old_session.from_).dec()
 
+
+
+def parse_arguments() -> None:
+    global FETCH_INTERVAL, SERVER_PORT, SERVER_HOST, WATCHFILE
+
+    parser = argparse.ArgumentParser(
+        prog='python prometheus-ssh-exporter.py',
+        description='Prometheus exporter for info about SSH sessions')
+    parser.add_argument('-H', '--host', type=str,
+                        default='0.0.0.0', help='Hostname to bind to')
+    parser.add_argument('-p', '--port', type=int, default=9999,
+                        help='Port for the server to listen to')
+    parser.add_argument('-i', '--interval', type=int, default=15,
+                        help='Interval in seconds to fetch SSH sessions data')
+    parser.add_argument('-f', '--file', type=str, default='/var/run/utmp',
+                        help='File that changes every time a new SSH session is opened or closed')
+
+    args = parser.parse_args()
+    FETCH_INTERVAL = args.interval
+    SERVER_PORT = args.port
+    SERVER_HOST = args.host
+    WATCHFILE = args.file
+
+
+
+if __name__ == '__main__':
+    """
+    This program exports the number of SSH sessions as a metric "ssh_num_sessions" for prometheus.
+    It applies a label to each increment or decrement of that number, containing the remote IP address.
+    That way we can filter by the remote IP in Grafana, getting the number of SSH sessions by IP address,
+    or sum them up to get the total number of sessions.
+    """
+
+
+    parse_arguments()
+
+    # Start up the server to expose the metrics.
+    prometheus_client.start_http_server(SERVER_PORT)
+    print("Started metrics server bound to {}:{}".format(SERVER_HOST, SERVER_PORT))
+    gauge_num_sessions = prometheus_client.Gauge(
+        'ssh_num_sessions', 'Number of SSH sessions', ['remote_ip'])
+    
+    # session_data contains the current list of sessions
+    session_data = get_utmp_data()
+    num_sessions = len(session_data)
+    
+    # Initial metrics
+    for session in session_data:
+        gauge_num_sessions.labels(remote_ip=session.from_).inc()
+        print("Initial connection: {}".format(session.from_))
+
+
+    """
+    Start the watchdog to monitor the WATCHDOG file for changes. 
+    This is used to immediately look for changes in the SSH sessions when a new session is opened or closed
+    to prevent missing any sessions that lasted less than the FETCH_INTERVAL.
+    """
+    event_handler = FileOpenedHandler()
+    observer = Observer()
+    observer.schedule(event_handler, path=WATCHFILE, recursive=False)
+    observer.start()
+
+    # Generate some requests.
+    print("Looking for SSH connection changes at interval {}".format(FETCH_INTERVAL))
+    try:
+
+        while True:
+            # Keep looking for changes in the SSH sessions in case the watchdog missed something
+            handle_sessions_changed()
             time.sleep(FETCH_INTERVAL)
+
+    except KeyboardInterrupt:
+        print("Manual interrupt...")
+        observer.stop()
+    observer.join()

requirements.txt

@@ -1,2 +1,3 @@
 prometheus_client==0.16.0
 utmp==21.10.0
+watchdog==3.0.0