#!/usr/bin/env python
# Copyright (c) PLUMgrid, Inc.
# Licensed under the Apache License, Version 2.0 (the "License")

# test program to count the packets sent to a device in a .5
# second period

from ctypes import c_uint, c_ulong, Structure
from netaddr import IPAddress
from bcc import BPF
from subprocess import check_call
import sys
from unittest import main, TestCase

arg1 = sys.argv.pop(1)
arg2 = ""
if len(sys.argv) > 1:
  arg2 = sys.argv.pop(1)

Key = None
Leaf = None
if arg1.endswith(".b"):
    class Key(Structure):
        _fields_ = [("dip", c_uint),
                    ("sip", c_uint)]
    class Leaf(Structure):
        _fields_ = [("rx_pkts", c_ulong),
                    ("tx_pkts", c_ulong)]

class TestBPFSocket(TestCase):
    def setUp(self):
        b = BPF(arg1, arg2, debug=0)
        fn = b.load_func("on_packet", BPF.SOCKET_FILTER)
        BPF.attach_raw_socket(fn, "eth0")
        self.stats = b.get_table("stats", Key, Leaf)

    def test_ping(self):
        cmd = ["ping", "-f", "-c", "100", "172.16.1.1"]
        check_call(cmd)
        #for key, leaf in self.stats.items():
        #    print(IPAddress(key.sip), "=>", IPAddress(key.dip),
        #          "rx", leaf.rx_pkts, "tx", leaf.tx_pkts)
        key = self.stats.Key(IPAddress("172.16.1.2").value, IPAddress("172.16.1.1").value)
        leaf = self.stats[key]
        self.assertEqual(leaf.rx_pkts, 100)
        self.assertEqual(leaf.tx_pkts, 100)
        del self.stats[key]
        with self.assertRaises(KeyError):
            x = self.stats[key]
        with self.assertRaises(KeyError):
            del self.stats[key]
        self.stats.clear()
        self.assertEqual(len(self.stats), 0)
        self.stats[key] = leaf
        self.assertEqual(len(self.stats), 1)
        self.stats.clear()
        self.assertEqual(len(self.stats), 0)

    def test_empty_key(self):
        # test with a 0 key
        self.stats.clear()
        self.stats[self.stats.Key()] = self.stats.Leaf(100, 200)
        x = self.stats.popitem()
        self.stats[self.stats.Key(10, 20)] = self.stats.Leaf(300, 400)
        with self.assertRaises(KeyError):
            x = self.stats[self.stats.Key()]
        (_, x) = self.stats.popitem()
        self.assertEqual(x.rx_pkts, 300)
        self.assertEqual(x.tx_pkts, 400)
        self.stats.clear()
        self.assertEqual(len(self.stats), 0)
        self.stats[self.stats.Key()] = x
        self.stats[self.stats.Key(0, 1)] = x
        self.stats[self.stats.Key(0, 2)] = x
        self.stats[self.stats.Key(0, 3)] = x
        self.assertEqual(len(self.stats), 4)

if __name__ == "__main__":
    main()