basic ui with labels

1 files changed, 63 insertions, 0 deletions

diff --git a/bandwidthmon.py b/bandwidthmon.py
index c99aaac..3b45f4a 100755
--- a/bandwidthmon.py
+++ b/bandwidthmon.py
@@ -4,6 +4,60 @@ import argparse
 import sys
 import glib
 import gtk
+import time
+
+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 = []
+        self.window_size = window_size
+
+    def add_packet(self, data):
+        num_bytes = len(data)
+        # TODO - other stuff then just length
+        self.total += num_bytes
+        now = time.time()
+        indices = []
+        # remove too old data points
+        for i, (t, x) in enumerate(self.window):
+            if t + self.window_size < now:
+                indices.append(i)
+        for i in reversed(indices):
+            del self.window[i]
+        # add new datapoint
+        self.window.append((now, num_bytes))
+        return self.average()
+
+    def average(self):
+        # suboptimal
+        return float(sum([x for t,x in self.window])) / self.window_size
+
+    def global_average(self):
+        if self.start_time == self.last_time:
+            return 0
+        return float(self.total) / (self.last_time - self.start_time)
+
+class UI(object):
+    def __init__(self):
+        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.window.show_all()
+
+    def set_average_in(self, avg):
+        self.label_avg_in.set_label(str(avg))
+
+    def set_average_out(self, avg):
+        self.label_avg_out.set_label(str(avg))
 
 def main():
     debug = True
@@ -17,6 +71,9 @@ def main():
     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)
     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)
@@ -52,6 +109,12 @@ def main():
             return fd == accepter
         (src, dst, data, completer) = result
         completer()
+        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())
         print repr((src, dst, data))
         return True
     def update_fds():