list of sources

1 files changed, 110 insertions, 42 deletions

diff --git a/bandwidthmon b/bandwidthmon
index 0531012..63c37d2 100755
--- a/bandwidthmon
+++ b/bandwidthmon
@@ -5,23 +5,42 @@ import sys
 import glib
 import gtk
 import time
+import socket
+
+debug = True
+def dprint(x):
+    if not debug:
+        return
+    print x
 
 class Bandwidth(object):
     """ Compute average bandwidth over last X seconds
     """
 
-    def __init__(self, window_size):
-        self.total = 0
-        self.start_time = time.time()
-        self.last_time = self.start_time
-        self.window = []
+    def __init__(self, window_size, average_callback = None):
         self.window_size = window_size
+        self.average_callback = average_callback
+        self.reset()
 
     def add_packet(self, data):
         num_bytes = len(data)
         # TODO - other stuff then just length
         self.total += num_bytes
         now = time.time()
+        # add new datapoint
+        self.window.append((now, num_bytes))
+        self.remove_old(now)
+        glib.timeout_add_seconds(self.window_size, self.remove_old,
+                                 now + self.window_size + 0.0001)
+        return self.average()
+
+    def reset(self):
+        self.total = 0
+        self.start_time = time.time()
+        self.last_time = self.start_time
+        self.window = []
+
+    def remove_old(self, now):
         indices = []
         # remove too old data points
         for i, (t, x) in enumerate(self.window):
@@ -29,9 +48,11 @@ class Bandwidth(object):
                 indices.append(i)
         for i in reversed(indices):
             del self.window[i]
-        # add new datapoint
-        self.window.append((now, num_bytes))
-        return self.average()
+        self.on_average_updated()
+
+    def on_average_updated(self):
+        if self.average_callback:
+            self.average_callback(self)
 
     def average(self):
         # suboptimal
@@ -43,42 +64,76 @@ class Bandwidth(object):
         return float(self.total) / (self.last_time - self.start_time)
 
 class UI(object):
-    def __init__(self):
+    def __init__(self, pairs):
+        """
+             |        hbox
+        vbox | in_label out_label
+             ; reset_button
+        """
+        self.pairs = pairs
         self.window = gtk.Window()
-        self.hbox = gtk.HBox()
-        self.window.add(self.hbox)
-        self.label_avg_in = gtk.Label('0')
-        self.label_avg_out = gtk.Label('0')
-        self.hbox.add(self.label_avg_in)
-        self.hbox.add(self.label_avg_out)
+        self.vbox = gtk.VBox()
+        self.window.add(self.vbox)
+        self.port_in_bw = {}
+        self.port_out_bw = {}
+        self.bw_to_ui = {}
+        for src, dst in self.pairs:
+            self.add_in_out_pair(src, dst, self.vbox)
+        self.bw_reset_button = gtk.Button('reset')
+        self.bw_reset_button.connect("clicked", self.reset, None)
+        self.vbox.add(self.bw_reset_button)
         self.window.set_size_request(200, 80)
         self.window.show_all()
+    
+    def add_in_out_pair(self, src, dst, vbox):
+        #import pdb; pdb.set_trace()
+        hbox1 = gtk.HBox()
+        hbox2 = gtk.HBox()
+        vbox.add(hbox1)
+        vbox.add(hbox2)
+        label_avg_in = gtk.Label('0')
+        label_avg_out = gtk.Label('0')
+        label_avg_in_from_reset = gtk.Label('0')
+        label_avg_out_from_reset = gtk.Label('0')
+        hbox1.add(label_avg_in)
+        hbox1.add(label_avg_out)
+        hbox2.add(label_avg_in_from_reset)
+        hbox2.add(label_avg_out_from_reset)
+        in_bw = Bandwidth(1, self.on_bw_average)
+        out_bw = Bandwidth(1, self.on_bw_average)
+        self.port_in_bw[src] = in_bw
+        self.port_out_bw[src] = out_bw
+        self.port_in_bw[dst] = in_bw
+        self.port_out_bw[dst] = out_bw
+        # TODO - use a class
+        self.bw_to_ui[in_bw] = dict(label_avg=label_avg_in,
+            label_avg_from_reset=label_avg_in_from_reset)
+        self.bw_to_ui[out_bw] = dict(label_avg=label_avg_out,
+            label_avg_from_reset=label_avg_out_from_reset)
 
