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