# Path: tests/utils.py
# def ping(host_from, host_to, count=4, interval=0.2):
#     """
#     ICMP packets has to be sent by a process owned by the root, /bin/ping has
#     setuid so I'm using this to avoid juggling with root privileges.
#     This is blocking.
#     """
#     assert interval>=0.2, "Interval cannot be less than 0.2 seconds"
#     cmd = 'ping -c {count} -I {0} -i {interval} {1}'\
#                     .format(host_from, host_to, interval=interval, count=count)
#     p = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE,
#                             stderr=subprocess.PIPE)
#     stdout, stderr = p.communicate()
#     if p.returncode == 0:
#         return float(stdout.splitlines()[-1].split('/')[4])
#     raise PingError(p.returncode)
# 
# def random_file(size):
#     with open('/dev/urandom', 'rb') as f:
#         return bytearray(f.read(size))
# 
# def eta(filesize, a, b):
#         """
#         Determines how long the transfer of filesize bytes from A to B should take.
#         units for filesize are bytes, a and b is speed in kilobytes
#         """
#         return filesize/(min(a, b)*1024)
# 
# Path: tests/mixins.py
# class ServerMixin(object):
#     """
#     Expected members of a class this object is mixed into:
#     filesize: size of transferred file in bytes
#     server_addr: (address,port) the server should bind to
#     """
#     class Server(threading.Thread):
#         def __init__(self, server, filesize):
#             threading.Thread.__init__(self)
#             self.server_addr = server
#             self.filesize = filesize
#         
#         def run(self):
#             s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
#             s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
#             s.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 50688) # 49.5 KB
#             s.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 87380) # 85.3 KB
#             s.bind(self.server_addr)
#             s.listen(1)
#             conn, addr = s.accept()
#             data = bytearray()
#             while len(data) < self.filesize:
#                 data += conn.recv(1024)
# 
#             sent = 0
#             while sent < self.filesize:
#                 sent += conn.send(data[sent:sent+4096])
#             
#             conn.recv(1)
#             conn.close()
#             s.close()
#         
#     def run_server(self):
#         self.server_thread = ServerMixin.Server(self.server_addr, self.filesize)
#         self.server_thread.daemon = True
#         self.server_thread.start()
#         time.sleep(0.1) # wait for server to initialize
# 
# class ClientMixin(object):
#     """
#     Expected members of a class this object is mixed into:
#     randomfile: the array to be transmitted over the "virtual network"
#     filesize: size of transferred file in bytes
#     server_addr: (address,port) of a server the client will connect to
#     client_addr: (address,port) the client will bind to
#     """
#     
#     def make_transfer(self):
#         self.sock_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
#         # SO_REUSEADDR: http://stackoverflow.com/questions/3229860/what-is-the-meaning-of-so-reuseaddr-setsockopt-option-linux
#         s = self.sock_client
#         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
#         s.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 50688) # 49.5 KB
#         s.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 87380) # 85.3 KB
#         s.bind(self.client_addr)
#         s.connect(self.server_addr)
#         start = datetime.now()
#         
#         ## client -> server
#         sent = 0
#         while sent < self.filesize:
#             sent += s.send(self.randomfile[sent:sent+4096])
#         
#         # We have to wait until the server finishes reading data from its socket
#         # and closes the connection.
#         rcvd = s.recv(1)
#         time_up = total_seconds(datetime.now() - start)
#         
#         time.sleep(getattr(self, 'sleep', 0)) # wait for a bucket to fill again
#         
#         ## server -> client
#         start = datetime.now()
#         
#         while len(rcvd) < self.filesize:
#             rcvd += s.recv(1024)
#         
#         time_down = total_seconds(datetime.now() - start)
#         s.close()
#         return time_up, time_down


import unittest
from shapy.framework.qdisc import *
from shapy.framework.tcclass import *
from shapy.framework.filter import *
from shapy.framework.interface import *
from tests.utils import ping, random_file, eta
from tests.mixins import ServerMixin, ClientMixin



class TestHTB(unittest.TestCase, ServerMixin, ClientMixin):
    def setUp(self):
        self.server_addr = ('127.0.0.2', 55000)
        self.client_addr = ('127.0.0.3', 55001)
        self.server_speed = 256 # download speed in kbyte/s
        self.client_speed = 128 
        
        self.i = Interface('lo')
        try:
            self.i.teardown()
        except OSError:
            pass
        h1 = HTBQdisc('1:')#, default_class='1ff')
        h1.add( FlowFilter('dst 127.0.0.2', '1:2', prio=1) )
        h1.add( FlowFilter('dst 127.0.0.3', '1:1', prio=1) )
        h1.add( HTBClass('1:1', rate=self.server_speed, ceil=self.server_speed))
        h1.add( HTBClass('1:2', rate=self.client_speed, ceil=self.client_speed))
        self.i.add( h1 )
        self.i.set_shaping()
        
        self.filesize = 2**20

