#!/usr/libexec/platform-python
#
# Copyright (c) 2022-2023 - Adjacent Link LLC, Bridgewater, New Jersey
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#  * Neither the name of Adjacent Link LLC nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# See toplevel COPYING for more information.
#

import sys
import select
import signal
import socket
import struct
import os
import logging
import re
import subprocess
import copy
import json
import ctypes
import fcntl

from argparse import ArgumentParser
from dataclasses import dataclass
from dataclasses import astuple
from dataclasses import asdict
from collections import deque
from collections import defaultdict
import daemon
import daemon.pidfile
import traceback

@dataclass
class ReceiveFlow:
    source: str
    flow: int
    destination: str
    packets: int = 0
    bytes_accum: int = 0
    latency_accum: int = 0
    dup_packets: int = 0
    dup_bytes_accum: int = 0
    min_sequence: int = -1
    max_sequence: int = -1

@dataclass
class TransmitFlow:
    source_port: int
    flow: int
    destination: str
    packets: int = 0
    bytes_accum: int = 0
    min_sequence: int = -1
    max_sequence: int = -1

class Connection(object):
    def __init__(self,sock,address,rx_udp,tx_udp,rx_tcp,tx_tcp):
        self.sock=sock
        self.address=address
        self.rx_udp = copy.deepcopy(rx_udp)
        self.tx_udp = copy.deepcopy(tx_udp)
        self.rx_tcp = copy.deepcopy(rx_tcp)
        self.rx_tcp = copy.deepcopy(tx_tcp)
        self.rx_sequence_cache = defaultdict(lambda : defaultdict(lambda : defaultdict(lambda : set())))
        self.tx_sequence_cache = defaultdict(lambda : defaultdict(lambda : defaultdict(lambda : set())))
        self.pending_command = ''

    def update(self,rx_udp,tx_udp,rx_tcp,tx_tcp):
        self.rx_udp = copy.deepcopy(rx_udp)
        self.tx_udp = copy.deepcopy(tx_udp)
        self.rx_tcp = copy.deepcopy(rx_tcp)
        self.rx_tcp = copy.deepcopy(tx_tcp)
        self.rx_sequence_cache = defaultdict(lambda : defaultdict(lambda : defaultdict(lambda : set())))
        self.tx_sequence_cache = defaultdict(lambda : defaultdict(lambda : defaultdict(lambda : set())))

argument_parser = ArgumentParser()

argument_parser.add_argument('listen-address',
                             type=str,
                             help='listen endpoint format <address>')

argument_parser.add_argument('listen-port',
                             type=int,
                             help='listen endpoint format <port>')

argument_parser.add_argument('mgen-output-file',
                            type=str,
                            help='MGEN output file.')

argument_parser.add_argument('--log-file',
                            type=str,
                            metavar='FILE',
                            help='log file.')

argument_parser.add_argument('--log-level',
                             type=str,
                             metavar='LEVEL',
                             choices=['critical',
                                      'error',
                                      'warning',
                                      'info',
                                      'debug',
                                      'notset'],
                             default='info',
                             help='log level [default: %(default)s].')

argument_parser.add_argument('--pid-file',
                            type=str,
                            default=None,
                            help='write pid file')

argument_parser.add_argument('--daemonize',
                            '-d',
                            action='store_true',
                            dest='daemonize',
                            default=False,
                            help='daemonize application [default: %(default)s]')

ns = argument_parser.parse_args()

args = vars(ns)

eventfd = None

def shutdown_handler(signum,frame):
    if eventfd != None:
        os.write(eventfd,struct.pack('Q',1))

