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