services/qos/handler.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.service.util import Restorable, ConfigWriter, InitdHandler, \
   9                                TransactionalHandler, SubHandler, call, \
  10                                get_network_devices, ListComposedSubHandler, \
  11                                DictComposedSubHandler, ExecutionError
  12
  13 __all__ = ('QoSHandler')
  14
  15
  16 class DeviceError(HandlerError):
  17
  18     def __init__(self, dev):
  19         self.message = u'Devive error : "%s"' % dev
  20
  21
  22 class DeviceNotFoundError(DeviceError):
  23
  24     def __init__(self, dev):
  25         self.message = u'Device not found : "%s"' % dev
  26
  27
  28 class ClassError(HandlerError):
  29
  30     def __init__(self, cls):
  31         self.message = u'Class error : "%s"' % cls
  32
  33
  34 class ClassNotFoundError(ClassError):
  35
  36     def __init__(self, cls):
  37         self.message = u'Class not found : "%s"' % cls
  38
  39
  40 class ClassAlreadyExistsError(ClassError):
  41
  42     def __init__(self, cls):
  43         self.message = u'Class already exists : "%s"' % cls
  44
  45
  46 class HostError(HandlerError):
  47
  48     def __init__(self, host):
  49         self.message = u'Host error : "%s"' % host
  50
  51
  52 class HostNotFoundError(HostError):
  53
  54     def __init__(self, ip):
  55         self.message = u'Host not found : "%s"' % host
  56
  57
  58 class HostAlreadyExistsError(HostError):
  59
  60     def __init__(self, ip):
  61         self.message = u'Host already exists : "%s"' % host
  62
  63
  64 class Class(Sequence):
  65
  66     def __init__(self, cid, rate=None):
  67         self.cid = cid
  68         self.rate = rate
  69         self.hosts = dict()
  70
  71     def as_tuple(self):
  72         return (self.cid, self.rate)
  73
  74     def __cmp__(self, other):
  75         if self.cid == other.cid:
  76             return 0
  77         return cmp(id(self), id(other))
  78
  79
  80 class ClassHandler(Handler):
  81
  82     def __init__(self, parent):
  83         self.parent = parent
  84
  85     @handler('Adds a class : add <id> <device> <rate>')
  86     def add(self, dev, cid, rate):
  87         if not dev in self.parent.devices:
  88             raise DeviceNotFoundError(dev)
  89
  90         try:
  91             self.parent.devices[dev].classes[cid] = Class(cid, rate)
  92         except ValueError:
  93             raise ClassAlreadyExistsError(cid  + ' -> ' + dev)
  94
  95     @handler(u'Deletes a class : delete <id> <device>')
  96     def delete(self, dev, cid):
  97         if not dev in self.parent.devices:
  98             raise DeviceNotFoundError(dev)
  99
 100         try:
 101             del self.parent.devices[dev].classes[cid]
 102         except KeyError:
 103             raise ClassNotFoundError(cid + ' -> ' + dev)
 104
 105     @handler(u'Lists classes : list <dev>')
 106     def list(self, dev):
 107         try:
 108             k = self.parent.devices[dev].classes.items()
 109         except KeyError:
 110             k = dict()
 111         return k
 112
 113
 114 class Host(Sequence):
 115
 116     def __init__(self, ip):
 117         self.ip = ip
 118
 119     def as_tuple(self):
 120         return (self.ip)
 121
 122     def __cmp__(self, other):
 123         if self.ip == other.ip:
 124             return 0
 125         return cmp(id(self), id(other))
 126
 127
 128 class HostHandler(SubHandler):
 129
 130     def __init__(self, parent):
 131         self.parent = parent
 132
 133     @handler('Adds a host to a class : add <device> <class id> <ip>')
 134     def add(self, dev, cid, ip):
 135         if not dev in self.parent.devices:
 136             raise DeviceNotFoundError(dev)
 137
 138         if not cid in self.parent.devices[dev].classes:
 139             raise ClassNotFoundError(cid)
 140
 141         try:
 142             self.parent.devices[dev].classes[cid].hosts[ip] = Host(ip)
 143         except ValueError:
 144             raise HostAlreadyExistsError(h  + ' -> ' + dev)
 145
 146     @handler(u'Lists hosts : list <dev> <class id>')
 147     def list(self, dev, cid):
 148         try:
 149             k = self.parent.devices[dev].classes[cid].hosts.keys()
 150         except KeyError:
 151             k = dict()
 152         return k
 153
 154
 155 class Device(Sequence):
 156
 157     def __init__(self, name, mac):
 158         self.name = name
 159         self.mac = mac
 160         self.classes = dict()
 161
 162     def as_tuple(self):
 163         return (self.name, self.mac)
 164
 165
 166 class DeviceHandler(SubHandler):
 167
 168     handler_help = u"Manage network devices"
 169
 170     def __init__(self, parent):
 171         # FIXME remove templates to execute commands
 172         from mako.template import Template
 173         self.parent = parent
 174         template_dir = path.join(path.dirname(__file__), 'templates')
 175         dev_fn = path.join(template_dir, 'device')
 176         self.device_template = Template(filename=dev_fn)
 177
 178     @handler(u'Bring the device up')
 179     def up(self, name):
 180         if name in self.parent.devices:
 181             try:
 182                 call(self.device_template.render(dev=name, action='add'), shell=True)
 183             except ExecutionError:
 184                 pass
 185         else:
 186             raise DeviceNotFoundError(name)
 187
 188     @handler(u'Bring the device down')
 189     def down(self, name):
 190         if name in self.parent.devices:
 191             try:
 192                 call(self.device_template.render(dev=name, action='del'), shell=True)
 193             except ExecutionError:
 194                 pass
 195         else:
 196             raise DeviceNotFoundError(name)
 197
 198     @handler(u'List all devices')
 199     def list(self):
 200         return self.parent.devices.keys()
 201
 202     @handler(u'Get information about a device')
 203     def show(self):
 204         return self.parent.devices.items()
 205
 206
 207 class QoSHandler(Restorable, ConfigWriter, TransactionalHandler):
 208
 209     handler_help = u"Manage QoS devices, classes and hosts"
 210
 211     _persistent_attrs = ('devices')
 212
 213     _restorable_defaults = dict(
 214                             devices=dict((dev, Device(dev, mac))
 215                                 for (dev, mac) in get_network_devices().items())
 216                             )
 217
 218     _config_writer_files = ('device', 'class_add', 'class_del', 'host_add')
 219
 220     _config_writer_tpl_dir = path.join(path.dirname(__file__), 'templates')
 221
 222     def __init__(self, pickle_dir='.', config_dir='.'):
 223         r"Initialize QoSHandler object, see class documentation for details."
 224         self._persistent_dir = pickle_dir
 225         self._config_writer_cfg_dir = config_dir
 226         self._config_build_templates()
 227         self._restore()
 228         self._write_config()
 229         self.dev = DeviceHandler(self)
 230         self.cls = ClassHandler(self)
 231         self.host = HostHandler(self)
 232
 233     def _write_config(self):
 234         r"_write_config() -> None :: Execute all commands."
 235         for device in self.devices.values():
 236             try:
 237                 call(self._render_config('device', dict(dev=device.name, action='del')), shell=True)
 238             except ExecutionError:
 239                 pass
 240
 241             try:
 242                 call(self._render_config('device', dict(dev=device.name, action='add')), shell=True)
 243             except ExecutionError:
 244                 pass
 245
 246             for cls in device.classes.values():
 247                 try:
 248                     call(self._render_config('class_add', dict(
 249                         dev = device.name,
 250                         cid = cls.cid,
 251                         rate = cls.rate
 252                         )
 253                     ), shell=True)
 254                 except ExecutionError:
 255                     pass
 256
 257                 for host in cls.hosts.values():
 258                     try:
 259                         call(self._render_config('host_add', dict(
 260                             dev = device.name,
 261                             ip = host.ip,
 262                             cid = cls.cid
 263                             )
 264                         ), shell=True)
 265                     except ExecutionError:
 266                         pass
 267
 268     def handle_timer(self):
 269         self.refresh_devices()
 270
 271     def refresh_devices(self):
 272         devices = get_network_devices()
 273         #add not registered devices
 274         for k, v in devices.items():
 275             if k not in self.devices:
 276                 self.devices[k] = Device(k, v)
 277         #delete dead devices
 278         for k in self.devices.keys():
 279             if k not in devices:
 280                 del self.devices[k]
 281
 282
 283 if __name__ == '__main__':
 284
 285     qos = QoSHandler()
 286     print '----------------------'
 287     qos.commit()