#!/usr/bin/env python

# Copyright (c) PLUMgrid, Inc.

# Licensed under the Apache License, Version 2.0 (the "License")

from netaddr import IPAddress

from bcc import BPF

from pyroute2 import IPRoute, protocols

from socket import socket, AF_INET, SOCK_DGRAM

from subprocess import call

import sys

from time import sleep

from unittest import main, TestCase

arg1 = sys.argv.pop(1)

arg2 = ""

if len(sys.argv) > 1:

arg2 = sys.argv.pop(1)

class TestBPFFilter(TestCase):

def setUp(self):

b = BPF(arg1, arg2, debug=0)

fn = b.load_func("on_packet", BPF.SCHED_ACT)

ip = IPRoute()

ifindex = ip.link_lookup(ifname="eth0")[0]

# set up a network to change the flow:

# outside | inside

# 172.16.1.1 - 172.16.1.2 | 192.168.1.1 - 192.16.1.2

ip.addr("del", index=ifindex, address="172.16.1.2", mask=24)

ip.addr("add", index=ifindex, address="192.168.1.2", mask=24)

# add an ingress and egress qdisc

ip.tc("add", "ingress", ifindex, "ffff:")

ip.tc("add", "sfq", ifindex, "1:")

# add same program to both ingress/egress, so pkt is translated in both directions

action = {"kind": "bpf", "fd": fn.fd, "name": fn.name, "action": "ok"}

ip.tc("add-filter", "u32", ifindex, ":1", parent="ffff:", action=[action],

protocol=protocols.ETH_P_ALL, classid=1, target=0x10002, keys=['0x0/0x0+0'])

ip.tc("add-filter", "u32", ifindex, ":2", parent="1:", action=[action],

protocol=protocols.ETH_P_ALL, classid=1, target=0x10002, keys=['0x0/0x0+0'])

self.xlate = b.get_table("xlate")

def test_xlate(self):

key1 = self.xlate.Key(IPAddress("172.16.1.2").value, IPAddress("172.16.1.1").value)

leaf1 = self.xlate.Leaf(IPAddress("192.168.1.2").value, IPAddress("192.168.1.1").value, 0, 0)

self.xlate[key1] = leaf1

key2 = self.xlate.Key(IPAddress("192.168.1.1").value, IPAddress("192.168.1.2").value)

leaf2 = self.xlate.Leaf(IPAddress("172.16.1.1").value, IPAddress("172.16.1.2").value, 0, 0)

self.xlate[key2] = leaf2

call(["ping", "-c1", "192.168.1.1"])

leaf = self.xlate[key1]

self.assertGreater(leaf.ip_xlated_pkts, 0)

self.assertGreater(leaf.arp_xlated_pkts, 0)

leaf = self.xlate[key2]

self.assertGreater(leaf.ip_xlated_pkts, 0)

self.assertGreater(leaf.arp_xlated_pkts, 0)

if __name__ == "__main__":

main()