bwmon/monitor.py

   1 # -*- coding: utf-8 -*-
   2
   3 from __future__ import absolute_import
   4
   5 from bwmon import proc
   6 from bwmon import util
   7
   8 import collections
   9 import time
  10 import sys
  11 import copy
  12 import re
  13
  14 class Monitor(object):
  15     def __init__(self):
  16         self.fd_map = {}
  17         self.conntrack = {}
  18         self.last_conntrack = {}
  19         self.connections = {}
  20         self.tracking = {}
  21         self.include_filter = []
  22         self.exclude_filter = []
  23
  24     def update(self):
  25         self.fd_map.update(proc.get_fd_map())
  26         self.last_conntrack = copy.deepcopy(self.conntrack)
  27         self.conntrack.update(proc.parse_ip_conntrack())
  28         self.connections.update(proc.get_connections())
  29
  30     def set_filter(self, include_filter, exclude_filter):
  31         self.include_filter = [re.compile(f) for f in include_filter] if include_filter else []
  32         self.exclude_filter = [re.compile(f) for f in exclude_filter] if exclude_filter else []
  33
  34     def convert(self):
  35         for con in self.connections.itervalues():
  36             inode = con.get('inode', None)
  37             process = self.fd_map.get(inode, None)
  38
  39             if process is None:
  40                 continue
  41
  42             if self.include_filter and not any([f.search(process['cmd']) for f in self.include_filter]):
  43                 continue
  44
  45             if self.exclude_filter and any([f.search(process['cmd']) for f in self.exclude_filter]):
  46                 continue
  47
  48             key_in  = proc.ip_hash(con['remote'], con['local'])
  49             key_out = proc.ip_hash(con['local'], con['remote'])
  50             keys = {'in': key_in, 'out': key_out}
  51             new_byte = {'in': 0, 'out': 0}
  52
  53             for direction in ('in', 'out'):
  54                 k = keys[direction]
  55                 if k in self.conntrack:
  56                     if key_in in self.last_conntrack:
  57                         new_byte[direction] = int(self.conntrack[k]['bytes']) - int(self.last_conntrack[k]['bytes'])
  58                     else:
  59                         new_byte[direction] = int(self.conntrack[k]['bytes'])
  60
  61             if process['cmd'] in self.tracking:
  62                 old_in, old_out = self.tracking[process['cmd']]
  63             else:
  64                 old_in = 0
  65                 old_out = 0
  66
  67             self.tracking[process['cmd']] = (old_in + new_byte['in'], old_out + new_byte['out'])
  68
  69     def output(self):
  70         def compare(a, b):
  71             return cmp(a[1], b[1])
  72
  73         util.clear()
  74         for cmd, bytes in sorted(self.tracking.iteritems(), cmp=compare):
  75             if len(cmd) > 60:
  76                 cmd = cmd[:57] + '...'
  77             print '%10d / %10d -- %s' % (bytes[0], bytes[1], cmd)
  78             sys.stdout.flush()
  79
  80     def loop(self):
  81         while True:
  82             self.update()
  83             self.convert()
  84             self.output()
  85             time.sleep(1)
  86