# Path: tests/utils.py
# def ping(host_from, host_to, count=4, interval=0.2):
#     """
#     ICMP packets has to be sent by a process owned by the root, /bin/ping has
#     setuid so I'm using this to avoid juggling with root privileges.
#     This is blocking.
#     """
#     assert interval>=0.2, "Interval cannot be less than 0.2 seconds"
#     cmd = 'ping -c {count} -I {0} -i {interval} {1}'\
#                     .format(host_from, host_to, interval=interval, count=count)
#     p = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE,
#                             stderr=subprocess.PIPE)
#     stdout, stderr = p.communicate()
#     if p.returncode == 0:
#         return float(stdout.splitlines()[-1].split('/')[4])
#     raise PingError(p.returncode)
# 
# def random_file(size):
#     with open('/dev/urandom', 'rb') as f:
#         return bytearray(f.read(size))
# 
# def eta(filesize, a, b):
#         """
#         Determines how long the transfer of filesize bytes from A to B should take.
#         units for filesize are bytes, a and b is speed in kilobytes
#         """
#         return filesize/(min(a, b)*1024)
# 
# Path: tests/mixins.py
# class ServerMixin(object):
#     """
#     Expected members of a class this object is mixed into:
#     filesize: size of transferred file in bytes
#     server_addr: (address,port) the server should bind to
#     """
#     class Server(threading.Thread):
#         def __init__(self, server, filesize):
#             threading.Thread.__init__(self)
#             self.server_addr = server
#             self.filesize = filesize
#         
#         def run(self):
#             s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
#             s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
#             s.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 50688) # 49.5 KB
#             s.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 87380) # 85.3 KB
#             s.bind(self.server_addr)
#             s.listen(1)
#             conn, addr = s.accept()
#             data = bytearray()
#             while len(data) < self.filesize:
#                 data += conn.recv(1024)
# 
#             sent = 0
#             while sent < self.filesize:
#                 sent += conn.send(data[sent:sent+4096])
#             
#             conn.recv(1)
#             conn.close()
#             s.close()
#         
#     def run_server(self):
#         self.server_thread = ServerMixin.Server(self.server_addr, self.filesize)
#         self.server_thread.daemon = True
#         self.server_thread.start()
#         time.sleep(0.1) # wait for server to initialize
# 
# class ClientMixin(object):
#     """
#     Expected members of a class this object is mixed into:
#     randomfile: the array to be transmitted over the "virtual network"
#     filesize: size of transferred file in bytes
#     server_addr: (address,port) of a server the client will connect to
#     client_addr: (address,port) the client will bind to
#     """
#     
#     def make_transfer(self):
#         self.sock_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
#         # SO_REUSEADDR: http://stackoverflow.com/questions/3229860/what-is-the-meaning-of-so-reuseaddr-setsockopt-option-linux
#         s = self.sock_client
#         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
#         s.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 50688) # 49.5 KB
#         s.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 87380) # 85.3 KB
#         s.bind(self.client_addr)
#         s.connect(self.server_addr)
#         start = datetime.now()
#         
#         ## client -> server
#         sent = 0
#         while sent < self.filesize:
#             sent += s.send(self.randomfile[sent:sent+4096])
#         
#         # We have to wait until the server finishes reading data from its socket
#         # and closes the connection.
#         rcvd = s.recv(1)
#         time_up = total_seconds(datetime.now() - start)
#         
#         time.sleep(getattr(self, 'sleep', 0)) # wait for a bucket to fill again
#         
#         ## server -> client
#         start = datetime.now()
#         
#         while len(rcvd) < self.filesize:
#             rcvd += s.recv(1024)
#         
#         time_down = total_seconds(datetime.now() - start)
#         s.close()
#         return time_up, time_down


# Path: tests/framework/test_htb.py
import unittest
from shapy.framework.qdisc import *
from shapy.framework.tcclass import *
from shapy.framework.filter import *
from shapy.framework.interface import *
from tests.utils import ping, random_file, eta
from tests.mixins import ServerMixin, ClientMixin




class TestHTB(unittest.TestCase, ServerMixin, ClientMixin):
    def setUp(self):
        self.server_addr = ('127.0.0.2', 55000)
        self.client_addr = ('127.0.0.3', 55001)
        self.server_speed = 256 # download speed in kbyte/s
        self.client_speed = 128 
        
        self.i = Interface('lo')
        try:
            self.i.teardown()
        except OSError:
            pass
        h1 = HTBQdisc('1:')#, default_class='1ff')
        h1.add( FlowFilter('dst 127.0.0.2', '1:2', prio=1) )
        h1.add( FlowFilter('dst 127.0.0.3', '1:1', prio=1) )
        h1.add( HTBClass('1:1', rate=self.server_speed, ceil=self.server_speed))
        h1.add( HTBClass('1:2', rate=self.client_speed, ceil=self.client_speed))
        self.i.add( h1 )
        self.i.set_shaping()
        
        self.filesize = 2**20

        self.randomfile = random_file(self.filesize)