nando/nando

#!/usr/bin/env python3

import socket, sys, time, os, configobj

RC = "/etc/nandorc"
BLOCK_SIZE = 10240


def setup_options():
    """ Setup the command line options """
    from argparse import ArgumentParser

    parser = ArgumentParser(
        description="Alive notifier. Options first read from /etc/nandorc, and overriden by switches."
    )
    parser.add_argument(
        "--host",
        action="store",
        dest="HOST",
        type=str,
        default=None,
        help="Host name/IP of the server",
    )
    parser.add_argument(
        "--port",
        action="store",
        dest="PORT",
        default=None,
        type=int,
        help="Port number of the server",
    )
    parser.add_argument(
        "-t",
        action="store",
        dest="TIMEOUT",
        default=5,
        type=int,
        help="Connection timeout",
    )
    parser.add_argument(
        "-i",
        action="store",
        dest="INTERVAL",
        default=0,
        type=int,
        help="Send signal every i seconds. If 0, send only once and exit.",
    )
    parser.add_argument(
        "-m",
        action="store",
        dest="MSGFILE",
        type=str,
        default=None,
        help="Read a text file as description. Max length 512 chars",
    )
    parser.add_argument(
        "--query-ip",
        action="store",
        dest="QUERY_HOSTNAME",
        default=None,
        type=str,
        help="Query the IP of a host name",
    )
    parser.add_argument(
        "--query-host",
        action="store",
        dest="QUERY_IP",
        default=None,
        type=str,
        help="Query the hostname of an IP",
    )
    parser.add_argument(
        "command",
        type=str,
        action="store",
        default="",
        nargs="?",
        choices=["", "list", "alive", "lost"],
        help="Command to send. If empty, send standard alive-notice. Commands: list, alive",
    )
    options = parser.parse_args()
    if os.path.exists(RC):
        rc = configobj.ConfigObj(RC)
        if options.HOST is None:
            if "HOST" in rc.keys():
                options.HOST = rc["HOST"]
        if options.PORT is None:
            if "PORT" in rc.keys():
                options.PORT = int(rc["PORT"])
        if options.MSGFILE is None:
            if "MSGFILE" in rc.keys():
                options.MSGFILE = rc["MSGFILE"]

    if options.HOST is None:
        parser.error("Host name required")
    if options.PORT is None:
        parser.error("Port number required")
    return options


def print_table(data):
    lengths = [0, 0, 0, 0]
    for row in data.split("\n"):
        cols = row.split("|", 4)
        if len(cols) != 4:
            continue
        for c in range(4):
            lengths[c] = max(lengths[c], len(cols[c]))
    lengths = [str(x) for x in lengths]
    for row in data.split("\n"):
        cols = row.split("|", 4)
        if len(cols) != 4:
            continue
        print(
            (
                "{0:<"
                + lengths[0]
                + "} | {1:<"
                + lengths[1]
                + "} | {2:>"
                + lengths[2]
                + "} | {3}"
            ).format(*cols)
        )


def query_ip(sock, opts):
    sock.sendto("list", (opts.HOST, opts.PORT))
    received = sock.recv(BLOCK_SIZE)
    for row in received.split("\n"):
        cols = row.split("|", 4)
        if cols[0] == opts.QUERY_HOSTNAME:
            sys.stdout.write(cols[1])
            sys.exit(0)
    sys.exit(1)


def query_host(sock, opts):
    sock.sendto("list", (opts.HOST, opts.PORT))
    received = sock.recv(BLOCK_SIZE)
    for row in received.split("\n"):
        cols = row.split("|", 4)
        if cols[1] == opts.QUERY_IP:
            sys.stdout.write(cols[0])
            sys.exit(0)
    sys.exit(1)


def read_desc(opts):

    if not opts.MSGFILE:
        return ""
    if not os.path.isfile(opts.MSGFILE):
        return ""

    DESC = open(opts.MSGFILE, "rb").read(1024)
    DESC = "".join([str(c) for c in DESC.decode("utf-8") if c not in ["\n", "\r", "|"]])
    return DESC


opts = setup_options()

errors = 0
while True:
    try:
        MYNAME = socket.gethostname()
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.settimeout(opts.TIMEOUT)
        sock.connect((opts.HOST, opts.PORT))
        MYIP = sock.getsockname()[0]

        if opts.QUERY_HOSTNAME:
            query_ip(sock, opts)
        if opts.QUERY_IP:
            query_host(sock, opts)

        if opts.command == "":
            DESC = read_desc(opts)
            MSG = "{0}|{1}|{2}\n".format(MYNAME, MYIP, DESC)
        else:
            MSG = opts.command

        sock.sendto(MSG.encode("utf8"), (opts.HOST, opts.PORT))
        received = sock.recv(BLOCK_SIZE)
        if opts.command != "":
            print_table(received.decode("utf8"))
        if opts.INTERVAL == 0 and opts.command == "":
            sock.sendto("list".encode("utf8"), (opts.HOST, opts.PORT))
            received = sock.recv(BLOCK_SIZE)
            print_table(received.decode("utf8"))
        errors = 0
    except socket.error:
        print("Didn't get reply from {0}:{1}".format(opts.HOST, opts.PORT))
        errors = 1
    except IndexError:
        print(received)
    if opts.INTERVAL == 0:
        break
    time.sleep(opts.INTERVAL)

sys.exit(errors)