pymin/services/qos/__init__.py

   1 # vim: set encoding=utf-8 et sw=4 sts=4 :
   2
   3 from subprocess import Popen, PIPE
   4 from os import path
   5
   6 from pymin.seqtools import Sequence
   7 from pymin.dispatcher import handler, HandlerError, Handler
   8 from pymin.services.util import Restorable, ConfigWriter, InitdHandler, \
   9                                 TransactionalHandler, SubHandler, call, \
  10                                 get_network_devices, ListComposedSubHandler, \
  11                                 DictComposedSubHandler
  12
  13 __ALL__ = ('QoSHandler',)
  14
  15 class ClassError(HandlerError):
  16
  17     def __init__(self, qosclass):
  18         self.message = u'Class error : "%s"' % qosclass
  19
  20
  21 class ClassNotFoundError(ClassError):
  22
  23     def __init__(self, qosclass):
  24         self.message = u'Class not found : "%s"' % qosclass
  25
  26
  27 class ClassAlreadyExistsError(ClassError):
  28
  29     def __init__(self, qosclass):
  30         self.message = u'Class already exists : "%s"' % qosclass
  31
  32
  33 class Class(Sequence):
  34
  35     def __init__(self, id, rate=None):
  36         self.id = id
  37         self.rate = rate
  38         self.hosts = list()
  39
  40     def as_tuple(self):
  41         return (self.id, self.rate)
  42
  43     def __cmp__(self, other):
  44         if self.id == other.id
  45             return 0
  46         return cmp(id(self), id(other))
  47
  48
  49 class ClassHandler(Handler):
  50
  51     def __init__(self, parent):
  52         self.parent = parent
  53
  54     @handler('Adds a class : add <id> <device> <rate>')
  55     def add(self, id, dev, rate):
  56         if not dev in self.parent.devices:
  57             raise DeviceNotFoundError(device)
  58         c = Class(id, dev, rate)
  59         try:
  60             self.parent.classes.index(c)
  61             raise ClassAlreadyExistsError(id  + '->' + dev)
  62         except ValueError:
  63             self.parent.classes.append(c)
  64
  65     @handler(u'Deletes a class : delete <id> <device>')
  66     def delete(self, id, dev):
  67         if not dev in self.parent.devices:
  68             raise DeviceNotFoundError(device)
  69         c = Class(id, dev)
  70         try:
  71             self.parent.classes.remove(c)
  72         except ValueError:
  73             raise ClassNotFoundError(id + '->' + dev)
  74
  75     @handler(u'Lists classes : list <dev>')
  76     def list(self, device):
  77         try:
  78             k = self.parent.classes.keys()
  79         except ValueError:
  80             k = list()
  81         return k
  82
  83     @handler(u'Get information about all classes: show <dev>')
  84     def show(self, device):
  85         try:
  86             k = self.parent.classes.values()
  87         except ValueError:
  88             k = list()
  89         return k
  90
  91
  92 class Host(Sequence):
  93
  94     def __init__(self, ip):
  95         self.ip = ip
  96
  97     def as_tuple(self):
  98         return (self.ip)
  99
 100
 101 class HostHandler(DictComposedSubHandler):
 102     handler_help = u"Manage Hosts"
 103     _comp_subhandler_cont = 'classes'
 104     _comp_subhandler_attr = 'hosts'
 105     _comp_subhandler_class = Host
 106
 107
 108 class Device(Sequence):
 109
 110     def __init__(self, name, mac):
 111         self.name = name
 112         self.mac = mac
 113         self.classes = list()
 114
 115     def as_tuple(self):
 116         return (self.name, self.mac)
 117
 118
 119 class DeviceHandler(SubHandler):
 120
 121     handler_help = u"Manage network devices"
 122
 123     def __init__(self, parent):
 124         # FIXME remove templates to execute commands
 125         from mako.template import Template
 126         self.parent = parent
 127         template_dir = path.join(path.dirname(__file__), 'templates')
 128         dev_fn = path.join(template_dir, 'device')
 129         self.device_template = Template(filename=dev_fn)
 130
 131     @handler(u'Bring the device up')
 132     def up(self, name):
 133         if name in self.parent.devices:
 134             call(self.device_template.render(dev=name, action='add'), shell=True)
 135         else:
 136             raise DeviceNotFoundError(name)
 137
 138     @handler(u'Bring the device down')
 139     def down(self, name):
 140         if name in self.parent.devices:
 141             call(self.device_template.render(dev=name, action='del'), shell=True)
 142         else:
 143             raise DeviceNotFoundError(name)
 144
 145     @handler(u'List all devices')
 146     def list(self):
 147         return self.parent.devices.keys()
 148
 149     @handler(u'Get information about a device')
 150     def show(self):
 151         return self.parent.devices.items()
 152
 153
 154 class QoSHandler(Restorable, ConfigWriter, TransactionalHandler):
 155
 156     handler_help = u"Manage QoS devices, classes and hosts"
 157
 158     _persistent_attrs = ('devices','classes','hosts')
 159
 160     _restorable_defaults = dict(
 161                             devices=dict((dev, Device(dev, mac))
 162                                 for (dev, mac) in get_network_devices().items()),
 163                             classes = list()
 164                             hosts = list()
 165                             )
 166
 167     _config_writer_files = ('device', 'class_add', 'class_del', 'host_add')
 168
 169     _config_writer_tpl_dir = path.join(path.dirname(__file__), 'templates')
 170
 171     def __init__(self, pickle_dir='.', config_dir='.'):
 172         r"Initialize QoSHandler object, see class documentation for details."
 173         self._persistent_dir = pickle_dir
 174         self._config_writer_cfg_dir = config_dir
 175         self._config_build_templates()
 176         self._restore()
 177         self._write_config()
 178         self.dev = DeviceHandler(self)
 179         self.classes = ClassHandler(self)
 180         self.hosts = HostHandler(self)
 181
 182     def _write_config(self):
 183         r"_write_config() -> None :: Execute all commands."
 184         for device in self.devices.values():
 185             call(self._render_config('devices', dict(dev=device.name, action='del')), shell=True)
 186             call(self._render_config('devices', dict(dev=device.name, action='add')), shell=True)
 187             for qosclass in device.classes:
 188                 call(self._render_config('class_add', dict(
 189                         dev = device.name,
 190                         id = qosclass.id,
 191                         rate = qosclass.rate
 192                     )
 193                 ), shell=True)
 194             for host in classes.hosts:
 195                 call(self._render_config('host_add', dict(
 196                         dev = device.name,
 197                         ip = host.ip,
 198                         id = qosclass.id
 199                     )
 200                 ), shell=True)
 201
 202     def handle_timer(self):
 203         self.refresh_devices()
 204
 205     def refresh_devices(self):
 206         devices = get_network_devices()
 207         #add not registered devices
 208         for k, v in devices.items():
 209             if k not in self.devices:
 210                 self.devices[k] = Device(k, v)
 211         #delete dead devices
 212         for k in self.devices.keys():
 213             if k not in devices:
 214                 del self.devices[k]
 215
 216
 217 if __name__ == '__main__':
 218
 219     qos = QoSHandler()
 220     print '----------------------'
 221     qos.commit()