#!/usr/bin/env python
# coding: utf-8

import sys
import telnetlib
import re
import time
import logging
import argparse

LOG = logging.getLogger('check-bgp-neighbor-state')

TELNET_TIMEOUT = 5


def setup_logger(logger, debug):
    if debug:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)
    handler = logging.StreamHandler(sys.stdout)
    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)


def get_bgpd_param(conf_path, keys):
    ret = {}
    with open(conf_path, "r") as conf:
        lines = conf.readlines()
        for line in lines:
            for key in keys:
                found = re.search("^%s " % key, line)
                if not found:
                    continue

                words = line.split()
                if key in ret:
                    ret[key].append(words[1])
                else:
                    ret[key] = [words[1]]

    return ret


def login_bgpd(host, port, password, prompt):
    telnet = telnetlib.Telnet(host, port)

    # login to bpgd
    telnet.read_until("Password: ", TELNET_TIMEOUT)
    telnet.write(password + "\n")
    telnet.read_until(prompt, TELNET_TIMEOUT)
    return telnet


def get_bgp_summary(telnet, prompt):
    telnet.write("show ip bgp summary\n")
    return telnet.read_until(prompt, TELNET_TIMEOUT)


def check_neighbor_state(summary, neighbors):
    """ Check if all neighbor's states are 1.

    Args:
        summary: output of command 'show ip bgp summary'. summary contains
            text like below:
            Neighbor        V         AS MsgRcvd MsgSent   TblVer  InQ OutQ Up/Down  State/PfxRcd
            169.254.24.245  4 10124       6       6        0    0    0 00:00:28        1
            169.254.26.49   4 10124       0       0        0    0    0 never    Connect
        neighbors: array of neighbors to be checked.

    Returns:
        If all neighbor's states are 1, return True.
    """
    for line in summary.split('\n'):
        for neighbor in neighbors:
            found = re.search("^%s " % neighbor, line)
            if not found:
                continue

            columns = line.split()
            if columns[9] == "1":
                neighbors.remove(neighbor)
                break

    return len(neighbors) == 0


def main():
    parser = argparse.ArgumentParser(description='''
This script checks if all BGP neighbor's status are established through bgpd
telnet interface. This script checks the status periodically until all status
become established or the number of retries has reached the limit.
If all status are established, the exit status is 0, otherwise 1.
''')
    parser.add_argument("-i", "--interval", default=5,
                        help="interval to check the status")
    parser.add_argument("-r", "--retry", default=12,
                        help="retry count to check the status")
    parser.add_argument("-H", "--host", default="localhost",
                        help="target host to telnet")
    parser.add_argument("-p", "--port", default=2605,
                        help="target port to telnet")
    parser.add_argument("-c", "--conf", default="/etc/quagga/bgpd.conf",
                        help="path to bpgd config")
    parser.add_argument("-d", "--debug", action='store_true', default=False)
    args = parser.parse_args()

    setup_logger(LOG, args.debug)

    # init parameter
    param = get_bgpd_param(args.conf, ["password", "hostname", "neighbor"])
    prompt = param["hostname"][0] + "> "

    telnet = login_bgpd(args.host, args.port, param["password"][0], prompt)

    retry = args.retry
    while retry >= 0:
        bgp_summary = get_bgp_summary(telnet, prompt)
        LOG.debug(bgp_summary)

        # delete duplication record
        neighbors = list(set(param["neighbor"]))
        ret = check_neighbor_state(bgp_summary, neighbors)
        if ret:
            LOG.info("All neighbors are established.")
            exit(0)

        LOG.info("%s is not established yet.", ",".join(neighbors))
        time.sleep(args.interval)
        retry -= 1

    exit(1)


if __name__ == '__main__':
    main()