#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from pyroute2.common import map_namespace
from pyroute2.netlink import genlmsg
from pyroute2.netlink.generic import GenericNetlinkSocket
from pyroute2.netlink.nlsocket import Marshal
from pyroute2.netlink import NLM_F_REQUEST
from pyroute2.netlink import NLM_F_DUMP
from socket import if_nametoindex
from socket import if_indextoname


BATADV_CMD_UNSPEC = 0
BATADV_CMD_GET_MESH_INFO = 1
BATADV_CMD_TP_METER = 2
BATADV_CMD_TP_METER_CANCEL = 3
BATADV_CMD_GET_ROUTING_ALGOS = 4
BATADV_CMD_GET_HARDIFS = 5
BATADV_CMD_GET_TRANSTABLE_LOCAL = 6
BATADV_CMD_GET_TRANSTABLE_GLOBAL = 7
BATADV_CMD_GET_ORIGINATORS = 8
BATADV_CMD_GET_NEIGHBORS = 9
BATADV_CMD_GET_GATEWAYS = 10
BATADV_CMD_GET_BLA_CLAIM = 11
BATADV_CMD_GET_BLA_BACKBONE = 12
BATADV_CMD_GET_DAT_CACHE = 13
BATADV_CMD_GET_MCAST_FLAGS = 14
BATADV_CMD_MAX = BATADV_CMD_GET_MCAST_FLAGS

(BATADV_NAMES, BATADV_VALUES) = map_namespace('BATADV_CMD_', globals())


class batadvcmd(genlmsg):
    prefix = 'BATADV_ATTR_'

    '''
    Message class that will be used to communicate
    with the kernel module
    '''
    nla_map = (('BATADV_ATTR_UNSPEC', 'none'),
               ('BATADV_ATTR_VERSION', 'asciiz'),
               ('BATADV_ATTR_ALGO_NAME', 'asciiz'),
               ('BATADV_ATTR_MESH_IFINDEX', 'uint32'),
               ('BATADV_ATTR_MESH_IFNAME', 'asciiz'),
               ('BATADV_ATTR_MESH_ADDRESS', 'l2addr'),
               ('BATADV_ATTR_HARD_IFINDEX', 'uint32'),
               ('BATADV_ATTR_HARD_IFNAME', 'asciiz'),
               ('BATADV_ATTR_HARD_ADDRESS', 'l2addr'),
               ('BATADV_ATTR_ORIG_ADDRESS', 'l2addr'),
               ('BATADV_ATTR_TPMETER_RESULT', 'uint8'),
               ('BATADV_ATTR_TPMETER_TEST_TIME', 'uint32'),
               ('BATADV_ATTR_TPMETER_BYTES', 'uint64'),
               ('BATADV_ATTR_TPMETER_COOKIE', 'uint32'),
               ('BATADV_ATTR_PAD', 'none'),
               ('BATADV_ATTR_ACTIVE', 'flag'),
               ('BATADV_ATTR_TT_ADDRESS', 'l2addr'),
               ('BATADV_ATTR_TT_TTVN', 'uint8'),
               ('BATADV_ATTR_TT_LAST_TTVN', 'uint8'),
               ('BATADV_ATTR_TT_CRC32', 'uint32'),
               ('BATADV_ATTR_TT_VID', 'uint16'),
               ('BATADV_ATTR_TT_FLAGS', 'uint32'),
               ('BATADV_ATTR_FLAG_BEST', 'flag'),
               ('BATADV_ATTR_LAST_SEEN_MSECS', 'uint32'),
               ('BATADV_ATTR_NEIGH_ADDRESS', 'l2addr'),
               ('BATADV_ATTR_TQ', 'uint8'),
               ('BATADV_ATTR_THROUGHPUT', 'uint32'),
               ('BATADV_ATTR_BANDWIDTH_UP', 'uint32'),
               ('BATADV_ATTR_BANDWIDTH_DOWN', 'uint32'),
               ('BATADV_ATTR_ROUTER', 'l2addr'),
               ('BATADV_ATTR_BLA_OWN', 'flag'),
               ('BATADV_ATTR_BLA_ADDRESS', 'l2addr'),
               ('BATADV_ATTR_BLA_VID', 'uint16'),
               ('BATADV_ATTR_BLA_BACKBONE', 'l2addr'),
               ('BATADV_ATTR_BLA_CRC', 'uint16'),
               ('BATADV_ATTR_DAT_CACHE_IP4ADDRESS', 'ip4addr'),
               ('BATADV_ATTR_DAT_CACHE_HWADDRESS', 'l2addr'),
               ('BATADV_ATTR_DAT_CACHE_VID', 'uint16'),
               ('BATADV_ATTR_MCAST_FLAGS', 'uint32'),
               ('BATADV_ATTR_MCAST_FLAGS_PRIV', 'uint32'))


