#
# Copyright (C) 2018 Andreas Gerstmayr <andreas@gerstmayr.me>
# Based on the tcpretrans BCC tool by Brendan Gregg:
# https://github.com/iovisor/bcc/blob/master/tools/tcpretrans.py
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
""" PCP BCC PMDA tcpretrans module """

# pylint: disable=invalid-name, too-few-public-methods, too-many-instance-attributes

import ctypes as ct
from socket import inet_ntop, AF_INET, AF_INET6
from struct import pack

from bcc import BPF

from pcp.pmapi import pmUnits
from cpmapi import PM_TYPE_U64, PM_SEM_COUNTER, PM_COUNT_ONE
from cpmda import PMDA_FETCH_NOVALUES

from modules.pcpbcc import PCPBCCBase

#
# BPF program
#
bpf_src = "modules/tcpretrans_count.bpf"

#
# PCP BCC PMDA constants
#
MODULE = 'tcpretrans'
BASENS = 'io.net.tcp.retrans.'
units_none = pmUnits(0, 0, 0, 0, 0, 0)
units_count = pmUnits(0, 0, 1, 0, 0, PM_COUNT_ONE)

class Data_ipv4(ct.Structure):
    """ IPv4 data struct """
    _fields_ = [
        ("pid", ct.c_ulonglong),
        ("ip", ct.c_ulonglong),
        ("saddr", ct.c_ulonglong),
        ("daddr", ct.c_ulonglong),
        ("lport", ct.c_ulonglong),
        ("dport", ct.c_ulonglong),
        ("state", ct.c_ulonglong),
        ("type", ct.c_ulonglong)
    ]

class Data_ipv6(ct.Structure):
    """ IPv6 data struct """
    _fields_ = [
        ("pid", ct.c_ulonglong),
        ("ip", ct.c_ulonglong),
        ("saddr", (ct.c_ulonglong * 2)),
        ("daddr", (ct.c_ulonglong * 2)),
        ("lport", ct.c_ulonglong),
        ("dport", ct.c_ulonglong),
        ("state", ct.c_ulonglong),
        ("type", ct.c_ulonglong)
    ]

#
# PCP BCC Module
#
class PCPBCCModule(PCPBCCBase):
    """ PCP BCC tcpretrans module """
    def __init__(self, config, log, err, _):
        """ Constructor """
        PCPBCCBase.__init__(self, MODULE, config, log, err)

        self.include_tlp = False
        self.flow_count = 20
        self.buffer_page_count = 64

        for opt in self.config.options(MODULE):
            if opt == 'include_tlp':
                self.include_tlp = self.config.getboolean(MODULE, opt)
            if opt == 'flow_count':
                self.flow_count = int(self.config.get(MODULE, opt))
            if opt == 'buffer_page_count':
                self.buffer_page_count = int(self.config.get(MODULE, opt))
                if not self.buffer_page_count or \
                   self.buffer_page_count & (self.buffer_page_count - 1):
                    raise RuntimeError("Buffer page count is not power of two.")

        self.cache = {}
        self.insts = {}

        self.log("Initialized.")

    def metrics(self):
        """ Get metric definitions """
        name = BASENS
        self.items = (
            # Name - reserved - type - semantics - units - help
            (name + 'count', None, PM_TYPE_U64, PM_SEM_COUNTER, units_count, 'retransmits'),
        )
        return True, self.items

    def compile(self):
        """ Compile BPF """
        try:
            self.bpf = BPF(src_file=bpf_src)
            self.bpf.attach_kprobe(event="tcp_retransmit_skb", fn_name="trace_retransmit")
            if self.include_tlp:
                self.bpf.attach_kprobe(event="tcp_send_loss_probe", fn_name="trace_tlp")
            self.log("Compiled.")
        except Exception as error: # pylint: disable=broad-except
            self.bpf = None
            self.err(str(error))
            self.err("Module NOT active!")
            raise

    def cumulate_count_table(self, bcc_table, ip_version):
        """ Cumulate counts from BPF tables """
        counts = self.bpf[bcc_table]

        for k, v in counts.items():
            if ip_version == 4:
                saddr = inet_ntop(AF_INET, pack('I', k.saddr))
                daddr = inet_ntop(AF_INET, pack('I', k.daddr))
            else:
                saddr = inet_ntop(AF_INET6, k.saddr)
                daddr = inet_ntop(AF_INET6, k.daddr)

            key = '%s:%d::%s:%d' % (saddr, k.lport, daddr, k.dport)
            value = v.value if key not in self.cache else v.value + self.cache[key]
            self.cache[key] = value
            self.insts[key] = ct.c_int(1)

        counts.clear()

    def refresh(self):
        """ Refresh BPF data """
        if self.bpf is None:
            return None

        self.cumulate_count_table("ipv4_count", 4)
        self.cumulate_count_table("ipv6_count", 6)

        return self.insts

    def bpfdata(self, item, inst):
        """ Return BPF data as PCP metric value """
        try:
            key = self.pmdaIndom.inst_name_lookup(inst)
            return [self.cache[key], 1]
        except Exception: # pylint: disable=broad-except
            return [PMDA_FETCH_NOVALUES, 0]