-    def set_average_in(self, avg):
-        self.label_avg_in.set_label(str(avg))
+    def reset(self, widget, data=None):
+        for bw in self.bw_to_ui.keys():
+            bw.reset()
 
-    def set_average_out(self, avg):
-        self.label_avg_out.set_label(str(avg))
+    def on_bw_average(self, bw):
+        labels = self.bw_to_ui[bw]
+        labels['label_avg'].set_label(str(bw.average()))
+        labels['label_avg_from_reset'].set_label(str(bw.global_average()))
 
-def main():
-    debug = True
-    debug_proxy = False
-    def dprint(x):
-        if not debug:
+    def add_packet(self, src, dst, data):
+        if src in self.port_in_bw:
+            bw = self.port_in_bw[src]
+        elif dst in self.port_out_bw:
+            bw = self.port_out_bw[dst]
+        else:
+            print "got packet for unmonitored src port %d (%d->%d %d#)" % (
+                src, src, dst, len(data))
             return
-        print x
-    p = argparse.ArgumentParser(description="proxy multiple socket connections")
-    p.add_argument('--listen-port', required=True, type=int)
-    p.add_argument('--remote-port', required=True, type=int)
-    p.add_argument('--remote-host', default='127.0.0.1')
-    p.add_argument('--listen-host', default='127.0.0.1')
-    args = p.parse_args(sys.argv[1:])
-    ui = UI()
-    in_bw = Bandwidth(1)
-    out_bw = Bandwidth(1)
+        bw.add_packet(data)
+
+def setup_proxy(ui, listen_port, listen_host, remote_addr, debug_proxy):
     iterate_packets, handle_input, select_based_iterator, get_fds = proxy.make_proxy(
-        args.listen_port, (args.remote_host, args.remote_port),
-        args.listen_host, debug=debug_proxy)
+        listen_port, remote_addr, listen_host, debug=debug_proxy)
     assert(len(get_fds()) == 1) # only the accepting socket
     accepter = get_fds()[0]
     added_fds = set()
@@ -114,12 +169,7 @@ def main():
         (src, dst, data, completer) = result
         completer()
         dprint("%s: result %d->%d #%d" % (fd.fileno(), src, dst, len(data)))
-        if dst == args.remote_port:
-            in_bw.add_packet(data)
-            ui.set_average_in(in_bw.average())
-        else:
-            out_bw.add_packet(data)
-            ui.set_average_out(out_bw.average())
+        ui.add_packet(src, dst, data)
         return True
     def update_fds():
         fds = get_fds()
@@ -130,6 +180,24 @@ def main():
             added_fds.add(fd)
             new_fds.append(fd)
     update_fds()
+
+def main():
+    debug_proxy = False
+    p = argparse.ArgumentParser(description="proxy multiple socket connections")
+    p.add_argument('--listen-port', required=True, type=int, action='append')
+    p.add_argument('--remote-port', required=True, type=int, action='append')
+    # only allow same host for all ports - easier to implement.
+    p.add_argument('--remote-host', default='127.0.0.1')
+    p.add_argument('--listen-host', default='127.0.0.1')
+    args = p.parse_args(sys.argv[1:])
+    if len(args.listen_port) != len(args.remote_port):
+        print "must supply same amount of listening and remote ports"
+        sys.exit(1)
+    ui = UI(zip(args.listen_port, args.remote_port))
+    for listen_port, remote_port in zip(args.listen_port, args.remote_port):
+        setup_proxy(ui=ui, listen_port=listen_port, listen_host=args.listen_host,
+                    remote_addr=(args.remote_host, remote_port),
+                    debug_proxy=debug_proxy)
     gtk.main()
 
 if __name__ == '__main__':