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