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

# This example shows how a combination of BPF programs can be used to perform
# per-IP classification and rate limiting. The simulation in this example
# shows an example where N+M devices are combined and use 1 WAN. Traffic sent
# from/to the "neighbor" devices have their combined bandwidth capped at
# 128kbit, and the rest of the traffic can use an additional 1Mbit.

# This works by sharing a map between various tc ingress filters, each with
# a related set of bpf functions attached. The map stores a list of dynamically
# learned ip addresses that were seen on the neighbor devices and should be
# throttled.

#                          /------------\                        |
# neigh1 --|->->->->->->->-|            |                        |
# neigh2 --|->->->->->->->-|    <-128kb-|        /------\        |
# neigh3 --|->->->->->->->-|            |  wan0  | wan  |        |
#          | ^             |   br100    |-<-<-<--| sim  |        |
#          | clsfy_neigh() |            |   ^    \------/        |
# lan1 ----|->->->->->->->-|    <--1Mb--|   |                    |
# lan2 ----|->->->->->->->-|            |   classify_wan()       |
#            ^             \------------/                        |
#            pass()                                              |


from bpf import BPF
from pyroute2 import IPRoute, NetNS, IPDB, NSPopen
from simulation import Simulation
import sys
from time import sleep

ipr = IPRoute()
ipdb = IPDB(nl=ipr)
b = BPF(src_file="tc_neighbor_sharing.c", debug=0)

wan_fn = b.load_func("classify_wan", BPF.SCHED_CLS)
pass_fn = b.load_func("pass", BPF.SCHED_CLS)
neighbor_fn = b.load_func("classify_neighbor", BPF.SCHED_CLS)

num_neighbors = 3
num_locals = 2

# class to build the simulation network
class SharedNetSimulation(Simulation):

    def __init__(self, ipdb):
        super(SharedNetSimulation, self).__init__(ipdb)

        # Create the wan namespace, and attach an ingress filter for throttling
        # inbound (download) traffic
        wan_if = self._create_ns("wan0", ipaddr="172.16.1.5/24")[1]
        ipr.tc("add", "ingress", wan_if["index"], "ffff:")
        ipr.tc("add-filter", "bpf", wan_if["index"], ":1", fd=wan_fn.fd,
               prio=1, name=wan_fn.name, parent="ffff:", action="drop",
               classid=1, rate="128kbit", burst=1024 * 32, mtu=16 * 1024)
        ipr.tc("add-filter", "bpf", wan_if["index"], ":2", fd=pass_fn.fd,
               prio=2, name=pass_fn.name, parent="ffff:", action="drop",
               classid=2, rate="1024kbit", burst=1024 * 32, mtu=16 * 1024)
        self.wan_if = wan_if

    # start the namespaces that compose the network, interconnect them with the
    # bridge, and attach the tc filters
    def start(self):
        neighbor_list = []
        local_list = []
        cmd = ["netserver", "-D"]
        for i in range(0, num_neighbors):
            ipaddr = "172.16.1.%d/24" % (i + 100)
            ret = self._create_ns("neighbor%d" % i, ipaddr=ipaddr,
                                  fn=neighbor_fn, cmd=cmd)
            neighbor_list.append(ret)
        for i in range(0, num_locals):
            ipaddr = "172.16.1.%d/24" % (i + 150)
            ret = self._create_ns("local%d" % i, ipaddr=ipaddr,
                                  fn=pass_fn, cmd=cmd)
            local_list.append(ret)

        with ipdb.create(ifname="br100", kind="bridge") as br100:
            for x in neighbor_list:
                br100.add_port(x[1])
            for x in local_list:
                br100.add_port(x[1])
            br100.add_port(self.wan_if)
            br100.up()

try:
    sim = SharedNetSimulation(ipdb)
    sim.start()
    print("Network ready. Create a shell in the wan0 namespace and test with netperf")
    print("   (Neighbors are 172.16.1.100-%d, and LAN clients are 172.16.1.150-%d)"
            % (100 + num_neighbors - 1, 150 + num_locals - 1))
    print(" e.g.: ip netns exec wan0 netperf -H 172.16.1.100 -l 2")
    raw_input("Press enter when finished: ")
finally:
    if "sim" in locals(): sim.release()
    if "br100" in ipdb.interfaces: ipdb.interfaces.br100.remove().commit()
    ipdb.release()



@@ -1540,16 +1540,43 @@
 rt sleep
+%0Afrom builtins import input
 %0A%0Aipr = 
@@ -4099,12 +4099,8 @@
     
-raw_
 inpu

use builtins.input for backward/forward compability