#!/usr/bin/python3

"""
Monitoring Script for HAProxy Stats via Zabbix LLD.
by Marcel Jäpel (2021)

Usage:
    script.py -s <SOCKET> -t <TABLE>

Requirements:
    apt-get install python
"""

########################################################### Imports

try:
    import argparse
    import json
    import socket
except Exception as e:
    return_data = {"data": [], "status_ok": True, "status_errmsg": "ImportError: " + str(e)}
    print(return_data)

########################################################### Main


def main() -> None:
    parser = argparse.ArgumentParser(description="HAProxy Stats by Marcel Jäpel (2021)")
    parser.add_argument("--socket", "-s", action="store", default="/run/haproxy/admin.sock", dest="SOCKET", help="Path to UNIX socket")
    parser.add_argument("--table", "-t", action="store", default="lbcluster1-elasticsearch", dest="TABLE", help="Name of stats table")

    return_data = {"data": [], "keys": [], "status_ok": True, "status_errmsg": ""}

    ########################################## get data from socket
    try:
        args = parser.parse_args()

        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        sock.settimeout(5)
        sock.connect(args.SOCKET)

        msg = "show table " + args.TABLE + "\n"
        sock.sendall(msg.encode())

        buff_size = 4096  # 4 KiB
        data = b""
        while True:
            buff = sock.recv(buff_size)
            data += buff
            if len(buff) < buff_size:
                break

        data = data.decode()
        data = data.split("\n")
        data = data[1:-2]

    except Exception as e:
        return_data["status_ok"] = False
        return_data["status_errmsg"] = "Socket error: " + str(e)

    ########################################## parse data
    try:
        for line in data:
            line_dict = {}
            line_splitted = line.split(" ")[1:]

            for attr in line_splitted:
                attr_key, attr_value = attr.split("=", 2)
                attr_key = attr_key.split("(")[0]
                line_dict[attr_key] = attr_value

                if attr_key == "key":
                    line_dict["{#KEY}"] = attr_value

            return_data["data"].append(line_dict)
    except Exception as e:
        return_data["status_ok"] = False
        return_data["status_errmsg"] = "Socket error: " + str(e)

    ########################################## output
    print(json.dumps(return_data, indent=4, sort_keys=True))


if __name__ == "__main__":
    main()