/neatx/lib/utils.py

  1#
  2#
  3
  4# Copyright (C) 2009 Google Inc.
  5#
  6# This program is free software; you can redistribute it and/or modify
  7# it under the terms of the GNU General Public License as published by
  8# the Free Software Foundation; either version 2 of the License, or
  9# (at your option) any later version.
 10#
 11# This program is distributed in the hope that it will be useful, but
 12# WITHOUT ANY WARRANTY; without even the implied warranty of
 13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 14# General Public License for more details.
 15#
 16# You should have received a copy of the GNU General Public License
 17# along with this program; if not, write to the Free Software
 18# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 19# 02110-1301, USA.
 20
 21
 22"""Module for utility functions"""
 23
 24
 25import errno
 26import fcntl
 27import logging
 28import logging.handlers
 29import os
 30import os.path
 31import pwd
 32import resource
 33import re
 34import signal
 35import sys
 36import syslog
 37import tempfile
 38import termios
 39import time
 40
 41from neatx import constants
 42from neatx import errors
 43
 44
 45_SHELL_UNQUOTED_RE = re.compile('^[-.,=:/_+@A-Za-z0-9]+$')
 46
 47try:
 48  DEV_NULL = os.devnull
 49except AttributeError:
 50  DEV_NULL = "/dev/null"
 51
 52
 53class CustomSyslogHandler(logging.Handler):
 54  """Custom syslog handler.
 55
 56  logging.handlers.SysLogHandler doesn't support splitting exceptions into
 57  several lines.
 58
 59  This class should only be used once per process.
 60
 61  """
 62  _LEVEL_MAP = {
 63    logging.CRITICAL: syslog.LOG_CRIT,
 64    logging.ERROR: syslog.LOG_ERR,
 65    logging.WARNING: syslog.LOG_WARNING,
 66    logging.INFO: syslog.LOG_INFO,
 67    logging.DEBUG: syslog.LOG_DEBUG,
 68    }
 69
 70  def __init__(self, ident):
 71    """Initializes instances.
 72
 73    @type ident: string
 74    @param ident: String prepended to every message
 75
 76    """
 77    logging.Handler.__init__(self)
 78
 79    syslog.openlog(ident, syslog.LOG_PID, syslog.LOG_USER)
 80
 81  def close(self):
 82    syslog.closelog()
 83
 84  def _MapLogLevel(self, levelno):
 85    """Maps log level to syslog.
 86
 87    @type levelno: int
 88    @param levelno: Log level
 89
 90    Default is LOG_DEBUG.
 91
 92    """
 93    return self._LEVEL_MAP.get(levelno, syslog.LOG_DEBUG)
 94
 95  def emit(self, record):
 96    """Send a log record to syslog.
 97
 98    @type record: logging.LogRecord
 99    @param record: Log record
100
101    """
102    msg = self.format(record)
103
104    if record.exc_info:
105      messages = msg.split(os.linesep)
106    else:
107      messages = [msg]
108
109    priority = self._MapLogLevel(record.levelno)
110
111    for msg in messages:
112      syslog.syslog(priority, msg)
113
114
115class LoggingSetupOptions(object):
116  def __init__(self, debug, logtostderr):
117    """Initializes logging setup options class.
118
119    @type debug: bool
120    @param debug: Whether to enable debug log messages
121    @type logtostderr: bool
122    @param logtostderr: Whether to write log messages to stderr
123
124    """
125    self.debug = debug
126    self.logtostderr = logtostderr
127
128
129class LoggingSetup(object):
130  """Logging setup class.
131
132  This class should only be used once per process.
133
134  """
135  def __init__(self, program):
136    """Configures the logging module
137
138    @type program: str
139    @param program: the name under which we should log messages
140
141    """
142    self._program = program
143    self._options = None
144    self._root_logger = None
145    self._stderr_handler = None
146    self._syslog_handler = None
147
148  def Init(self):
149    """Initializes the logging module.
150
151    """
152    assert not self._root_logger
153    assert not self._stderr_handler
154    assert not self._syslog_handler
155
156    # Get root logger
157    self._root_logger = self._InitRootLogger()
158
159    # Create stderr handler
160    self._stderr_handler = logging.StreamHandler(sys.stderr)
161
162    # Create syslog handler
163    self._syslog_handler = CustomSyslogHandler(self._program)
164
165    self._ConfigureHandlers()
166
167  def SetOptions(self, options):
168    """Configure logging setup.
169
170    @type options: L{LoggingSetupOptions}
171    @param options: Configuration object
172
173    """
174    self._options = options
175    self._ConfigureHandlers()
176
177  @staticmethod
178  def _InitRootLogger():
179    """Initializes and returns the root logger.
180
181    """
182    root_logger = logging.getLogger("")
183    root_logger.setLevel(logging.NOTSET)
184
185    # Remove all previously setup handlers
186    for handler in root_logger.handlers:
187      root_logger.removeHandler(handler)
188      handler.close()
189
190    return root_logger
191
192  def _ConfigureHandlers(self):
193    """Set formatters and levels for handlers.
194
195    """
196    stderr_level = None
197    syslog_level = None
198
199    if self._options is None:
200      # Log error and above
201      stderr_level = logging.ERROR
202    else:
203      if self._options.debug:
204        # Log everything
205        level = logging.NOTSET
206      else:
207        # Log info and above (e.g. error)
208        level = logging.INFO
209
210      if self._options.logtostderr:
211        stderr_level = level
212      else:
213        syslog_level = level
214
215    # Configure handlers
216    self._ConfigureSingleHandler(self._stderr_handler, stderr_level, False)
217    self._ConfigureSingleHandler(self._syslog_handler, syslog_level, True)
218
219  def _ConfigureSingleHandler(self, handler, level, is_syslog):
220    """Configures a handler based on the parameters.
221
222    """
223    if level is None:
224      self._root_logger.removeHandler(handler)
225      return
226
227    debug = self._options is not None and self._options.debug
228
229    fmt = self._GetMessageFormat(self._program, is_syslog, debug)
230
231    handler.setLevel(level)
232    handler.setFormatter(logging.Formatter(fmt))
233
234    self._root_logger.addHandler(handler)
235
236  @staticmethod
237  def _GetMessageFormat(program, is_syslog, debug):
238    """Returns message format.
239
240    """
241    assert "%" not in program
242
243    if is_syslog:
244      fmt = ""
245    else:
246      fmt = "%(asctime)s: " + program + " pid=%(process)d "
247
248    fmt += "%(levelname)s"
249
250    if debug:
251      fmt += " %(module)s:%(lineno)s"
252
253    fmt += " %(message)s"
254
255    return fmt
256
257
258def WithoutTerminalEcho(fd, fn, *args, **kwargs):
259  """Calls function with ECHO flag disabled on passed file descriptor.
260
261  @type fd: file or int
262  @param fd: File descriptor
263  @type fn: callable
264  @param fn: Called function
265
266  """
267  assert callable(fn)
268
269  # Keep old terminal settings
270  try:
271    old = termios.tcgetattr(fd)
272  except termios.error, err:
273    if err.args[0] not in (errno.ENOTTY, errno.EINVAL):
274      raise
275    old = None
276
277  if old is not None:
278    new = old[:]
279
280    # Disable the echo flag in lflags (index 3)
281    new[3] &= ~termios.ECHO
282
283    termios.tcsetattr(fd, termios.TCSADRAIN, new)
284
285  try:
286    return fn(*args, **kwargs)
287  finally:
288    if old is not None:
289      termios.tcsetattr(fd, termios.TCSADRAIN, old)
290
291
292def NormalizeSpace(text):
293  """Replace all whitespace (\\s) with a single space.
294
295  Whitespace at start and end are also removed.
296
297  """
298  return re.sub(r"\s+", " ", text).strip()
299
300
301def RemoveFile(filename):
302  """Remove a file ignoring some errors.
303
304  Remove a file, ignoring non-existing ones or directories. Other
305  errors are passed.
306
307  @type filename: str
308  @param filename: the file to be removed
309
310  """
311  try:
312    os.unlink(filename)
313  except OSError, err:
314    if err.errno not in (errno.ENOENT, errno.EISDIR):
315      raise
316
317
318def SetCloseOnExecFlag(fd, enable):
319  """Sets or unsets the close-on-exec flag on a file descriptor.
320
321  @type fd: int
322  @param fd: File descriptor
323  @type enable: bool
324  @param enable: Whether to set or unset it.
325
326  """
327  flags = fcntl.fcntl(fd, fcntl.F_GETFD)
328
329  if enable:
330    flags |= fcntl.FD_CLOEXEC
331  else:
332    flags &= ~fcntl.FD_CLOEXEC
333
334  fcntl.fcntl(fd, fcntl.F_SETFD, flags)
335
336
337def SetNonblockFlag(fd, enable):
338  """Sets or unsets the O_NONBLOCK flag on on a file descriptor.
339
340  @type fd: int
341  @param fd: File descriptor
342  @type enable: bool
343  @param enable: Whether to set or unset it
344
345  """
346  flags = fcntl.fcntl(fd, fcntl.F_GETFL)
347
348  if enable:
349    flags |= os.O_NONBLOCK
350  else:
351    flags &= ~os.O_NONBLOCK
352
353  fcntl.fcntl(fd, fcntl.F_SETFL, flags)
354
355
356def ListVisibleFiles(path):
357  """Returns a list of visible files in a directory.
358
359  @type path: str
360  @param path: the directory to enumerate
361  @rtype: list
362  @return: the list of all files not starting with a dot
363
364  """
365  files = [i for i in os.listdir(path) if not i.startswith(".")]
366  files.sort()
367  return files
368
369
370def WriteFile(file_name, fn=None, data=None,
371              mode=None, uid=-1, gid=-1):
372  """(Over)write a file atomically.
373
374  The file_name and either fn (a function taking one argument, the
375  file descriptor, and which should write the data to it) or data (the
376  contents of the file) must be passed. The other arguments are
377  optional and allow setting the file mode, owner and group of the file.
378
379  If the function (WriteFile) doesn't raise an exception, it has succeeded and
380  the target file has the new contents. If the function has raised an
381  exception, an existing target file should be unmodified and the temporary
382  file should be removed.
383
384  @type file_name: str
385  @param file_name: the target filename
386  @type fn: callable
387  @param fn: content writing function, called with
388      file descriptor as parameter
389  @type data: str
390  @param data: contents of the file
391  @type mode: int
392  @param mode: file mode
393  @type uid: int
394  @param uid: the owner of the file
395  @type gid: int
396  @param gid: the group of the file
397
398  @raise errors.ProgrammerError: if any of the arguments are not valid
399
400  """
401  if not os.path.isabs(file_name):
402    raise errors.ProgrammerError("Path passed to WriteFile is not"
403                                 " absolute: '%s'" % file_name)
404
405  if [fn, data].count(None) != 1:
406    raise errors.ProgrammerError("fn or data required")
407
408  dir_name, base_name = os.path.split(file_name)
409  fd, new_name = tempfile.mkstemp(prefix=".tmp", suffix=base_name,
410                                  dir=dir_name)
411  try:
412    if uid != -1 or gid != -1:
413      os.chown(new_name, uid, gid)
414    if mode:
415      os.chmod(new_name, mode)
416    if data is not None:
417      os.write(fd, data)
418    else:
419      fn(fd)
420    os.fsync(fd)
421    os.rename(new_name, file_name)
422  finally:
423    os.close(fd)
424    # Make sure temporary file is removed in any case
425    RemoveFile(new_name)
426
427
428def FormatTable(data, columns):
429  """Formats a list of input data as a table.
430
431  Columns must be passed via the C{columns} parameter. It must be a list of
432  tuples, each containing the column caption, width and a function to retrieve
433  the value. If the width is negative, the value is aligned to the right. The
434  column function is called for every item.
435
436  Example:
437    >>> columns = [
438      ("Name", 10, lambda item: item[0]),
439      ("Value", -5, lambda item: item[1]),
440      ]
441    >>> data = [("Row%d" % i, i) for i in xrange(3)]
442    >>> print "\\n".join(utils.FormatTable(data, columns))
443    Name       Value
444    ---------- -----
445    Row0           0
446    Row1           1
447    Row2           2
448
449  @type data: list
450  @param data: Input data
451  @type columns: list of tuples
452  @param columns: Column definitions
453  @rtype: list of strings
454  @return: Rows as strings
455
456  """
457  col_width = []
458  header_row = []
459  dashes_row = []
460  format_fields = []
461
462  for idx, (header, width, _) in enumerate(columns):
463    if idx == (len(columns) - 1) and width >= 0:
464      # Last column
465      col_width.append(None)
466      fmt = "%s"
467    else:
468      col_width.append(abs(width))
469      if width < 0:
470        fmt = "%*s"
471      else:
472        fmt = "%-*s"
473
474    format_fields.append(fmt)
475
476    if col_width[idx] is not None:
477      header_row.append(col_width[idx])
478      dashes_row.append(col_width[idx])
479
480    header_row.append(header)
481    dashes_row.append(abs(width) * "-")
482
483  format = " ".join(format_fields)
484
485  rows = [header_row, dashes_row]
486  for item in data:
487    row = []
488    for idx, (_, width, fn) in enumerate(columns):
489      if col_width[idx] is not None:
490        row.append(col_width[idx])
491      row.append(fn(item))
492    rows.append(row)
493
494  return [format % tuple(row) for row in rows]
495
496
497class RetryTimeout(Exception):
498  """Retry loop timed out.
499
500  """
501
502
503class RetryAgain(Exception):
504  """Retry again.
505
506  """
507
508
509def Retry(fn, start, factor, limit, timeout, _time=time):
510  """Call a function repeatedly until it succeeds.
511
512  The function C{fn} is called repeatedly until it doesn't throw L{RetryAgain}
513  anymore. Between calls a delay starting at C{start} and multiplied by
514  C{factor} on each run until it's above C{limit} is inserted. After a total of
515  C{timeout} seconds, the retry loop fails with L{RetryTimeout}.
516
517  @type fn: callable
518  @param fn: Function to be called (no parameters supported)
519  @type start: float
520  @param start: Initial value for delay
521  @type factor: float
522  @param factor: Factor for delay increase
523  @type limit: float
524  @param limit: Upper limit for delay
525  @type timeout: float
526  @param timeout: Total timeout
527  @return: Return value of function
528
529  """
530  assert start > 0
531  assert factor > 1.0
532  assert limit >= 0
533
534  end_time = _time.time() + timeout
535  delay = start
536
537  while True:
538    try:
539      return fn()
540    except RetryAgain:
541      pass
542
543    if _time.time() > end_time:
544      raise RetryTimeout()
545
546    _time.sleep(delay)
547
548    if delay < limit:
549      delay *= factor
550
551
552def ShellQuote(value):
553  """Quotes shell argument according to POSIX.
554
555  @type value: str
556  @param value: the argument to be quoted
557  @rtype: str
558  @return: the quoted value
559
560  """
561  if _SHELL_UNQUOTED_RE.match(value):
562    return value
563  else:
564    return "'%s'" % value.replace("'", "'\\''")
565
566
567def ShellQuoteArgs(args):
568  """Quotes a list of shell arguments.
569
570  @type args: list
571  @param args: list of arguments to be quoted
572  @rtype: str
573  @return: the quoted arguments concatenaned with spaces
574
575  """
576  return " ".join([ShellQuote(i) for i in args])
577
578
579def CloseFd(fd, retries=5):
580  """Close a file descriptor ignoring errors.
581
582  @type fd: int
583  @param fd: File descriptor
584  @type retries: int
585  @param retries: How many retries to make, in case we get any
586    other error than EBADF
587
588  """
589  while retries > 0:
590    retries -= 1
591    try:
592      os.close(fd)
593    except OSError, err:
594      if err.errno != errno.EBADF:
595        continue
596    break
597
598
599def GetMaxFd():
600  """Determine max file descriptor number.
601
602  @rtype: int
603  @return: Max file descriptor number
604
605  """
606  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
607  if maxfd == resource.RLIM_INFINITY:
608    # Default maximum for the number of available file descriptors.
609    maxfd = 1024
610
611    try:
612      maxfd = os.sysconf("SC_OPEN_MAX")
613    except ValueError:
614      pass
615
616  if maxfd < 0:
617    maxfd = 1024
618
619  return maxfd
620
621
622def StartDaemon(fn):
623  """Start a daemon process.
624
625  Starts a daemon process by double-forking and invoking a function. The
626  function should then use exec*(2) to start the process.
627
628  """
629  # First fork
630  pid = os.fork()
631  if pid != 0:
632    # Parent process
633
634    # Try to avoid zombies
635    try:
636      os.waitpid(pid, 0)
637    except OSError:
638      pass
639
640    return
641
642  # First child process
643  os.setsid()
644
645  # Second fork
646  pid = os.fork()
647  if pid != 0:
648    # Second parent process
649    os._exit(0)
650
651  # Second child process
652  os.chdir("/")
653  os.umask(077)
654
655  # Close all file descriptors
656  for fd in xrange(GetMaxFd()):
657    CloseFd(fd)
658
659  # Open /dev/null
660  fd = os.open(DEV_NULL, os.O_RDWR)
661
662  # Redirect stdio to /dev/null
663  os.dup2(fd, constants.STDIN_FILENO)
664  os.dup2(fd, constants.STDOUT_FILENO)
665  os.dup2(fd, constants.STDERR_FILENO)
666
667  try:
668    # Call function starting daemon
669    fn()
670    os._exit(0)
671  except (SystemExit, KeyboardInterrupt):
672    raise
673  except:
674    os._exit(1)
675
676
677def CallWithSignalHandlers(sigtbl, fn, *args, **kwargs):
678  previous = {}
679  try:
680    for (signum, handler) in sigtbl.iteritems():
681      # Setup handler
682      prev_handler = signal.signal(signum, handler)
683      try:
684        previous[signum] = prev_handler
685      except Exception:
686        # Restore previous handler
687        signal.signal(signum, prev_handler)
688        raise
689
690    return fn(*args, **kwargs)
691
692  finally:
693    for (signum, prev_handler) in previous.items():
694      signal.signal(signum, prev_handler)
695
696      # If successful, remove from dict
697      del previous[signum]
698
699    assert not previous
700
701
702def _GetSignalNumberTable(_signal=signal):
703  table = {}
704
705  for name in dir(_signal):
706    if name.startswith("SIG") and not name.startswith("SIG_"):
707      signum = getattr(_signal, name)
708      if isinstance(signum, (int, long)):
709        table[signum] = name
710
711  return table
712
713
714def GetSignalName(signum, _signal=signal):
715  """Returns signal name by signal number.
716
717  If the signal number is not known, "Signal 123" is returned (with the passed
718  signal number instead of 123).
719
720  @type signum: int
721  @param signum: Signal number
722  @rtype: str
723  @return: Signal name
724
725  """
726  # This table could be cached
727  table = _GetSignalNumberTable(_signal=_signal)
728
729  try:
730    return table[signum]
731  except KeyError:
732    return "Signal %s" % signum
733
734
735def _ConvertVersionPart(value):
736  """Tries to convert a value to an integer.
737
738  """
739  try:
740    return int(value)
741  except ValueError:
742    return value
743
744
745def _GetVersionSplitter(sep, count):
746  """Returns callable to split wanted parts from version.
747
748  @type sep: string
749  @param sep: String of separator characters
750  @type count: int
751  @param count: How many parts to return
752
753  """
754  assert sep
755  assert count == -1 or count > 0
756
757  # str.split is a lot faster but doesn't provide the right semantics when
758  # the caller wants more than one possible separator.
759  #if len(sep) == 1:
760  #  if count == -1:
761  #    return lambda ver: ver.split(sep)
762  #  else:
763  #    return lambda ver: ver.split(sep, count)[:count]
764
765  re_split = re.compile("[%s]" % re.escape(sep)).split
766  if count == -1:
767    return lambda ver: re_split(ver)
768  else:
769    return lambda ver: re_split(ver, count)[:count]
770
771
772def GetVersionComparator(sep, count=-1):
773  """Returns a cmp-compatible function to compare two version strings.
774
775  @type sep: string
776  @param sep: String of separator characters, similar to strsep(3)
777  @type count: int
778  @param count: How many parts to compare (-1 for all)
779
780  """
781  # TODO: Support for versions such as "1.2~alpha0"
782
783  split_fn = _GetVersionSplitter(sep, count)
784  split_version = lambda ver: map(_ConvertVersionPart, split_fn(ver))
785
786  return lambda x, y: cmp(split_version(x), split_version(y))
787
788
789def ParseVersion(version, sep, digits):
790  """Parses a version and converts it to a number.
791
792  Example:
793    >>> ParseVersion("3.3.7.2-1", ".-", [2, 2, 2])
794    30307
795    >>> ParseVersion("3.3.5.2-1", ".-", [2, 2, 4])
796    3030005
797    >>> ParseVersion("3.3.9.2-1", ".-", [2, 2, 2, 2, 2])
798    303090201
799    >>> ParseVersion("12.1", ".-", [2, 2, 2, 2])
800    12010000
801    >>> ParseVersion("23.193", ".-", [2, 4])
802    230193
803
804  @type version: str
805  @param version: Version string
806  @type sep: string
807  @param sep: String of separator characters, similar to strsep(3)
808  @type digits: list
809  @param digits: List of digits to be used per part
810
811  """
812  split_fn = _GetVersionSplitter(sep, len(digits))
813  parts = split_fn(version)
814
815  version = 0
816  total_exp = 0
817
818  for idx, exp in reversed(list(enumerate(digits))):
819    try:
820      value = int(parts[idx])
821    except IndexError:
822      value = 0
823
824    if value > 10 ** exp:
825      raise ValueError("Version part %s (%r) too long for %s digits" %
826                       (idx, value, exp))
827
828    version += (10 ** total_exp) * value
829
830    total_exp += exp
831
832  return version
833
834
835def FormatVersion(version, sep, digits):
836  """Format version number.
837
838  @type version: int
839  @param version: Version number (as returned by L{ParseVersion})
840  @type sep: string
841  @param sep: Separator string between digits
842  @type digits: list
843  @param digits: List of digits used per part in the version numbers
844
845  """
846  # TODO: Implement support for more than one separator
847  parts = []
848  next = version
849
850  for exp in reversed(digits):
851    (next, value) = divmod(next, 10 ** exp)
852
853    parts.append(str(value))
854
855  if next > 0:
856    raise ValueError("Invalid version number (%r) for given digits (%r)" %
857                     (version, digits))
858
859  parts.reverse()
860
861  return sep.join(parts)
862
863
864def LogFunctionWithPrefix(fn, prefix):
865  return lambda msg, *args, **kwargs: fn(prefix + msg, *args, **kwargs)
866
867
868def GetExitcodeSignal(status):
869  if status < 0:
870    return (None, -status)
871  else:
872    return (status, None)
873
874
875def GetCurrentUserName():
876  """Returns the name of the user that owns the current process.
877
878  """
879  return pwd.getpwuid(os.getuid())[0]