services/util.py

   1 # vim: set encoding=utf-8 et sw=4 sts=4 :
   2
   3 import subprocess
   4 from mako.template import Template
   5 from mako.runtime import Context
   6 from os import path
   7 try:
   8     import cPickle as pickle
   9 except ImportError:
  10     import pickle
  11
  12 from dispatcher import Handler, handler, HandlerError
  13
  14 #DEBUG = False
  15 DEBUG = True
  16
  17 __ALL__ = ('ServiceHandler', 'InitdHandler', 'Persistent', 'ConfigWriter',
  18             'Error', 'ReturnNot0Error', 'ExecutionError', 'call')
  19
  20 class Error(HandlerError):
  21     r"""
  22     Error(message) -> Error instance :: Base ServiceHandler exception class.
  23
  24     All exceptions raised by the ServiceHandler inherits from this one, so
  25     you can easily catch any ServiceHandler exception.
  26
  27     message - A descriptive error message.
  28     """
  29
  30     def __init__(self, message):
  31         r"Initialize the object. See class documentation for more info."
  32         self.message = message
  33
  34     def __str__(self):
  35         return self.message
  36
  37 class ReturnNot0Error(Error):
  38     r"""
  39     ReturnNot0Error(return_value) -> ReturnNot0Error instance.
  40
  41     A command didn't returned the expected 0 return value.
  42
  43     return_value - Return value returned by the command.
  44     """
  45
  46     def __init__(self, return_value):
  47         r"Initialize the object. See class documentation for more info."
  48         self.return_value = return_value
  49
  50     def __str__(self):
  51         return 'The command returned %d' % self.return_value
  52
  53 class ExecutionError(Error):
  54     r"""
  55     ExecutionError(command, error) -> ExecutionError instance.
  56
  57     Error executing a command.
  58
  59     command - Command that was tried to execute.
  60
  61     error - Error received when trying to execute the command.
  62     """
  63
  64     def __init__(self, command, error):
  65         r"Initialize the object. See class documentation for more info."
  66         self.command = command
  67         self.error = error
  68
  69     def __str__(self):
  70         command = self.command
  71         if not isinstance(self.command, basestring):
  72             command = ' '.join(command)
  73         return "Can't execute command %s: %s" % (command, self.error)
  74
  75 class ParameterError(Error, KeyError):
  76     r"""
  77     ParameterError(paramname) -> ParameterError instance
  78
  79     This is the base exception for all DhcpHandler parameters related errors.
  80     """
  81
  82     def __init__(self, paramname):
  83         r"Initialize the object. See class documentation for more info."
  84         self.message = 'Parameter error: "%s"' % paramname
  85
  86 class ParameterNotFoundError(ParameterError):
  87     r"""
  88     ParameterNotFoundError(hostname) -> ParameterNotFoundError instance
  89
  90     This exception is raised when trying to operate on a parameter that doesn't
  91     exists.
  92     """
  93
  94     def __init__(self, paramname):
  95         r"Initialize the object. See class documentation for more info."
  96         self.message = 'Parameter not found: "%s"' % paramname
  97
  98 def call(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
  99             stderr=subprocess.PIPE, close_fds=True, universal_newlines=True,
 100             **kw):
 101     if DEBUG:
 102         if not isinstance(command, basestring):
 103             command = ' '.join(command)
 104         print 'Executing command:', command
 105         return
 106     try:
 107         r = subprocess.call(command, stdin=stdin, stdout=stdout, stderr=stderr,
 108                                 universal_newlines=universal_newlines,
 109                                 close_fds=close_fds, **kw)
 110     except Exception, e:
 111         raise ExecutionError(command, e)
 112     if r is not 0:
 113         raise ExecutionError(command, ReturnNot0Error(r))
 114
 115 class Persistent:
 116     r"""Persistent([attrs[, dir[, ext]]]) -> Persistent.
 117
 118     This is a helper class to inherit from to automatically handle data
 119     persistence using pickle.
 120
 121     The variables attributes to persist (attrs), and the pickle directory (dir)
 122     and file extension (ext) can be defined by calling the constructor or in a
 123     more declarative way as class attributes, like:
 124
 125     class TestHandler(Persistent):
 126         _persistent_attrs = ('some_attr', 'other_attr')
 127         _persistent_dir = 'persistent-data'
 128         _persistent_ext = '.pickle'
 129
 130     The default dir is '.' and the default extension is '.pkl'. There are no
 131     default variables, and they should be specified as string if a single
 132     attribute should be persistent or as a tuple of strings if they are more.
 133     The strings should be the attribute names to be persisted. For each
 134     attribute a separated pickle file is generated in the pickle directory.
 135
 136     You can call _dump() and _load() to write and read the data respectively.
 137     """
 138     # TODO implement it using metaclasses to add the handlers method by demand
 139     # (only for specifieds commands).
 140
 141     _persistent_attrs = ()
 142     _persistent_dir = '.'
 143     _persistent_ext = '.pkl'
 144
 145     def __init__(self, attrs=None, dir=None, ext=None):
 146         r"Initialize the object, see the class documentation for details."
 147         if attrs is not None:
 148             self._persistent_attrs = attrs
 149         if dir is not None:
 150             self._persistent_dir = dir
 151         if ext is not None:
 152             self._persistent_ext = ext
 153
 154     def _dump(self):
 155         r"_dump() -> None :: Dump all persistent data to pickle files."
 156         if isinstance(self._persistent_attrs, basestring):
 157             self._persistent_attrs = (self._persistent_attrs,)
 158         for attrname in self._persistent_attrs:
 159             self._dump_attr(attrname)
 160
 161     def _load(self):
 162         r"_load() -> None :: Load all persistent data from pickle files."
 163         if isinstance(self._persistent_attrs, basestring):
 164             self._persistent_attrs = (self._persistent_attrs,)
 165         for attrname in self._persistent_attrs:
 166             self._load_attr(attrname)
 167
 168     def _dump_attr(self, attrname):
 169         r"_dump_attr() -> None :: Dump a specific variable to a pickle file."
 170         f = file(self._pickle_filename(attrname), 'wb')
 171         pickle.dump(getattr(self, attrname), f, 2)
 172         f.close()
 173
 174     def _load_attr(self, attrname):
 175         r"_load_attr() -> object :: Load a specific pickle file."
 176         f = file(self._pickle_filename(attrname))
 177         setattr(self, attrname, pickle.load(f))
 178         f.close()
 179
 180     def _pickle_filename(self, name):
 181         r"_pickle_filename() -> string :: Construct a pickle filename."
 182         return path.join(self._persistent_dir, name) + self._persistent_ext
 183
 184 class Restorable(Persistent):
 185     r"""Restorable([defaults]) -> Restorable.
 186
 187     This is a helper class to inherit from that provides a nice _restore()
 188     method to restore the persistent data if any, or load some nice defaults
 189     if not.
 190
 191     The defaults can be defined by calling the constructor or in a more
 192     declarative way as class attributes, like:
 193
 194     class TestHandler(Restorable):
 195         _persistent_attrs = ('some_attr', 'other_attr')
 196         _restorable_defaults = dict(
 197                 some_attr = 'some_default',
 198                 other_attr = 'other_default')
 199
 200     The defaults is a dictionary, very coupled with the _persistent_attrs
 201     attribute inherited from Persistent. The defaults keys should be the
 202     values from _persistent_attrs, and the values the default values.
 203
 204     The _restore() method returns True if the data was restored successfully
 205     or False if the defaults were loaded (in case you want to take further
 206     actions). If a _write_config method if found, it's executed when a restore
 207     fails too.
 208     """
 209     # TODO implement it using metaclasses to add the handlers method by demand
 210     # (only for specifieds commands).
 211
 212     _restorable_defaults = dict()
 213
 214     def __init__(self, defaults=None):
 215         r"Initialize the object, see the class documentation for details."
 216         if defaults is not None:
 217             self._restorable_defaults = defaults
 218
 219     def _restore(self):
 220         r"_restore() -> bool :: Restore persistent data or create a default."
 221         try:
 222             self._load()
 223             return True
 224         except IOError:
 225             for (k, v) in self._restorable_defaults.items():
 226                 setattr(self, k, v)
 227             self._dump()
 228             if hasattr(self, '_write_config'):
 229                 self._write_config()
 230             return False
 231
 232 class ConfigWriter:
 233     r"""ConfigWriter([initd_name[, initd_dir]]) -> ConfigWriter.
 234
 235     This is a helper class to inherit from to automatically handle
 236     configuration generation. Mako template system is used for configuration
 237     files generation.
 238
 239     The configuration filenames, the generated configuration files directory
 240     and the templates directory can be defined by calling the constructor or
 241     in a more declarative way as class attributes, like:
 242
 243     class TestHandler(ConfigWriter):
 244         _config_writer_files = ('base.conf', 'custom.conf')
 245         _config_writer_cfg_dir = '/etc/service'
 246         _config_writer_tpl_dir = 'templates'
 247
 248     The generated configuration files directory defaults to '.' and the
 249     templates directory to 'templates'. _config_writer_files has no default and
 250     must be specified in either way. It can be string or a tuple if more than
 251     one configuration file must be generated.
 252
 253     The template filename and the generated configuration filename are both the
 254     same (so if you want to generate some /etc/config, you should have some
 255     templates/config template). That's why _config_writer_cfg_dir and
 256     _config_writer_tpl_dir can't be the same.
 257
 258     When you write your Handler, you should call _config_build_templates() in
 259     you Handler constructor to build the templates.
 260
 261     To write the configuration files, you must use the _write_config() method.
 262     To know what variables to replace in the template, you have to provide a
 263     method called _get_config_vars(tamplate_name), which should return a
 264     dictionary of variables to pass to the template system to be replaced in
 265     the template for the configuration file 'config_file'.
 266     """
 267     # TODO implement it using metaclasses to add the handlers method by demand
 268     # (only for specifieds commands).
 269
 270     _config_writer_files = ()
 271     _config_writer_cfg_dir = '.'
 272     _config_writer_tpl_dir = 'templates'
 273
 274     def __init__(self, files=None, cfg_dir=None, tpl_dir=None):
 275         r"Initialize the object, see the class documentation for details."
 276         if files is not None:
 277             self._config_writer_files = files
 278         if cfg_dir is not None:
 279             self._config_writer_cfg_dir = cfg_dir
 280         if tpl_dir is not None:
 281             self._config_writer_tpl_dir = tpl_dir
 282         self._config_build_templates()
 283
 284     def _config_build_templates(self):
 285         r"_config_writer_templates() -> None :: Build the template objects."
 286         if isinstance(self._config_writer_files, basestring):
 287             self._config_writer_files = (self._config_writer_files,)
 288         if not hasattr(self, '_config_writer_templates') \
 289                                         or not self._config_writer_templates:
 290             self._config_writer_templates = dict()
 291             for t in self._config_writer_files:
 292                 f = path.join(self._config_writer_tpl_dir, t)
 293                 self._config_writer_templates[t] = Template(filename=f)
 294
 295     def _render_config(self, template_name, vars=None):
 296         r"""_render_config(template_name[, config_filename[, vars]]).
 297
 298         Render a single config file using the template 'template_name'. If
 299         vars is specified, it's used as the dictionary with the variables
 300         to replace in the templates, if not, it looks for a
 301         _get_config_vars() method to get it.
 302         """
 303         if vars is None:
 304             if hasattr(self, '_get_config_vars'):
 305                 vars = self._get_config_vars(template_name)
 306             else:
 307                 vars = dict()
 308         elif callable(vars):
 309             vars = vars(template_name)
 310         return self._config_writer_templates[template_name].render(**vars)
 311
 312     def _write_single_config(self, template_name, config_filename=None, vars=None):
 313         r"""_write_single_config(template_name[, config_filename[, vars]]).
 314
 315         Write a single config file using the template 'template_name'. If no
 316         config_filename is specified, the config filename will be the same as
 317         the 'template_name' (but stored in the generated config files
 318         directory). If it's specified, the generated config file is stored in
 319         the file called 'config_filename' (also in the generated files
 320         directory). If vars is specified, it's used as the dictionary with the
 321         variables to replace in the templates, if not, it looks for a
 322         _get_config_vars() method to get it.
 323         """
 324         if not config_filename:
 325             config_filename = template_name
 326         if vars is None:
 327             if hasattr(self, '_get_config_vars'):
 328                 vars = self._get_config_vars(template_name)
 329             else:
 330                 vars = dict()
 331         elif callable(vars):
 332             vars = vars(template_name)
 333         f = file(path.join(self._config_writer_cfg_dir, config_filename), 'w')
 334         ctx = Context(f, **vars)
 335         self._config_writer_templates[template_name].render_context(ctx)
 336         f.close()
 337
 338     def _write_config(self):
 339         r"_write_config() -> None :: Generate all the configuration files."
 340         for t in self._config_writer_files:
 341             self._write_single_config(t)
 342
 343 class ServiceHandler(Handler):
 344     r"""ServiceHandler([start[, stop[, restart[, reload]]]]) -> ServiceHandler.
 345
 346     This is a helper class to inherit from to automatically handle services
 347     with start, stop, restart, reload actions.
 348
 349     The actions can be defined by calling the constructor with all the
 350     parameters or in a more declarative way as class attributes, like:
 351
 352     class TestHandler(ServiceHandler):
 353         _service_start = ('command', 'start')
 354         _service_stop = ('command', 'stop')
 355         _service_restart = ('command', 'restart')
 356         _service_reload = 'reload-command'
 357
 358     Commands are executed without using the shell, that's why they are specified
 359     as tuples (where the first element is the command and the others are the
 360     command arguments). If only a command is needed (without arguments) a single
 361     string can be specified.
 362
 363     All commands must be specified.
 364     """
 365     # TODO implement it using metaclasses to add the handlers method by demand
 366     # (only for specifieds commands).
 367
 368     def __init__(self, start=None, stop=None, restart=None, reload=None):
 369         r"Initialize the object, see the class documentation for details."
 370         for (name, action) in dict(start=start, stop=stop, restart=restart,
 371                                                     reload=reload).items():
 372             if action is not None:
 373                 setattr(self, '_service_%s' % name, action)
 374
 375     @handler(u'Start the service.')
 376     def start(self):
 377         r"start() -> None :: Start the service."
 378         call(self._service_start)
 379
 380     @handler(u'Stop the service.')
 381     def stop(self):
 382         r"stop() -> None :: Stop the service."
 383         call(self._service_stop)
 384
 385     @handler(u'Restart the service.')
 386     def restart(self):
 387         r"restart() -> None :: Restart the service."
 388         call(self._service_restart)
 389
 390     @handler(u'Reload the service config (without restarting, if possible).')
 391     def reload(self):
 392         r"reload() -> None :: Reload the configuration of the service."
 393         call(self._service_reload)
 394
 395 class InitdHandler(Handler):
 396     r"""InitdHandler([initd_name[, initd_dir]]) -> InitdHandler.
 397
 398     This is a helper class to inherit from to automatically handle services
 399     with start, stop, restart, reload actions using a /etc/init.d like script.
 400
 401     The name and directory of the script can be defined by calling the
 402     constructor or in a more declarative way as class attributes, like:
 403
 404     class TestHandler(ServiceHandler):
 405         _initd_name = 'some-service'
 406         _initd_dir = '/usr/local/etc/init.d'
 407
 408     The default _initd_dir is '/etc/init.d', _initd_name has no default and
 409     must be specified in either way.
 410
 411     Commands are executed without using the shell.
 412     """
 413     # TODO implement it using metaclasses to add the handlers method by demand
 414     # (only for specifieds commands).
 415
 416     _initd_dir = '/etc/init.d'
 417
 418     def __init__(self, initd_name=None, initd_dir=None):
 419         r"Initialize the object, see the class documentation for details."
 420         if initd_name is not None:
 421             self._initd_name = initd_name
 422         if initd_dir is not None:
 423             self._initd_dir = initd_dir
 424
 425     @handler(u'Start the service.')
 426     def start(self):
 427         r"start() -> None :: Start the service."
 428         call((path.join(self._initd_dir, self._initd_name), 'start'))
 429
 430     @handler(u'Stop the service.')
 431     def stop(self):
 432         r"stop() -> None :: Stop the service."
 433         call((path.join(self._initd_dir, self._initd_name), 'stop'))
 434
 435     @handler(u'Restart the service.')
 436     def restart(self):
 437         r"restart() -> None :: Restart the service."
 438         call((path.join(self._initd_dir, self._initd_name), 'restart'))
 439
 440     @handler(u'Reload the service config (without restarting, if possible).')
 441     def reload(self):
 442         r"reload() -> None :: Reload the configuration of the service."
 443         call((path.join(self._initd_dir, self._initd_name), 'reload'))
 444
 445 class TransactionalHandler(Handler):
 446     r"""Handle command transactions providing a commit and rollback commands.
 447
 448     This is a helper class to inherit from to automatically handle
 449     transactional handlers, which have commit and rollback commands.
 450
 451     The handler should provide a reload() method (see ServiceHandler and
 452     InitdHandler for helper classes to provide this) which will be called
 453     when a commit command is issued (if a reload() command is present).
 454     The persistent data will be written too (if a _dump() method is provided,
 455     see Persistent and Restorable for that), and the configuration files
 456     will be generated (if a _write_config method is present, see ConfigWriter).
 457     """
 458     # TODO implement it using metaclasses to add the handlers method by demand
 459     # (only for specifieds commands).
 460
 461     @handler(u'Commit the changes (reloading the service, if necessary).')
 462     def commit(self):
 463         r"commit() -> None :: Commit the changes and reload the service."
 464         if hasattr(self, '_dump'):
 465             self._dump()
 466         if hasattr(self, '_write_config'):
 467             self._write_config()
 468         if hasattr(self, '_reload'):
 469             self.reload()
 470
 471     @handler(u'Discard all the uncommited changes.')
 472     def rollback(self):
 473         r"rollback() -> None :: Discard the changes not yet commited."
 474         if hasattr(self, '_load'):
 475             self._load()
 476
 477 class ParametersHandler(Handler):
 478     r"""ParametersHandler([attr]) -> ParametersHandler.
 479
 480     This is a helper class to inherit from to automatically handle
 481     service parameters, providing set, get, list and show commands.
 482
 483     The attribute that holds the parameters can be defined by calling the
 484     constructor or in a more declarative way as class attributes, like:
 485
 486     class TestHandler(ServiceHandler):
 487         _parameters_attr = 'some_attr'
 488
 489     The default is 'params' and it should be a dictionary.
 490     """
 491     # TODO implement it using metaclasses to add the handlers method by demand
 492     # (only for specifieds commands).
 493
 494     _parameters_attr = 'params'
 495
 496     def __init__(self, attr=None):
 497         r"Initialize the object, see the class documentation for details."
 498         if attr is not None:
 499             self._parameters_attr = attr
 500
 501     @handler(u'Set a service parameter.')
 502     def set(self, param, value):
 503         r"set(param, value) -> None :: Set a service parameter."
 504         if not param in self.params:
 505             raise ParameterNotFoundError(param)
 506         self.params[param] = value
 507
 508     @handler(u'Get a service parameter.')
 509     def get(self, param):
 510         r"get(param) -> None :: Get a service parameter."
 511         if not param in self.params:
 512             raise ParameterNotFoundError(param)
 513         return self.params[param]
 514
 515     @handler(u'List all available service parameters.')
 516     def list(self):
 517         r"list() -> tuple :: List all the parameter names."
 518         return self.params.keys()
 519
 520     @handler(u'Get all service parameters, with their values.')
 521     def show(self):
 522         r"show() -> (key, value) tuples :: List all the parameters."
 523         return self.params.items()
 524
 525
 526 if __name__ == '__main__':
 527
 528     # Execution tests
 529     class STestHandler1(ServiceHandler):
 530         _service_start = ('service', 'start')
 531         _service_stop = ('service', 'stop')
 532         _service_restart = ('ls', '/')
 533         _service_reload = ('cp', '/la')
 534     class STestHandler2(ServiceHandler):
 535         def __init__(self):
 536             ServiceHandler.__init__(self, 'cmd-start', 'cmd-stop',
 537                                         'cmd-restart', 'cmd-reload')
 538     class ITestHandler1(InitdHandler):
 539         _initd_name = 'test1'
 540     class ITestHandler2(InitdHandler):
 541         def __init__(self):
 542             InitdHandler.__init__(self, 'test2', '/usr/local/etc/init.d')
 543     handlers = [
 544         STestHandler1(),
 545         STestHandler2(),
 546         ITestHandler1(),
 547         ITestHandler2(),
 548     ]
 549     for h in handlers:
 550         print h.__class__.__name__
 551         try:
 552             h.start()
 553         except ExecutionError, e:
 554             print e
 555         try:
 556             h.stop()
 557         except ExecutionError, e:
 558             print e
 559         try:
 560             h.restart()
 561         except ExecutionError, e:
 562             print e
 563         try:
 564             h.reload()
 565         except ExecutionError, e:
 566             print e
 567         print
 568
 569     # Persistent test
 570     print 'PTestHandler'
 571     class PTestHandler(Persistent):
 572         _persistent_attrs = 'vars'
 573         def __init__(self):
 574             self.vars = dict(a=1, b=2)
 575     h = PTestHandler()
 576     print h.vars
 577     h._dump()
 578     h.vars['x'] = 100
 579     print h.vars
 580     h._load()
 581     print h.vars
 582     h.vars['x'] = 100
 583     h._dump()
 584     print h.vars
 585     del h.vars['x']
 586     print h.vars
 587     h._load()
 588     print h.vars
 589     print
 590
 591     # Restorable test
 592     print 'RTestHandler'
 593     class RTestHandler(Restorable):
 594         _persistent_attrs = 'vars'
 595         _restorable_defaults = dict(vars=dict(a=1, b=2))
 596         def __init__(self):
 597             self._restore()
 598     h = RTestHandler()
 599     print h.vars
 600     h.vars['x'] = 100
 601     h._dump()
 602     h = RTestHandler()
 603     print h.vars
 604     print
 605
 606     # ConfigWriter test
 607     print 'CTestHandler'
 608     import os
 609     os.mkdir('templates')
 610     f = file('templates/config', 'w')
 611     f.write('Hello, ${name}! You are ${what}.')
 612     f.close()
 613     print 'template:'
 614     print file('templates/config').read()
 615     class CTestHandler(ConfigWriter):
 616         _config_writer_files = 'config'
 617         def __init__(self):
 618             self._config_build_templates()
 619         def _get_config_vars(self, config_file):
 620             return dict(name='you', what='a parrot')
 621     h = CTestHandler()
 622     h._write_config()
 623     print 'config:'
 624     print file('config').read()
 625     os.unlink('config')
 626     os.unlink('templates/config')
 627     os.rmdir('templates')
 628     print
 629