class MarshalBatadv(Marshal):
    msg_map = {BATADV_CMD_UNSPEC: batadvcmd,
               BATADV_CMD_GET_MESH_INFO: batadvcmd,
               BATADV_CMD_TP_METER: batadvcmd,
               BATADV_CMD_TP_METER_CANCEL: batadvcmd,
               BATADV_CMD_GET_ROUTING_ALGOS: batadvcmd,
               BATADV_CMD_GET_HARDIFS: batadvcmd,
               BATADV_CMD_GET_TRANSTABLE_LOCAL: batadvcmd,
               BATADV_CMD_GET_TRANSTABLE_GLOBAL: batadvcmd,
               BATADV_CMD_GET_ORIGINATORS: batadvcmd,
               BATADV_CMD_GET_NEIGHBORS: batadvcmd,
               BATADV_CMD_GET_GATEWAYS: batadvcmd,
               BATADV_CMD_GET_BLA_CLAIM: batadvcmd,
               BATADV_CMD_GET_BLA_BACKBONE: batadvcmd,
               BATADV_CMD_GET_DAT_CACHE: batadvcmd,
               BATADV_CMD_GET_MCAST_FLAGS: batadvcmd}


class BatadvSocket(GenericNetlinkSocket):
    def __init__(self):
        GenericNetlinkSocket.__init__(self)
        self.marshal = MarshalBatadv()

    def bind(self, groups=0, **kwarg):
        GenericNetlinkSocket.bind(self, 'batadv', batadvcmd,
                                  groups, None, **kwarg)

class Batadv(BatadvSocket):

    __meshif = 'bat0'
    __meshif_index = 0

    def __init__(self, *argv, **kwarg):
        # continue with init
        super(Batadv, self).__init__(*argv, **kwarg)

        if 'meshif' in kwarg:
            self.__meshif = kwarg.pop('meshif')

        self.__meshif_index = if_nametoindex(self.__meshif)

        try:
            self.bind()
        except:
            super(Batadv, self).close()
            raise

    def get_neighbors(self):
        msg = batadvcmd()
        msg['cmd'] = BATADV_NAMES['BATADV_CMD_GET_NEIGHBORS']
        msg['attrs'] = [
            ['BATADV_ATTR_MESH_IFINDEX', self.__meshif_index],
        ]

        return self.nlm_request(msg,
                                msg_type=self.prid,
                                msg_flags=NLM_F_REQUEST | NLM_F_DUMP)

    def get_originators(self):
        msg = batadvcmd()
        msg['cmd'] = BATADV_NAMES['BATADV_CMD_GET_ORIGINATORS']
        msg['attrs'] = [
            ['BATADV_ATTR_MESH_IFINDEX', self.__meshif_index],
        ]

        return self.nlm_request(msg,
                                msg_type=self.prid,
                                msg_flags=NLM_F_REQUEST | NLM_F_DUMP)

def main():
    batadv = Batadv()

    orig2tq = {}
    for q in batadv.get_originators():
        tq = q.get_attr('BATADV_ATTR_TQ')
        if not tq:
            continue

        neigh = q.get_attr('BATADV_ATTR_NEIGH_ADDRESS')
        if not neigh:
            continue

        orig = q.get_attr('BATADV_ATTR_ORIG_ADDRESS')
        if not orig:
            continue

        hardif = q.get_attr('BATADV_ATTR_HARD_IFINDEX')
        if not hardif:
            continue

        if orig != neigh:
            continue

        orig2tq[(neigh, hardif)] = tq

    print("IF             Neighbor              last-seen	TQ")

    # this is not really necessary when we actually have done the orig == neigh
    # check before - so printing orig2tq is nearly enough
    for q in batadv.get_neighbors():
        tq = '?'

        hardif = q.get_attr('BATADV_ATTR_HARD_IFINDEX')
        if not hardif:
            continue

        last_seen = q.get_attr('BATADV_ATTR_LAST_SEEN_MSECS')
        if not last_seen:
            continue

        neigh = q.get_attr('BATADV_ATTR_NEIGH_ADDRESS')
        if not neigh:
            continue

        if (neigh, hardif) in orig2tq:
            tq = str(orig2tq[(neigh, hardif)])

        ifname = if_indextoname(hardif)
        print("   %10s	  %s %4i.%03is	%s" % (ifname, neigh,
                                          last_seen / 1000,
                                          last_seen % 1000, tq))

    batadv.close()

if __name__ == "__main__":
    main()