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