code.grnet.gr Git - ganeti-local/blob - lib/utils.py

   1 #
   2 #
   3
   4 # Copyright (C) 2006, 2007 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 """Ganeti small utilities
  23
  24 """
  25
  26
  27 import sys
  28 import os
  29 import sha
  30 import time
  31 import subprocess
  32 import re
  33 import socket
  34 import tempfile
  35 import shutil
  36 import errno
  37 import pwd
  38 import itertools
  39 import select
  40 import fcntl
  41
  42 from cStringIO import StringIO
  43
  44 from ganeti import logger
  45 from ganeti import errors
  46 from ganeti import constants
  47
  48
  49 _locksheld = []
  50 _re_shell_unquoted = re.compile('^[-.,=:/_+@A-Za-z0-9]+$')
  51
  52 debug = False
  53
  54 class RunResult(object):
  55   """Simple class for holding the result of running external programs.
  56
  57   Instance variables:
  58     exit_code: the exit code of the program, or None (if the program
  59                didn't exit())
  60     signal: numeric signal that caused the program to finish, or None
  61             (if the program wasn't terminated by a signal)
  62     stdout: the standard output of the program
  63     stderr: the standard error of the program
  64     failed: a Boolean value which is True in case the program was
  65             terminated by a signal or exited with a non-zero exit code
  66     fail_reason: a string detailing the termination reason
  67
  68   """
  69   __slots__ = ["exit_code", "signal", "stdout", "stderr",
  70                "failed", "fail_reason", "cmd"]
  71
  72
  73   def __init__(self, exit_code, signal, stdout, stderr, cmd):
  74     self.cmd = cmd
  75     self.exit_code = exit_code
  76     self.signal = signal
  77     self.stdout = stdout
  78     self.stderr = stderr
  79     self.failed = (signal is not None or exit_code != 0)
  80
  81     if self.signal is not None:
  82       self.fail_reason = "terminated by signal %s" % self.signal
  83     elif self.exit_code is not None:
  84       self.fail_reason = "exited with exit code %s" % self.exit_code
  85     else:
  86       self.fail_reason = "unable to determine termination reason"
  87
  88     if debug and self.failed:
  89       logger.Debug("Command '%s' failed (%s); output: %s" %
  90                    (self.cmd, self.fail_reason, self.output))
  91
  92   def _GetOutput(self):
  93     """Returns the combined stdout and stderr for easier usage.
  94
  95     """
  96     return self.stdout + self.stderr
  97
  98   output = property(_GetOutput, None, None, "Return full output")
  99
 100
 101 def _GetLockFile(subsystem):
 102   """Compute the file name for a given lock name."""
 103   return "%s/ganeti_lock_%s" % (constants.LOCK_DIR, subsystem)
 104
 105
 106 def Lock(name, max_retries=None, debug=False):
 107   """Lock a given subsystem.
 108
 109   In case the lock is already held by an alive process, the function
 110   will sleep indefintely and poll with a one second interval.
 111
 112   When the optional integer argument 'max_retries' is passed with a
 113   non-zero value, the function will sleep only for this number of
 114   times, and then it will will raise a LockError if the lock can't be
 115   acquired. Passing in a negative number will cause only one try to
 116   get the lock. Passing a positive number will make the function retry
 117   for approximately that number of seconds.
 118
 119   """
 120   lockfile = _GetLockFile(name)
 121
 122   if name in _locksheld:
 123     raise errors.LockError('Lock "%s" already held!' % (name,))
 124
 125   errcount = 0
 126
 127   retries = 0
 128   while True:
 129     try:
 130       fd = os.open(lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_SYNC)
 131       break
 132     except OSError, creat_err:
 133       if creat_err.errno != errno.EEXIST:
 134         raise errors.LockError("Can't create the lock file. Error '%s'." %
 135                                str(creat_err))
 136
 137       try:
 138         pf = open(lockfile, 'r')
 139       except IOError, open_err:
 140         errcount += 1
 141         if errcount >= 5:
 142           raise errors.LockError("Lock file exists but cannot be opened."
 143                                  " Error: '%s'." % str(open_err))
 144         time.sleep(1)
 145         continue
 146
 147       try:
 148         pid = int(pf.read())
 149       except ValueError:
 150         raise errors.LockError("Invalid pid string in %s" %
 151                                (lockfile,))
 152
 153       if not IsProcessAlive(pid):
 154         raise errors.LockError("Stale lockfile %s for pid %d?" %
 155                                (lockfile, pid))
 156
 157       if max_retries and max_retries <= retries:
 158         raise errors.LockError("Can't acquire lock during the specified"
 159                                " time, aborting.")
 160       if retries == 5 and (debug or sys.stdin.isatty()):
 161         logger.ToStderr("Waiting for '%s' lock from pid %d..." % (name, pid))
 162
 163       time.sleep(1)
 164       retries += 1
 165       continue
 166
 167   os.write(fd, '%d\n' % (os.getpid(),))
 168   os.close(fd)
 169
 170   _locksheld.append(name)
 171
 172
 173 def Unlock(name):
 174   """Unlock a given subsystem.
 175
 176   """
 177   lockfile = _GetLockFile(name)
 178
 179   try:
 180     fd = os.open(lockfile, os.O_RDONLY)
 181   except OSError:
 182     raise errors.LockError('Lock "%s" not held.' % (name,))
 183
 184   f = os.fdopen(fd, 'r')
 185   pid_str = f.read()
 186
 187   try:
 188     pid = int(pid_str)
 189   except ValueError:
 190     raise errors.LockError('Unable to determine PID of locking process.')
 191
 192   if pid != os.getpid():
 193     raise errors.LockError('Lock not held by me (%d != %d)' %
 194                            (os.getpid(), pid,))
 195
 196   os.unlink(lockfile)
 197   _locksheld.remove(name)
 198
 199
 200 def LockCleanup():
 201   """Remove all locks.
 202
 203   """
 204   for lock in _locksheld:
 205     Unlock(lock)
 206
 207
 208 def RunCmd(cmd):
 209   """Execute a (shell) command.
 210
 211   The command should not read from its standard input, as it will be
 212   closed.
 213
 214   Args:
 215     cmd: command to run. (str)
 216
 217   Returns: `RunResult` instance
 218
 219   """
 220   if isinstance(cmd, list):
 221     cmd = [str(val) for val in cmd]
 222     strcmd = " ".join(cmd)
 223     shell = False
 224   else:
 225     strcmd = cmd
 226     shell = True
 227   env = os.environ.copy()
 228   env["LC_ALL"] = "C"
 229   poller = select.poll()
 230   child = subprocess.Popen(cmd, shell=shell,
 231                            stderr=subprocess.PIPE,
 232                            stdout=subprocess.PIPE,
 233                            stdin=subprocess.PIPE,
 234                            close_fds=True, env=env)
 235
 236   child.stdin.close()
 237   poller.register(child.stdout, select.POLLIN)
 238   poller.register(child.stderr, select.POLLIN)
 239   out = StringIO()
 240   err = StringIO()
 241   fdmap = {
 242     child.stdout.fileno(): (out, child.stdout),
 243     child.stderr.fileno(): (err, child.stderr),
 244     }
 245   for fd in fdmap:
 246     status = fcntl.fcntl(fd, fcntl.F_GETFL)
 247     fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
 248
 249   while fdmap:
 250     for fd, event in poller.poll():
 251       if event & select.POLLIN or event & select.POLLPRI:
 252         data = fdmap[fd][1].read()
 253         # no data from read signifies EOF (the same as POLLHUP)
 254         if not data:
 255           poller.unregister(fd)
 256           del fdmap[fd]
 257           continue
 258         fdmap[fd][0].write(data)
 259       if (event & select.POLLNVAL or event & select.POLLHUP or
 260           event & select.POLLERR):
 261         poller.unregister(fd)
 262         del fdmap[fd]
 263
 264   out = out.getvalue()
 265   err = err.getvalue()
 266
 267   status = child.wait()
 268   if status >= 0:
 269     exitcode = status
 270     signal = None
 271   else:
 272     exitcode = None
 273     signal = -status
 274
 275   return RunResult(exitcode, signal, out, err, strcmd)
 276
 277
 278 def RunCmdUnlocked(cmd):
 279   """Execute a shell command without the 'cmd' lock.
 280
 281   This variant of `RunCmd()` drops the 'cmd' lock before running the
 282   command and re-aquires it afterwards, thus it can be used to call
 283   other ganeti commands.
 284
 285   The argument and return values are the same as for the `RunCmd()`
 286   function.
 287
 288   Args:
 289     cmd - command to run. (str)
 290
 291   Returns:
 292     `RunResult`
 293
 294   """
 295   Unlock('cmd')
 296   ret = RunCmd(cmd)
 297   Lock('cmd')
 298
 299   return ret
 300
 301
 302 def RemoveFile(filename):
 303   """Remove a file ignoring some errors.
 304
 305   Remove a file, ignoring non-existing ones or directories. Other
 306   errors are passed.
 307
 308   """
 309   try:
 310     os.unlink(filename)
 311   except OSError, err:
 312     if err.errno not in (errno.ENOENT, errno.EISDIR):
 313       raise
 314
 315
 316 def _FingerprintFile(filename):
 317   """Compute the fingerprint of a file.
 318
 319   If the file does not exist, a None will be returned
 320   instead.
 321
 322   Args:
 323     filename - Filename (str)
 324
 325   """
 326   if not (os.path.exists(filename) and os.path.isfile(filename)):
 327     return None
 328
 329   f = open(filename)
 330
 331   fp = sha.sha()
 332   while True:
 333     data = f.read(4096)
 334     if not data:
 335       break
 336
 337     fp.update(data)
 338
 339   return fp.hexdigest()
 340
 341
 342 def FingerprintFiles(files):
 343   """Compute fingerprints for a list of files.
 344
 345   Args:
 346     files - array of filenames.  ( [str, ...] )
 347
 348   Return value:
 349     dictionary of filename: fingerprint for the files that exist
 350
 351   """
 352   ret = {}
 353
 354   for filename in files:
 355     cksum = _FingerprintFile(filename)
 356     if cksum:
 357       ret[filename] = cksum
 358
 359   return ret
 360
 361
 362 def CheckDict(target, template, logname=None):
 363   """Ensure a dictionary has a required set of keys.
 364
 365   For the given dictionaries `target` and `template`, ensure target
 366   has all the keys from template. Missing keys are added with values
 367   from template.
 368
 369   Args:
 370     target   - the dictionary to check
 371     template - template dictionary
 372     logname  - a caller-chosen string to identify the debug log
 373                entry; if None, no logging will be done
 374
 375   Returns value:
 376     None
 377
 378   """
 379   missing = []
 380   for k in template:
 381     if k not in target:
 382       missing.append(k)
 383       target[k] = template[k]
 384
 385   if missing and logname:
 386     logger.Debug('%s missing keys %s' %
 387                  (logname, ', '.join(missing)))
 388
 389
 390 def IsProcessAlive(pid):
 391   """Check if a given pid exists on the system.
 392
 393   Returns: true or false, depending on if the pid exists or not
 394
 395   Remarks: zombie processes treated as not alive
 396
 397   """
 398   try:
 399     f = open("/proc/%d/status" % pid)
 400   except IOError, err:
 401     if err.errno in (errno.ENOENT, errno.ENOTDIR):
 402       return False
 403
 404   alive = True
 405   try:
 406     data = f.readlines()
 407     if len(data) > 1:
 408       state = data[1].split()
 409       if len(state) > 1 and state[1] == "Z":
 410         alive = False
 411   finally:
 412     f.close()
 413
 414   return alive
 415
 416
 417 def MatchNameComponent(key, name_list):
 418   """Try to match a name against a list.
 419
 420   This function will try to match a name like test1 against a list
 421   like ['test1.example.com', 'test2.example.com', ...]. Against this
 422   list, 'test1' as well as 'test1.example' will match, but not
 423   'test1.ex'. A multiple match will be considered as no match at all
 424   (e.g. 'test1' against ['test1.example.com', 'test1.example.org']).
 425
 426   Args:
 427     key: the name to be searched
 428     name_list: the list of strings against which to search the key
 429
 430   Returns:
 431     None if there is no match *or* if there are multiple matches
 432     otherwise the element from the list which matches
 433
 434   """
 435   mo = re.compile("^%s(\..*)?$" % re.escape(key))
 436   names_filtered = [name for name in name_list if mo.match(name) is not None]
 437   if len(names_filtered) != 1:
 438     return None
 439   return names_filtered[0]
 440
 441
 442 class HostInfo:
 443   """Class implementing resolver and hostname functionality
 444
 445   """
 446   def __init__(self, name=None):
 447     """Initialize the host name object.
 448
 449     If the name argument is not passed, it will use this system's
 450     name.
 451
 452     """
 453     if name is None:
 454       name = self.SysName()
 455
 456     self.query = name
 457     self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
 458     self.ip = self.ipaddrs[0]
 459
 460   def ShortName(self):
 461     """Returns the hostname without domain.
 462
 463     """
 464     return self.name.split('.')[0]
 465
 466   @staticmethod
 467   def SysName():
 468     """Return the current system's name.
 469
 470     This is simply a wrapper over socket.gethostname()
 471
 472     """
 473     return socket.gethostname()
 474
 475   @staticmethod
 476   def LookupHostname(hostname):
 477     """Look up hostname
 478
 479     Args:
 480       hostname: hostname to look up
 481
 482     Returns:
 483       a tuple (name, aliases, ipaddrs) as returned by socket.gethostbyname_ex
 484       in case of errors in resolving, we raise a ResolverError
 485
 486     """
 487     try:
 488       result = socket.gethostbyname_ex(hostname)
 489     except socket.gaierror, err:
 490       # hostname not found in DNS
 491       raise errors.ResolverError(hostname, err.args[0], err.args[1])
 492
 493     return result
 494
 495
 496 def ListVolumeGroups():
 497   """List volume groups and their size
 498
 499   Returns:
 500      Dictionary with keys volume name and values the size of the volume
 501
 502   """
 503   command = "vgs --noheadings --units m --nosuffix -o name,size"
 504   result = RunCmd(command)
 505   retval = {}
 506   if result.failed:
 507     return retval
 508
 509   for line in result.stdout.splitlines():
 510     try:
 511       name, size = line.split()
 512       size = int(float(size))
 513     except (IndexError, ValueError), err:
 514       logger.Error("Invalid output from vgs (%s): %s" % (err, line))
 515       continue
 516
 517     retval[name] = size
 518
 519   return retval
 520
 521
 522 def BridgeExists(bridge):
 523   """Check whether the given bridge exists in the system
 524
 525   Returns:
 526      True if it does, false otherwise.
 527
 528   """
 529   return os.path.isdir("/sys/class/net/%s/bridge" % bridge)
 530
 531
 532 def NiceSort(name_list):
 533   """Sort a list of strings based on digit and non-digit groupings.
 534
 535   Given a list of names ['a1', 'a10', 'a11', 'a2'] this function will
 536   sort the list in the logical order ['a1', 'a2', 'a10', 'a11'].
 537
 538   The sort algorithm breaks each name in groups of either only-digits
 539   or no-digits. Only the first eight such groups are considered, and
 540   after that we just use what's left of the string.
 541
 542   Return value
 543     - a copy of the list sorted according to our algorithm
 544
 545   """
 546   _SORTER_BASE = "(\D+|\d+)"
 547   _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
 548                                                   _SORTER_BASE, _SORTER_BASE,
 549                                                   _SORTER_BASE, _SORTER_BASE,
 550                                                   _SORTER_BASE, _SORTER_BASE)
 551   _SORTER_RE = re.compile(_SORTER_FULL)
 552   _SORTER_NODIGIT = re.compile("^\D*$")
 553   def _TryInt(val):
 554     """Attempts to convert a variable to integer."""
 555     if val is None or _SORTER_NODIGIT.match(val):
 556       return val
 557     rval = int(val)
 558     return rval
 559
 560   to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
 561              for name in name_list]
 562   to_sort.sort()
 563   return [tup[1] for tup in to_sort]
 564
 565
 566 def CheckDaemonAlive(pid_file, process_string):
 567   """Check wether the specified daemon is alive.
 568
 569   Args:
 570    - pid_file: file to read the daemon pid from, the file is
 571                expected to contain only a single line containing
 572                only the PID
 573    - process_string: a substring that we expect to find in
 574                      the command line of the daemon process
 575
 576   Returns:
 577    - True if the daemon is judged to be alive (that is:
 578       - the PID file exists, is readable and contains a number
 579       - a process of the specified PID is running
 580       - that process contains the specified string in its
 581         command line
 582       - the process is not in state Z (zombie))
 583    - False otherwise
 584
 585   """
 586   try:
 587     pid_file = file(pid_file, 'r')
 588     try:
 589       pid = int(pid_file.readline())
 590     finally:
 591       pid_file.close()
 592
 593     cmdline_file_path = "/proc/%s/cmdline" % (pid)
 594     cmdline_file = open(cmdline_file_path, 'r')
 595     try:
 596       cmdline = cmdline_file.readline()
 597     finally:
 598       cmdline_file.close()
 599
 600     if not process_string in cmdline:
 601       return False
 602
 603     stat_file_path =  "/proc/%s/stat" % (pid)
 604     stat_file = open(stat_file_path, 'r')
 605     try:
 606       process_state = stat_file.readline().split()[2]
 607     finally:
 608       stat_file.close()
 609
 610     if process_state == 'Z':
 611       return False
 612
 613   except (IndexError, IOError, ValueError):
 614     return False
 615
 616   return True
 617
 618
 619 def TryConvert(fn, val):
 620   """Try to convert a value ignoring errors.
 621
 622   This function tries to apply function `fn` to `val`. If no
 623   ValueError or TypeError exceptions are raised, it will return the
 624   result, else it will return the original value. Any other exceptions
 625   are propagated to the caller.
 626
 627   """
 628   try:
 629     nv = fn(val)
 630   except (ValueError, TypeError), err:
 631     nv = val
 632   return nv
 633
 634
 635 def IsValidIP(ip):
 636   """Verifies the syntax of an IP address.
 637
 638   This function checks if the ip address passes is valid or not based
 639   on syntax (not ip range, class calculations or anything).
 640
 641   """
 642   unit = "(0|[1-9]\d{0,2})"
 643   return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
 644
 645
 646 def IsValidShellParam(word):
 647   """Verifies is the given word is safe from the shell's p.o.v.
 648
 649   This means that we can pass this to a command via the shell and be
 650   sure that it doesn't alter the command line and is passed as such to
 651   the actual command.
 652
 653   Note that we are overly restrictive here, in order to be on the safe
 654   side.
 655
 656   """
 657   return bool(re.match("^[-a-zA-Z0-9._+/:%@]+$", word))
 658
 659
 660 def BuildShellCmd(template, *args):
 661   """Build a safe shell command line from the given arguments.
 662
 663   This function will check all arguments in the args list so that they
 664   are valid shell parameters (i.e. they don't contain shell
 665   metacharaters). If everything is ok, it will return the result of
 666   template % args.
 667
 668   """
 669   for word in args:
 670     if not IsValidShellParam(word):
 671       raise errors.ProgrammerError("Shell argument '%s' contains"
 672                                    " invalid characters" % word)
 673   return template % args
 674
 675
 676 def FormatUnit(value):
 677   """Formats an incoming number of MiB with the appropriate unit.
 678
 679   Value needs to be passed as a numeric type. Return value is always a string.
 680
 681   """
 682   if value < 1024:
 683     return "%dM" % round(value, 0)
 684
 685   elif value < (1024 * 1024):
 686     return "%0.1fG" % round(float(value) / 1024, 1)
 687
 688   else:
 689     return "%0.1fT" % round(float(value) / 1024 / 1024, 1)
 690
 691
 692 def ParseUnit(input_string):
 693   """Tries to extract number and scale from the given string.
 694
 695   Input must be in the format NUMBER+ [DOT NUMBER+] SPACE* [UNIT]. If no unit
 696   is specified, it defaults to MiB. Return value is always an int in MiB.
 697
 698   """
 699   m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', input_string)
 700   if not m:
 701     raise errors.UnitParseError("Invalid format")
 702
 703   value = float(m.groups()[0])
 704
 705   unit = m.groups()[1]
 706   if unit:
 707     lcunit = unit.lower()
 708   else:
 709     lcunit = 'm'
 710
 711   if lcunit in ('m', 'mb', 'mib'):
 712     # Value already in MiB
 713     pass
 714
 715   elif lcunit in ('g', 'gb', 'gib'):
 716     value *= 1024
 717
 718   elif lcunit in ('t', 'tb', 'tib'):
 719     value *= 1024 * 1024
 720
 721   else:
 722     raise errors.UnitParseError("Unknown unit: %s" % unit)
 723
 724   # Make sure we round up
 725   if int(value) < value:
 726     value += 1
 727
 728   # Round up to the next multiple of 4
 729   value = int(value)
 730   if value % 4:
 731     value += 4 - value % 4
 732
 733   return value
 734
 735
 736 def AddAuthorizedKey(file_name, key):
 737   """Adds an SSH public key to an authorized_keys file.
 738
 739   Args:
 740     file_name: Path to authorized_keys file
 741     key: String containing key
 742   """
 743   key_fields = key.split()
 744
 745   f = open(file_name, 'a+')
 746   try:
 747     nl = True
 748     for line in f:
 749       # Ignore whitespace changes
 750       if line.split() == key_fields:
 751         break
 752       nl = line.endswith('\n')
 753     else:
 754       if not nl:
 755         f.write("\n")
 756       f.write(key.rstrip('\r\n'))
 757       f.write("\n")
 758       f.flush()
 759   finally:
 760     f.close()
 761
 762
 763 def RemoveAuthorizedKey(file_name, key):
 764   """Removes an SSH public key from an authorized_keys file.
 765
 766   Args:
 767     file_name: Path to authorized_keys file
 768     key: String containing key
 769   """
 770   key_fields = key.split()
 771
 772   fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
 773   try:
 774     out = os.fdopen(fd, 'w')
 775     try:
 776       f = open(file_name, 'r')
 777       try:
 778         for line in f:
 779           # Ignore whitespace changes while comparing lines
 780           if line.split() != key_fields:
 781             out.write(line)
 782
 783         out.flush()
 784         os.rename(tmpname, file_name)
 785       finally:
 786         f.close()
 787     finally:
 788       out.close()
 789   except:
 790     RemoveFile(tmpname)
 791     raise
 792
 793
 794 def SetEtcHostsEntry(file_name, ip, hostname, aliases):
 795   """Sets the name of an IP address and hostname in /etc/hosts.
 796
 797   """
 798   # Ensure aliases are unique
 799   aliases = UniqueSequence([hostname] + aliases)[1:]
 800
 801   fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
 802   try:
 803     out = os.fdopen(fd, 'w')
 804     try:
 805       f = open(file_name, 'r')
 806       try:
 807         written = False
 808         for line in f:
 809           fields = line.split()
 810           if fields and not fields[0].startswith('#') and ip == fields[0]:
 811             continue
 812           out.write(line)
 813
 814         out.write("%s\t%s" % (ip, hostname))
 815         if aliases:
 816           out.write(" %s" % ' '.join(aliases))
 817         out.write('\n')
 818
 819         out.flush()
 820         os.fsync(out)
 821         os.rename(tmpname, file_name)
 822       finally:
 823         f.close()
 824     finally:
 825       out.close()
 826   except:
 827     RemoveFile(tmpname)
 828     raise
 829
 830
 831 def RemoveEtcHostsEntry(file_name, hostname):
 832   """Removes a hostname from /etc/hosts.
 833
 834   IP addresses without names are removed from the file.
 835   """
 836   fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
 837   try:
 838     out = os.fdopen(fd, 'w')
 839     try:
 840       f = open(file_name, 'r')
 841       try:
 842         for line in f:
 843           fields = line.split()
 844           if len(fields) > 1 and not fields[0].startswith('#'):
 845             names = fields[1:]
 846             if hostname in names:
 847               while hostname in names:
 848                 names.remove(hostname)
 849               if names:
 850                 out.write("%s %s\n" % (fields[0], ' '.join(names)))
 851               continue
 852
 853           out.write(line)
 854
 855         out.flush()
 856         os.fsync(out)
 857         os.rename(tmpname, file_name)
 858       finally:
 859         f.close()
 860     finally:
 861       out.close()
 862   except:
 863     RemoveFile(tmpname)
 864     raise
 865
 866
 867 def CreateBackup(file_name):
 868   """Creates a backup of a file.
 869
 870   Returns: the path to the newly created backup file.
 871
 872   """
 873   if not os.path.isfile(file_name):
 874     raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
 875                                 file_name)
 876
 877   prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
 878   dir_name = os.path.dirname(file_name)
 879
 880   fsrc = open(file_name, 'rb')
 881   try:
 882     (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
 883     fdst = os.fdopen(fd, 'wb')
 884     try:
 885       shutil.copyfileobj(fsrc, fdst)
 886     finally:
 887       fdst.close()
 888   finally:
 889     fsrc.close()
 890
 891   return backup_name
 892
 893
 894 def ShellQuote(value):
 895   """Quotes shell argument according to POSIX.
 896
 897   """
 898   if _re_shell_unquoted.match(value):
 899     return value
 900   else:
 901     return "'%s'" % value.replace("'", "'\\''")
 902
 903
 904 def ShellQuoteArgs(args):
 905   """Quotes all given shell arguments and concatenates using spaces.
 906
 907   """
 908   return ' '.join([ShellQuote(i) for i in args])
 909
 910
 911
 912 def TcpPing(source, target, port, timeout=10, live_port_needed=False):
 913   """Simple ping implementation using TCP connect(2).
 914
 915   Try to do a TCP connect(2) from the specified source IP to the specified
 916   target IP and the specified target port. If live_port_needed is set to true,
 917   requires the remote end to accept the connection. The timeout is specified
 918   in seconds and defaults to 10 seconds
 919
 920   """
 921   sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 922
 923   sucess = False
 924
 925   try:
 926     sock.bind((source, 0))
 927   except socket.error, (errcode, errstring):
 928     if errcode == errno.EADDRNOTAVAIL:
 929       success = False
 930
 931   sock.settimeout(timeout)
 932
 933   try:
 934     sock.connect((target, port))
 935     sock.close()
 936     success = True
 937   except socket.timeout:
 938     success = False
 939   except socket.error, (errcode, errstring):
 940     success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
 941
 942   return success
 943
 944
 945 def ListVisibleFiles(path):
 946   """Returns a list of all visible files in a directory.
 947
 948   """
 949   files = [i for i in os.listdir(path) if not i.startswith(".")]
 950   files.sort()
 951   return files
 952
 953
 954 def GetHomeDir(user, default=None):
 955   """Try to get the homedir of the given user.
 956
 957   The user can be passed either as a string (denoting the name) or as
 958   an integer (denoting the user id). If the user is not found, the
 959   'default' argument is returned, which defaults to None.
 960
 961   """
 962   try:
 963     if isinstance(user, basestring):
 964       result = pwd.getpwnam(user)
 965     elif isinstance(user, (int, long)):
 966       result = pwd.getpwuid(user)
 967     else:
 968       raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
 969                                    type(user))
 970   except KeyError:
 971     return default
 972   return result.pw_dir
 973
 974
 975 def NewUUID():
 976   """Returns a random UUID.
 977
 978   """
 979   f = open("/proc/sys/kernel/random/uuid", "r")
 980   try:
 981     return f.read(128).rstrip("\n")
 982   finally:
 983     f.close()
 984
 985
 986 def WriteFile(file_name, fn=None, data=None,
 987               mode=None, uid=-1, gid=-1,
 988               atime=None, mtime=None):
 989   """(Over)write a file atomically.
 990
 991   The file_name and either fn (a function taking one argument, the
 992   file descriptor, and which should write the data to it) or data (the
 993   contents of the file) must be passed. The other arguments are
 994   optional and allow setting the file mode, owner and group, and the
 995   mtime/atime of the file.
 996
 997   If the function doesn't raise an exception, it has succeeded and the
 998   target file has the new contents. If the file has raised an
 999   exception, an existing target file should be unmodified and the
1000   temporary file should be removed.
1001
1002   """
1003   if not os.path.isabs(file_name):
1004     raise errors.ProgrammerError("Path passed to WriteFile is not"
1005                                  " absolute: '%s'" % file_name)
1006
1007   if [fn, data].count(None) != 1:
1008     raise errors.ProgrammerError("fn or data required")
1009
1010   if [atime, mtime].count(None) == 1:
1011     raise errors.ProgrammerError("Both atime and mtime must be either"
1012                                  " set or None")
1013
1014
1015   dir_name, base_name = os.path.split(file_name)
1016   fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
1017   # here we need to make sure we remove the temp file, if any error
1018   # leaves it in place
1019   try:
1020     if uid != -1 or gid != -1:
1021       os.chown(new_name, uid, gid)
1022     if mode:
1023       os.chmod(new_name, mode)
1024     if data is not None:
1025       os.write(fd, data)
1026     else:
1027       fn(fd)
1028     os.fsync(fd)
1029     if atime is not None and mtime is not None:
1030       os.utime(new_name, (atime, mtime))
1031     os.rename(new_name, file_name)
1032   finally:
1033     os.close(fd)
1034     RemoveFile(new_name)
1035
1036
1037 def all(seq, pred=bool):
1038   "Returns True if pred(x) is True for every element in the iterable"
1039   for elem in itertools.ifilterfalse(pred, seq):
1040     return False
1041   return True
1042
1043
1044 def any(seq, pred=bool):
1045   "Returns True if pred(x) is True for at least one element in the iterable"
1046   for elem in itertools.ifilter(pred, seq):
1047     return True
1048   return False
1049
1050
1051 def UniqueSequence(seq):
1052   """Returns a list with unique elements.
1053
1054   Element order is preserved.
1055   """
1056   seen = set()
1057   return [i for i in seq if i not in seen and not seen.add(i)]
1058
1059
1060 def IsValidMac(mac):
1061   """Predicate to check if a MAC address is valid.
1062
1063   Checks wether the supplied MAC address is formally correct, only
1064   accepts colon separated format.
1065   """
1066   mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
1067   return mac_check.match(mac) is not None