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 from bwmon import model
   8
   9 import collections
  10 import time
  11 import sys
  12 import copy
  13 import re
  14
  15 BANDWIDTH, TRAFFIC = range(2)
  16
  17 class Monitor(object):
  18     def __init__(self):
  19         self.fd_map = {}
  20         self.sample_time = time.time()
  21         self.conntrack = {}
  22         self.last_conntrack = {}
  23         self.connections = {}
  24         self.entries = model.MonitorEntryCollection()
  25         self.include_filter = []
  26         self.exclude_filter = []
  27
  28     def update(self):
  29         self.fd_map.update(proc.get_fd_map())
  30         self.last_conntrack = copy.deepcopy(self.conntrack)
  31         self.conntrack.update(proc.parse_ip_conntrack())
  32         self.connections.update(proc.get_connections())
  33         self.entries.expire()
  34         self.sample_time = time.time()
  35
  36     def set_filter(self, include_filter, exclude_filter):
  37         self.include_filter = [re.compile(f) for f in include_filter] if include_filter else []
  38         self.exclude_filter = [re.compile(f) for f in exclude_filter] if exclude_filter else []
  39
  40     def convert(self):
  41         entries = collections.defaultdict(lambda: (0, 0))
  42         for con in self.connections.itervalues():
  43             inode = con.get('inode', None)
  44             process = self.fd_map.get(inode, None)
  45
  46             if process is None:
  47                 continue
  48
  49             if self.include_filter and not any([f.search(process['cmd']) for f in self.include_filter]):
  50                 continue
  51
  52             if self.exclude_filter and any([f.search(process['cmd']) for f in self.exclude_filter]):
  53                 continue
  54
  55             key_in  = proc.ip_hash(con['remote'], con['local'])
  56             key_out = proc.ip_hash(con['local'], con['remote'])
  57             keys = {'in': key_in, 'out': key_out}
  58             new_byte = {'in': 0, 'out': 0}
  59
  60             for direction in ('in', 'out'):
  61                 k = keys[direction]
  62                 if k in self.conntrack:
  63                     if key_in in self.last_conntrack:
  64                         new_byte[direction] = int(self.conntrack[k]['bytes']) - int(self.last_conntrack[k]['bytes'])
  65                     else:
  66                         new_byte[direction] = int(self.conntrack[k]['bytes'])
  67
  68             current_in, current_out = entries[process['cmd']]
  69             new_in, new_out = (new_byte['in'], new_byte['out'])
  70
  71             entries[process['cmd']] = (current_in + new_in, current_out + new_out)
  72
  73         for key in entries:
  74             new_in, new_out = entries[key]
  75             old_in, old_out = self.entries.get_last_bytes(key)
  76             entry = model.MonitorEntry(key, old_in + new_in, old_out + new_out, self.sample_time)
  77             self.entries.add(entry)
  78
  79     def output(self, mode=TRAFFIC):
  80         util.clear()
  81         if mode == TRAFFIC:
  82             entries = sorted(self.entries.get_traffic())
  83         else:
  84             entries = sorted(self.entries.get_usage())
  85
  86         for bytes_in, bytes_out, cmd in entries:
  87             if len(cmd) > 60:
  88                 cmd = cmd[:57] + '...'
  89             print '%10d / %10d -- %s' % (bytes_in, bytes_out, cmd)
  90             sys.stdout.flush()
  91
  92     def loop(self, mode):
  93         while True:
  94             self.update()
  95             self.convert()
  96             self.output(mode)
  97             time.sleep(1)
  98