services/qos/handler.py

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