def do_main():
    global eventfd

    logging.basicConfig(filename=args['log_file'],
                        format='%(asctime)s.%(msecs)03d %(levelname)s: %(message)s',
                        datefmt='%H:%M:%S',
                        level=getattr(logging,args['log_level'].upper()))

    pipefd_read, pipefd_write = os.pipe()

    fcntl.fcntl(pipefd_read,fcntl.F_SETFL,os.O_NONBLOCK)

    logging.info('tailing mgen output file: {}'.format(args['mgen-output-file']))

    p = subprocess.Popen(['tail',
                          '--retry',
                          '-q',
                          '-f',
                          args['mgen-output-file']],
                         bufsize=1048576,
                         stdout=pipefd_write,
                         stderr=pipefd_write)


    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR, 1)

    logging.info('listening on {}:{}'.format(args['listen-address'],
                                             args['listen-port']))

    sock.bind((args['listen-address'], args['listen-port']))

    poller = select.epoll()

    poller.register(sock.fileno(),select.EPOLLIN)

    poller.register(pipefd_read,select.EPOLLIN)

    libc = ctypes.CDLL(None)
    eventfd = libc.eventfd(ctypes.c_uint(0),ctypes.c_uint(0))

    poller.register(eventfd,select.EPOLLIN)

    sock.listen(1)


    log_regex = re.compile(r'(\d+:\d+:){0,1}(\d+.\d+) (SEND|RECV|LISTEN|START|STOP)( proto>(UDP|TCP) (flow>(\d+) seq>(\d+) ((src|srcPort)>([^\s]+) dst>([^\s]+))( sent>(\d+:\d+:){0,1}(\d+.\d+)){0,1} size>(\d+)|port>(\d+))){0,1}')

    rx_udp = {}
    tx_udp = {}
    rx_tcp = {}
    tx_tcp = {}
    seq_history = deque([],3)

    # drop oldest seq history set
    seq_history.append(set())
    clients = {}

    cancel = False
    buf = b''

    while not cancel:
        events = poller.poll()

        for fd, event in events:
            if fd == eventfd:
                cancel = True
                break

            elif fd == pipefd_read:
                buf += os.read(pipefd_read,256)

                eol = buf.find(b'\n')

                while eol >= 0:
                    line = buf[:eol]

                    if eol + 1 != len(buf):
                        buf = buf[eol + 1:]
                    else:
                        buf=b''

                    eol = buf.find(b'\n')

                    match = log_regex.match(line.decode('utf-8'))

                    if match:
                        log_timestamp = float(match.group(2))

                        if match.group(1):
                            hour,minute = [int(x) for x in match.group(1).split(':') if x]
                            log_timestamp += hour * 3600 + minute * 60

                        log_type = match.group(3)
                        proto = match.group(5)

                        if log_type == 'RECV':
                            tx_timestamp = float(match.group(15))

                            if match.group(14):
                                hour,minute = [int(x) for x in match.group(14).split(':') if x]
                                tx_timestamp += hour * 3600 + minute * 60

                            flow = int(match.group(7))
                            sequence = int(match.group(8))
                            size = int(match.group(16))
                            src = match.group(11).split('/')[0]
                            dst = match.group(12).split('/')[0]

                            if True:
                                if proto == 'UDP':
                                    store = rx_udp
                                else:
                                    store = rx_tcp

                                if src not in store:
                                    store[src] = {flow : ReceiveFlow(src,flow,dst)}
                                else:
                                    if flow not in store[src]:
                                        store[src][flow] = ReceiveFlow(src,flow,dst)

                                dup_found = False
                                for s in seq_history:
                                    if (src,flow,sequence) in s:
                                        dup_found = True
                                        break

                                if dup_found:
                                    store[src][flow].dup_packets += 1
                                    store[src][flow].dup_bytes_accum += size
                                else:
                                    store[src][flow].packets += 1
                                    store[src][flow].bytes_accum += size
                                    store[src][flow].latency_accum += log_timestamp - tx_timestamp

                                    # drop the oldest set of seq so we
                                    # look amoung the last 100 to 149
                                    # messages
                                    if len(seq_history[-1]) >= 50:
                                        seq_history.append(set())

                                    seq_history[-1].add((src,flow,sequence))

                                    for fd in clients:
                                        clients[fd].rx_sequence_cache[proto][src][flow].add(sequence)

                        elif log_type == 'SEND':
                            flow = int(match.group(7))
                            sequence = int(match.group(8))
                            size = int(match.group(16))

                            src_port = int(match.group(11))
                            dst = match.group(12).split('/')[0]

                            if True:
                                if proto == 'UDP':
                                    store = tx_udp
                                else:
                                    store = tx_tcp

                                if src_port not in store:
                                    store[src_port] = {flow : TransmitFlow(src_port,flow,dst)}
                                else:
                                    if flow not in store[src_port]:
                                        store[src_port][flow] = TransmitFlow(src_port,flow,dst)

                                store[src_port][flow].packets += 1
                                store[src_port][flow].bytes_accum += size

                                for fd in clients:
                                   clients[fd].tx_sequence_cache[proto][src_port][flow].add(sequence)
                        else:
                            logging.debug('ignoring: {}'.format(line.decode('utf-8')))

                    else:
                        logging.debug('not matching: {}'.format(line.decode('utf-8')))


            elif fd == sock.fileno():
                remote,address = sock.accept()
                poller.register(remote.fileno(),select.EPOLLIN)
                clients[remote.fileno()] = Connection(remote,
                                                      address,
                                                      rx_udp,
                                                      tx_udp,
                                                      rx_tcp,
                                                      tx_tcp)

                logging.info('new connection from {} {} {} {} {}'.format(address,
                                                                         rx_udp,
                                                                         tx_udp,
                                                                         rx_tcp,
                                                                         tx_tcp))


            elif fd in clients:
                command = clients[fd].sock.recv(65535)

                response = {'tx_udp' : [],
                            'tx_tcp' : [],
                            'rx_udp' : [],
                            'rx_tcp' : []}

                if len(command) and \
                   len(clients[fd].pending_command) + len('status\n') <= len('status\n'):

                    clients[fd].pending_command += command.decode()

                    if(clients[fd].pending_command == 'status\n'):

                        logging.debug('status from {}'.format(clients[fd].address))

                        for src in tx_udp:
                            for flow in tx_udp[src]:
                                tf = TransmitFlow(*astuple(tx_udp[src][flow]))

                                if src in clients[fd].tx_udp:
                                    if flow in clients[fd].tx_udp[src]:
                                        tf.packets = tf.packets - clients[fd].tx_udp[src][flow].packets
                                        tf.bytes_accum = tf.bytes_accum - clients[fd].tx_udp[src][flow].bytes_accum

                                if clients[fd].tx_sequence_cache['UDP'][src][flow]:
                                    tf.min_sequence = min(clients[fd].tx_sequence_cache['UDP'][src][flow])
                                    tf.max_sequence = max(clients[fd].tx_sequence_cache['UDP'][src][flow])

                                response['tx_udp'].append(asdict(tf))

                        for src in tx_tcp:
                            for flow in tx_tcp[src]:
                                tf = TransmitFlow(*astuple(tx_tcp[src][flow]))

                                if src in clients[fd].tx_tcp:
                                    if flow in clients[fd].tx_tcp[src]:
                                        tf.packets = tf.packets - clients[fd].tx_tcp[src][flow].packets
                                        tf.bytes_accum = tf.bytes_accum - clients[fd].tx_tcp[src][flow].bytes_accum

                                if clients[fd].tx_sequence_cache['TCP'][src][flow]:
                                    tf.min_sequence = min(clients[fd].tx_sequence_cache['TCP'][src][flow])
                                    tf.max_sequence = max(clients[fd].tx_sequence_cache['TCP'][src][flow])

                                response['tx_tcp'].append(asdict(tf))

                        for src in rx_udp:
                            for flow in rx_udp[src]:
                                rf = ReceiveFlow(*astuple(rx_udp[src][flow]))

                                if src in clients[fd].rx_udp:
                                    if flow in clients[fd].rx_udp[src]:
                                        rf.packets = rf.packets - clients[fd].rx_udp[src][flow].packets
                                        rf.bytes_accum = rf.bytes_accum - clients[fd].rx_udp[src][flow].bytes_accum
                                        rf.latency_accum = rf.latency_accum - clients[fd].rx_udp[src][flow].latency_accum
                                        rf.dup_packets = rf.dup_packets - clients[fd].rx_udp[src][flow].dup_packets
                                        rf.dup_bytes_accum = rf.dup_bytes_accum - clients[fd].rx_udp[src][flow].dup_bytes_accum

                                if clients[fd].rx_sequence_cache['UDP'][src][flow]:
                                    rf.min_sequence = min(clients[fd].rx_sequence_cache['UDP'][src][flow])
                                    rf.max_sequence = max(clients[fd].rx_sequence_cache['UDP'][src][flow])

                                response['rx_udp'].append(asdict(rf))

                        for src in rx_tcp:
                            for flow in rx_tcp[src]:
                                rf = ReceiveFlow(*astuple(rx_tcp[src][flow]))

                                if src in clients[fd].rx_tcp:
                                    if flow in clients[fd].rx_tcp[src]:
                                        rf.packets = rf.packets - clients[fd].rx_tcp[src][flow].packets
                                        rf.bytes_accum = rf.bytes_accum - clients[fd].rx_tcp[src][flow].bytes_accum
                                        rf.latency_accum = rf.latency_accum - clients[fd].rx_tcp[src][flow].latency_accum
                                        rf.dup_packets = rf.dup_packets - clients[fd].rx_tcp[src][flow].dup_packets
                                        rf.dup_bytes_accum = rf.dup_bytes_accum - clients[fd].rx_tcp[src][flow].dup_bytes_accum

                                if clients[fd].rx_sequence_cache['TCP'][src][flow]:
                                    rf.min_sequence = min(clients[fd].rx_sequence_cache['TCP'][src][flow])
                                    rf.max_sequence = max(clients[fd].rx_sequence_cache['TCP'][src][flow])

                                response['rx_tcp'].append(asdict(rf))

                        msg = json.dumps(response).encode()

                        os.write(fd,struct.pack('!I',len(msg)) + msg)

                        clients[fd].update(rx_udp,
                                           tx_udp,
                                           rx_tcp,
                                           tx_tcp)

                        clients[fd].pending_command = ''
                else:
                     poller.unregister(fd)
                     clients[fd].sock.close()
                     logging.info('connection to {} terminated'.format(clients[fd].address))
                     del clients[fd]

    logging.info('shutting down')

pid_file_context = None

if args['pid_file'] != None:
    pid_file_context = daemon.pidfile.PIDLockFile(args['pid_file'])

if args['log_file'] != None:
    args['log_file'] = os.path.realpath(args['log_file'])

with daemon.DaemonContext(pidfile=pid_file_context,
                          detach_process=args['daemonize'],
                          signal_map={signal.SIGINT:shutdown_handler,
                                      signal.SIGQUIT:shutdown_handler},
                          stdout=None if args['daemonize'] else sys.stdout,
                          stderr=None if args['daemonize'] else sys.stderr):
    try:
        do_main()
    except:
        logging.error(traceback.format_exc())
