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             # TODO tener en cuenta servicios que hay que levantar y los que no
 225             if hasattr(self, 'commit'): # TODO deberia ser reload y/o algo para comandos
 226                 self.commit()
 227             return True
 228         except IOError:
 229             for (k, v) in self._restorable_defaults.items():
 230                 setattr(self, k, v)
 231             # TODO tener en cuenta servicios que hay que levantar y los que no
 232             if hasattr(self, 'commit'):
 233                 self.commit()
 234                 return False
 235             self._dump()
 236             if hasattr(self, '_write_config'):
 237                 self._write_config()
 238             if hasattr(self, 'reload'):
 239                 self.reload()
 240             return False
 241
 242 class ConfigWriter:
 243     r"""ConfigWriter([initd_name[, initd_dir]]) -> ConfigWriter.
 244
 245     This is a helper class to inherit from to automatically handle
 246     configuration generation. Mako template system is used for configuration
 247     files generation.
 248
 249     The configuration filenames, the generated configuration files directory
 250     and the templates directory can be defined by calling the constructor or
 251     in a more declarative way as class attributes, like:
 252
 253     class TestHandler(ConfigWriter):
 254         _config_writer_files = ('base.conf', 'custom.conf')
 255         _config_writer_cfg_dir = '/etc/service'
 256         _config_writer_tpl_dir = 'templates'
 257
 258     The generated configuration files directory defaults to '.' and the
 259     templates directory to 'templates'. _config_writer_files has no default and
 260     must be specified in either way. It can be string or a tuple if more than
 261     one configuration file must be generated.
 262
 263     The template filename and the generated configuration filename are both the
 264     same (so if you want to generate some /etc/config, you should have some
 265     templates/config template). That's why _config_writer_cfg_dir and
 266     _config_writer_tpl_dir can't be the same.
 267
 268     When you write your Handler, you should call _config_build_templates() in
 269     you Handler constructor to build the templates.
 270
 271     To write the configuration files, you must use the _write_config() method.
 272     To know what variables to replace in the template, you have to provide a
 273     method called _get_config_vars(tamplate_name), which should return a
 274     dictionary of variables to pass to the template system to be replaced in
 275     the template for the configuration file 'config_file'.
 276     """
 277     # TODO implement it using metaclasses to add the handlers method by demand
 278     # (only for specifieds commands).
 279
 280     _config_writer_files = ()
 281     _config_writer_cfg_dir = '.'
 282     _config_writer_tpl_dir = 'templates'
 283
 284     def __init__(self, files=None, cfg_dir=None, tpl_dir=None):
 285         r"Initialize the object, see the class documentation for details."
 286         if files is not None:
 287             self._config_writer_files = files
 288         if cfg_dir is not None:
 289             self._config_writer_cfg_dir = cfg_dir
 290         if tpl_dir is not None:
 291             self._config_writer_tpl_dir = tpl_dir
 292         self._config_build_templates()
 293
 294     def _config_build_templates(self):
 295         r"_config_writer_templates() -> None :: Build the template objects."
 296         if isinstance(self._config_writer_files, basestring):
 297             self._config_writer_files = (self._config_writer_files,)
 298         if not hasattr(self, '_config_writer_templates') \
 299                                         or not self._config_writer_templates:
 300             self._config_writer_templates = dict()
 301             for t in self._config_writer_files:
 302                 f = path.join(self._config_writer_tpl_dir, t)
 303                 self._config_writer_templates[t] = Template(filename=f)
 304
 305     def _render_config(self, template_name, vars=None):
 306         r"""_render_config(template_name[, config_filename[, vars]]).
 307
 308         Render a single config file using the template 'template_name'. If
 309         vars is specified, it's used as the dictionary with the variables
 310         to replace in the templates, if not, it looks for a
 311         _get_config_vars() method to get it.
 312         """
 313         if vars is None:
 314             if hasattr(self, '_get_config_vars'):
 315                 vars = self._get_config_vars(template_name)
 316             else:
 317                 vars = dict()
 318         elif callable(vars):
 319             vars = vars(template_name)
 320         return self._config_writer_templates[template_name].render(**vars)
 321
 322     def _write_single_config(self, template_name, config_filename=None, vars=None):
 323         r"""_write_single_config(template_name[, config_filename[, vars]]).
 324
 325         Write a single config file using the template 'template_name'. If no
 326         config_filename is specified, the config filename will be the same as
 327         the 'template_name' (but stored in the generated config files
 328         directory). If it's specified, the generated config file is stored in
 329         the file called 'config_filename' (also in the generated files
 330         directory). If vars is specified, it's used as the dictionary with the
 331         variables to replace in the templates, if not, it looks for a
 332         _get_config_vars() method to get it.
 333         """
 334         if not config_filename:
 335             config_filename = template_name
 336         if vars is None:
 337             if hasattr(self, '_get_config_vars'):
 338                 vars = self._get_config_vars(template_name)
 339             else:
 340                 vars = dict()
 341         elif callable(vars):
 342             vars = vars(template_name)
 343         f = file(path.join(self._config_writer_cfg_dir, config_filename), 'w')
 344         ctx = Context(f, **vars)
 345         self._config_writer_templates[template_name].render_context(ctx)
 346         f.close()
 347
 348     def _write_config(self):
 349         r"_write_config() -> None :: Generate all the configuration files."
 350         for t in self._config_writer_files:
 351             self._write_single_config(t)
 352
 353
 354 class ServiceHandler(Handler):
 355     r"""ServiceHandler([start[, stop[, restart[, reload]]]]) -> ServiceHandler.
 356
 357     This is a helper class to inherit from to automatically handle services
 358     with start, stop, restart, reload actions.
 359
 360     The actions can be defined by calling the constructor with all the
 361     parameters or in a more declarative way as class attributes, like:
 362
 363     class TestHandler(ServiceHandler):
 364         _service_start = ('command', 'start')
 365         _service_stop = ('command', 'stop')
 366         _service_restart = ('command', 'restart')
 367         _service_reload = 'reload-command'
 368
 369     Commands are executed without using the shell, that's why they are specified
 370     as tuples (where the first element is the command and the others are the
 371     command arguments). If only a command is needed (without arguments) a single
 372     string can be specified.
 373
 374     All commands must be specified.
 375     """
 376     # TODO implement it using metaclasses to add the handlers method by demand
 377     # (only for specifieds commands).
 378
 379     def __init__(self, start=None, stop=None, restart=None, reload=None):
 380         r"Initialize the object, see the class documentation for details."
 381         for (name, action) in dict(start=start, stop=stop, restart=restart,
 382                                                     reload=reload).items():
 383             if action is not None:
 384                 setattr(self, '_service_%s' % name, action)
 385
 386     @handler(u'Start the service.')
 387     def start(self):
 388         r"start() -> None :: Start the service."
 389         call(self._service_start)
 390
 391     @handler(u'Stop the service.')
 392     def stop(self):
 393         r"stop() -> None :: Stop the service."
 394         call(self._service_stop)
 395
 396     @handler(u'Restart the service.')
 397     def restart(self):
 398         r"restart() -> None :: Restart the service."
 399         call(self._service_restart)
 400
 401     @handler(u'Reload the service config (without restarting, if possible).')
 402     def reload(self):
 403         r"reload() -> None :: Reload the configuration of the service."
 404         call(self._service_reload)
 405
 406 class InitdHandler(Handler):
 407     r"""InitdHandler([initd_name[, initd_dir]]) -> InitdHandler.
 408
 409     This is a helper class to inherit from to automatically handle services
 410     with start, stop, restart, reload actions using a /etc/init.d like script.
 411
 412     The name and directory of the script can be defined by calling the
 413     constructor or in a more declarative way as class attributes, like:
 414
 415     class TestHandler(ServiceHandler):
 416         _initd_name = 'some-service'
 417         _initd_dir = '/usr/local/etc/init.d'
 418
 419     The default _initd_dir is '/etc/init.d', _initd_name has no default and
 420     must be specified in either way.
 421
 422     Commands are executed without using the shell.
 423     """
 424     # TODO implement it using metaclasses to add the handlers method by demand
 425     # (only for specifieds commands).
 426
 427     _initd_dir = '/etc/init.d'
 428
 429     def __init__(self, initd_name=None, initd_dir=None):
 430         r"Initialize the object, see the class documentation for details."
 431         if initd_name is not None:
 432             self._initd_name = initd_name
 433         if initd_dir is not None:
 434             self._initd_dir = initd_dir
 435
 436     @handler(u'Start the service.')
 437     def start(self):
 438         r"start() -> None :: Start the service."
 439         call((path.join(self._initd_dir, self._initd_name), 'start'))
 440
 441     @handler(u'Stop the service.')
 442     def stop(self):
 443         r"stop() -> None :: Stop the service."
 444         call((path.join(self._initd_dir, self._initd_name), 'stop'))
 445
 446     @handler(u'Restart the service.')
 447     def restart(self):
 448         r"restart() -> None :: Restart the service."
 449         call((path.join(self._initd_dir, self._initd_name), 'restart'))
 450
 451     @handler(u'Reload the service config (without restarting, if possible).')
 452     def reload(self):
 453         r"reload() -> None :: Reload the configuration of the service."
 454         call((path.join(self._initd_dir, self._initd_name), 'reload'))
 455
 456 class TransactionalHandler(Handler):
 457     r"""Handle command transactions providing a commit and rollback commands.
 458
 459     This is a helper class to inherit from to automatically handle
 460     transactional handlers, which have commit and rollback commands.
 461
 462     The handler should provide a reload() method (see ServiceHandler and
 463     InitdHandler for helper classes to provide this) which will be called
 464     when a commit command is issued (if a reload() command is present).
 465     The persistent data will be written too (if a _dump() method is provided,
 466     see Persistent and Restorable for that), and the configuration files
 467     will be generated (if a _write_config method is present, see ConfigWriter).
 468     """
 469     # TODO implement it using metaclasses to add the handlers method by demand
 470     # (only for specifieds commands).
 471
 472     @handler(u'Commit the changes (reloading the service, if necessary).')
 473     def commit(self):
 474         r"commit() -> None :: Commit the changes and reload the service."
 475         if hasattr(self, '_dump'):
 476             self._dump()
 477         if hasattr(self, '_write_config'):
 478             self._write_config()
 479         if hasattr(self, 'reload'):
 480             self.reload()
 481
 482     @handler(u'Discard all the uncommited changes.')
 483     def rollback(self):
 484         r"rollback() -> None :: Discard the changes not yet commited."
 485         if hasattr(self, '_load'):
 486             self._load()
 487
 488 class ParametersHandler(Handler):
 489     r"""ParametersHandler([attr]) -> ParametersHandler.
 490
 491     This is a helper class to inherit from to automatically handle
 492     service parameters, providing set, get, list and show commands.
 493
 494     The attribute that holds the parameters can be defined by calling the
 495     constructor or in a more declarative way as class attributes, like:
 496
 497     class TestHandler(ServiceHandler):
 498         _parameters_attr = 'some_attr'
 499
 500     The default is 'params' and it should be a dictionary.
 501     """
 502     # TODO implement it using metaclasses to add the handlers method by demand
 503     # (only for specifieds commands).
 504
 505     _parameters_attr = 'params'
 506
 507     def __init__(self, attr=None):
 508         r"Initialize the object, see the class documentation for details."
 509         if attr is not None:
 510             self._parameters_attr = attr
 511
 512     @handler(u'Set a service parameter.')
 513     def set(self, param, value):
 514         r"set(param, value) -> None :: Set a service parameter."
 515         if not param in self.params:
 516             raise ParameterNotFoundError(param)
 517         self.params[param] = value
 518
 519     @handler(u'Get a service parameter.')
 520     def get(self, param):
 521         r"get(param) -> None :: Get a service parameter."
 522         if not param in self.params:
 523             raise ParameterNotFoundError(param)
 524         return self.params[param]
 525
 526     @handler(u'List all available service parameters.')
 527     def list(self):
 528         r"list() -> tuple :: List all the parameter names."
 529         return self.params.keys()
 530
 531     @handler(u'Get all service parameters, with their values.')
 532     def show(self):
 533         r"show() -> (key, value) tuples :: List all the parameters."
 534         return self.params.items()
 535
 536
 537 if __name__ == '__main__':
 538
 539     # Execution tests
 540     class STestHandler1(ServiceHandler):
 541         _service_start = ('service', 'start')
 542         _service_stop = ('service', 'stop')
 543         _service_restart = ('ls', '/')
 544         _service_reload = ('cp', '/la')
 545     class STestHandler2(ServiceHandler):
 546         def __init__(self):
 547             ServiceHandler.__init__(self, 'cmd-start', 'cmd-stop',
 548                                         'cmd-restart', 'cmd-reload')
 549     class ITestHandler1(InitdHandler):
 550         _initd_name = 'test1'
 551     class ITestHandler2(InitdHandler):
 552         def __init__(self):
 553             InitdHandler.__init__(self, 'test2', '/usr/local/etc/init.d')
 554     handlers = [
 555         STestHandler1(),
 556         STestHandler2(),
 557         ITestHandler1(),
 558         ITestHandler2(),
 559     ]
 560     for h in handlers:
 561         print h.__class__.__name__
 562         try:
 563             h.start()
 564         except ExecutionError, e:
 565             print e
 566         try:
 567             h.stop()
 568         except ExecutionError, e:
 569             print e
 570         try:
 571             h.restart()
 572         except ExecutionError, e:
 573             print e
 574         try:
 575             h.reload()
 576         except ExecutionError, e:
 577             print e
 578         print
 579
 580     # Persistent test
 581     print 'PTestHandler'
 582     class PTestHandler(Persistent):
 583         _persistent_attrs = 'vars'
 584         def __init__(self):
 585             self.vars = dict(a=1, b=2)
 586     h = PTestHandler()
 587     print h.vars
 588     h._dump()
 589     h.vars['x'] = 100
 590     print h.vars
 591     h._load()
 592     print h.vars
 593     h.vars['x'] = 100
 594     h._dump()
 595     print h.vars
 596     del h.vars['x']
 597     print h.vars
 598     h._load()
 599     print h.vars
 600     print
 601
 602     # Restorable test
 603     print 'RTestHandler'
 604     class RTestHandler(Restorable):
 605         _persistent_attrs = 'vars'
 606         _restorable_defaults = dict(vars=dict(a=1, b=2))
 607         def __init__(self):
 608             self._restore()
 609     h = RTestHandler()
 610     print h.vars
 611     h.vars['x'] = 100
 612     h._dump()
 613     h = RTestHandler()
 614     print h.vars
 615     print
 616
 617     # ConfigWriter test
 618     print 'CTestHandler'
 619     import os
 620     os.mkdir('templates')
 621     f = file('templates/config', 'w')
 622     f.write('Hello, ${name}! You are ${what}.')
 623     f.close()
 624     print 'template:'
 625     print file('templates/config').read()
 626     class CTestHandler(ConfigWriter):
 627         _config_writer_files = 'config'
 628         def __init__(self):
 629             self._config_build_templates()
 630         def _get_config_vars(self, config_file):
 631             return dict(name='you', what='a parrot')
 632     h = CTestHandler()
 633     h._write_config()
 634     print 'config:'
 635     print file('config').read()
 636     os.unlink('config')
 637     os.unlink('templates/config')
 638     os.rmdir('templates')
 639     print
 640