bwmon/__init__.py

   1 # -*- coding: utf-8 -*-
   2
   3 from __future__ import absolute_import
   4
   5 from bwmon import proc
   6
   7 import collections
   8 import time
   9 import sys
  10 import curses
  11 import copy
  12 import re
  13
  14 def clear():
  15     curses.setupterm()
  16     sys.stdout.write(curses.tigetstr("clear"))
  17     sys.stdout.flush()
  18
  19 class Monitor(object):
  20     def __init__(self):
  21         self.fd_map = {}
  22         self.conntrack = {}
  23         self.last_conntrack = {}
  24         self.connections = {}
  25         self.tracking = {}
  26         self.include_filter = []
  27         self.exclude_filter = []
  28
  29     def update(self):
  30         self.fd_map.update(proc.get_fd_map())
  31         self.last_conntrack = copy.deepcopy(self.conntrack)
  32         self.conntrack.update(proc.parse_ip_conntrack())
  33         self.connections.update(proc.get_connections())
  34
  35     def set_filter(self, include_filter, exclude_filter):
  36         self.include_filter = [re.compile(f) for f in include_filter] if include_filter else []
  37         self.exclude_filter = [re.compile(f) for f in exclude_filter] if exclude_filter else []
  38
  39     def convert(self):
  40         for con in self.connections.itervalues():
  41             inode = con.get('inode', None)
  42             process = self.fd_map.get(inode, None)
  43
  44             if process is None:
  45                 continue
  46
  47             if self.include_filter and not any([f.search(process['cmd']) for f in self.include_filter]):
  48                 continue
  49
  50             if self.exclude_filter and any([f.search(process['cmd']) for f in self.exclude_filter]):
  51                 continue
  52
  53             key_in  = proc.ip_hash(con['remote'], con['local'])
  54             key_out = proc.ip_hash(con['local'], con['remote'])
  55             keys = {'in': key_in, 'out': key_out}
  56             new_byte = {'in': 0, 'out': 0}
  57
  58             for direction in ('in', 'out'):
  59                 k = keys[direction]
  60                 if k in self.conntrack:
  61                     if key_in in self.last_conntrack:
  62                         new_byte[direction] = int(self.conntrack[k]['bytes']) - int(self.last_conntrack[k]['bytes'])
  63                     else:
  64                         new_byte[direction] = int(self.conntrack[k]['bytes'])
  65
  66             if process['cmd'] in self.tracking:
  67                 old_in, old_out = self.tracking[process['cmd']]
  68             else:
  69                 old_in = 0
  70                 old_out = 0
  71
  72             self.tracking[process['cmd']] = (old_in + new_byte['in'], old_out + new_byte['out'])
  73
  74     def output(self):
  75         def compare(a, b):
  76             return cmp(a[1], b[1])
  77
  78         clear()
  79         for cmd, bytes in sorted(self.tracking.iteritems(), cmp=compare):
  80             if len(cmd) > 60:
  81                 cmd = cmd[:57] + '...'
  82             print '%10d / %10d -- %s' % (bytes[0], bytes[1], cmd)
  83             sys.stdout.flush()
  84
  85     def loop(self):
  86         while True:
  87             self.update()
  88             self.convert()
  89             self.output()
  90             time.sleep(1)
  91
  92 #def loop():
  93 #    """
  94 #    prints a list of processes
  95 #    """
  96 #
  97 #    while True:
  98 #        for con in connections.itervalues():
  99 #
 100 #            if not con['inode']:
 101 #                continue
 102 #
 103 #            if not con['inode'] in fd_map:
 104 #                print 'inode %s not im map' % con['inode']
 105 #                continue
 106 #
 107 #            process = fd_map[con['inode']]
 108 #
 109 #
 110 #            key1 = proc.ip_hash(con['local'], con['remote'])
 111 #            key2 = proc.ip_hash(con['remote'], con['local'])
 112 #            key = None
 113 #
 114 #            if key1 in conntrack:
 115 #                key = key1
 116 #            elif key2 in conntrack:
 117 #                key = key2
 118 #
 119 #            if key and 'bytes' in conntrack[key]:
 120 #                byte_count = conntrack[key]['bytes']
 121 #            else:
 122 #                byte_count = '?'
 123 #
 124 #            print '%d: %s: %s -> %s (%s Bytes)' % (process['pid'], process['cmd'], con['local'], con['remote'], byte_count)
 125 #        time.sleep(1)
 126 #
 127 #
 128