import prometheus_client
import subprocess
import time
import argparse

SERVER_HOST = '0.0.0.0'
# Set this port to whatever you want this service to bind to
SERVER_PORT = 9999
FETCH_INTERVAL = 5

class Session:
    """
    This class is used to create a Session object containing info on an SSH session
    """
    def __init__(self, name, tty, from_, login, idle, jcpu, pcpu, what):
        self.name = name
        self.tty = tty
        self.from_ = from_
        self.login = login
        self.idle = idle
        self.jcpu = jcpu
        self.pcpu = pcpu
        self.what = what

    def __str__(self):
        return "%s %s" % (self.name, self.from_)

    def __repr__(self):
        return "%s %s" % (self.name, self.from_)

    def __eq__(self, other):
        return self.login == other.login and self.tty == other.tty and self.from_ == other.from_

    def to_dict(self):
        return {
            'name': self.name,
            'tty': self.tty,
            'from_': self.from_,
            'login': self.login,
            'idle': self.idle,
            'jcpu': self.jcpu,
            'pcpu': self.pcpu,
            'what': self.what
        }

    def to_list(self):
        return [self.name, self.tty, self.from_, self.login, self.idle, self.jcpu, self.pcpu, self.what]


def contains_user_list(user, other_user_list):
    for other_user in other_user_list:
        if are_equal(user, other_user):
            #print("Found equals: %s and %s" % (user, other_user))
            return True
    #print("Not found: %s in %s" % (user, other_user_list))
    return False


def are_equal(user_list, other_user_list):
    assert len(user_list) == len(other_user_list)
    for i in range(4):
        if user_list[i] != other_user_list[i]:
            return False
    return True


"""
Returns a list of Session objects
The string generated by the "w -h" command looks like follows:
USER          TTY      FROM             LOGIN@   IDLE   JCPU   PCPU   WHAT
<username>    pts/0    192.168.0.X      14:57    1:54m  0.12s  0.12s  -bash
"""
def get_w_data():
    # Retrieve command output and discard the header
    w_data = subprocess.check_output(['w', '-h']).decode('utf-8')
    w_data = w_data.split('\n')[:-1]
    users = []
    # Generate a list of Session objects from the output of the command
    for user_string in w_data:
        user_string_cleaned = user_string.split()
        if len(user_string_cleaned) == 8:
            users.append(Session(*user_string_cleaned))
    return users


def parse_arguments():
    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


"""
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.
"""

if __name__ == '__main__':

    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))
    num_sessions = []
    gauge_num_sessions = prometheus_client.Gauge('ssh_num_sessions', 'Number of SSH sessions', ['remote_ip'])
    data = get_w_data()
    list_data = [user.to_list() for user in data]

    # Initial metrics
    print("Connections at startup:")
    for i in range(len(list_data)):
        gauge_num_sessions.labels(remote_ip=list_data[i][2]).inc()
        print("Initial connection: {}".format(list_data[i]))

    # Generate some requests.
    print("Looking for SSH connection changes at interval {}".format(FETCH_INTERVAL))
    while True:

        list_old_data = list_data
        data = get_w_data()
        list_data = [user.to_list() for user in data]
        num_sessions = len(data)

        for i in range(num_sessions):
            # Looking for newly found SSH sessions
            if not contains_user_list(list_data[i], list_old_data):
                print("Session connected: %s" % list_data[i])
                gauge_num_sessions.labels(remote_ip=list_data[i][2]).dec()
           
        for i in range(len(list_old_data)):
            # Looking for SSH sessions that no longer exist
            if not contains_user_list(list_old_data[i], list_data):
                print("Session disconnected: %s" % list_old_data[i])
                gauge_num_sessions.labels(remote_ip=list_old_data[i][2]).inc()

        time.sleep(FETCH_INTERVAL)