Add support for command aliases
[ganeti-local] / lib / utils.py
index b43570f..f854f5a 100644 (file)
@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#
 #
 
 # Copyright (C) 2006, 2007 Google Inc.
@@ -20,6 +20,7 @@
 
 
 """Ganeti small utilities
+
 """
 
 
@@ -32,14 +33,24 @@ import re
 import socket
 import tempfile
 import shutil
-from errno import ENOENT, ENOTDIR, EISDIR, EEXIST
+import errno
+import pwd
+import itertools
+import select
+import fcntl
+
+from cStringIO import StringIO
 
 from ganeti import logger
 from ganeti import errors
+from ganeti import constants
+
 
 _locksheld = []
 _re_shell_unquoted = re.compile('^[-.,=:/_+@A-Za-z0-9]+$')
 
+debug = False
+
 class RunResult(object):
   """Simple class for holding the result of running external programs.
 
@@ -74,6 +85,10 @@ class RunResult(object):
     else:
       self.fail_reason = "unable to determine termination reason"
 
+    if debug and self.failed:
+      logger.Debug("Command '%s' failed (%s); output: %s" %
+                   (self.cmd, self.fail_reason, self.output))
+
   def _GetOutput(self):
     """Returns the combined stdout and stderr for easier usage.
 
@@ -85,7 +100,7 @@ class RunResult(object):
 
 def _GetLockFile(subsystem):
   """Compute the file name for a given lock name."""
-  return "/var/lock/ganeti_lock_%s" % subsystem
+  return "%s/ganeti_lock_%s" % (constants.LOCK_DIR, subsystem)
 
 
 def Lock(name, max_retries=None, debug=False):
@@ -115,7 +130,7 @@ def Lock(name, max_retries=None, debug=False):
       fd = os.open(lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_SYNC)
       break
     except OSError, creat_err:
-      if creat_err.errno != EEXIST:
+      if creat_err.errno != errno.EEXIST:
         raise errors.LockError("Can't create the lock file. Error '%s'." %
                                str(creat_err))
 
@@ -211,6 +226,7 @@ def RunCmd(cmd):
     shell = True
   env = os.environ.copy()
   env["LC_ALL"] = "C"
+  poller = select.poll()
   child = subprocess.Popen(cmd, shell=shell,
                            stderr=subprocess.PIPE,
                            stdout=subprocess.PIPE,
@@ -218,8 +234,35 @@ def RunCmd(cmd):
                            close_fds=True, env=env)
 
   child.stdin.close()
-  out = child.stdout.read()
-  err = child.stderr.read()
+  poller.register(child.stdout, select.POLLIN)
+  poller.register(child.stderr, select.POLLIN)
+  out = StringIO()
+  err = StringIO()
+  fdmap = {
+    child.stdout.fileno(): (out, child.stdout),
+    child.stderr.fileno(): (err, child.stderr),
+    }
+  for fd in fdmap:
+    status = fcntl.fcntl(fd, fcntl.F_GETFL)
+    fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
+
+  while fdmap:
+    for fd, event in poller.poll():
+      if event & select.POLLIN or event & select.POLLPRI:
+        data = fdmap[fd][1].read()
+        # no data from read signifies EOF (the same as POLLHUP)
+        if not data:
+          poller.unregister(fd)
+          del fdmap[fd]
+          continue
+        fdmap[fd][0].write(data)
+      if (event & select.POLLNVAL or event & select.POLLHUP or
+          event & select.POLLERR):
+        poller.unregister(fd)
+        del fdmap[fd]
+
+  out = out.getvalue()
+  err = err.getvalue()
 
   status = child.wait()
   if status >= 0:
@@ -266,7 +309,7 @@ def RemoveFile(filename):
   try:
     os.unlink(filename)
   except OSError, err:
-    if err.errno not in (ENOENT, EISDIR):
+    if err.errno not in (errno.ENOENT, errno.EISDIR):
       raise
 
 
@@ -355,7 +398,7 @@ def IsProcessAlive(pid):
   try:
     f = open("/proc/%d/status" % pid)
   except IOError, err:
-    if err.errno in (ENOENT, ENOTDIR):
+    if err.errno in (errno.ENOENT, errno.ENOTDIR):
       return False
 
   alive = True
@@ -396,33 +439,58 @@ def MatchNameComponent(key, name_list):
   return names_filtered[0]
 
 
-def LookupHostname(hostname):
-  """Look up hostname
+class HostInfo:
+  """Class implementing resolver and hostname functionality
 
-  Args:
-    hostname: hostname to look up, can be also be a non FQDN
+  """
+  def __init__(self, name=None):
+    """Initialize the host name object.
 
-  Returns:
-    Dictionary with keys:
-    - ip: IP addr
-    - hostname_full: hostname fully qualified
-    - hostname: hostname fully qualified (historic artifact)
+    If the name argument is not passed, it will use this system's
+    name.
 
-  """
-  try:
-    (fqdn, dummy, ipaddrs) = socket.gethostbyname_ex(hostname)
-    ipaddr = ipaddrs[0]
-  except socket.gaierror:
-    # hostname not found in DNS
-    return None
+    """
+    if name is None:
+      name = self.SysName()
 
-  returnhostname = {
-    "ip": ipaddr,
-    "hostname_full": fqdn,
-    "hostname": fqdn,
-    }
+    self.query = name
+    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
+    self.ip = self.ipaddrs[0]
+
+  def ShortName(self):
+    """Returns the hostname without domain.
+
+    """
+    return self.name.split('.')[0]
 
-  return returnhostname
+  @staticmethod
+  def SysName():
+    """Return the current system's name.
+
+    This is simply a wrapper over socket.gethostname()
+
+    """
+    return socket.gethostname()
+
+  @staticmethod
+  def LookupHostname(hostname):
+    """Look up hostname
+
+    Args:
+      hostname: hostname to look up
+
+    Returns:
+      a tuple (name, aliases, ipaddrs) as returned by socket.gethostbyname_ex
+      in case of errors in resolving, we raise a ResolverError
+
+    """
+    try:
+      result = socket.gethostbyname_ex(hostname)
+    except socket.gaierror, err:
+      # hostname not found in DNS
+      raise errors.ResolverError(hostname, err.args[0], err.args[1])
+
+    return result
 
 
 def ListVolumeGroups():
@@ -702,21 +770,98 @@ def RemoveAuthorizedKey(file_name, key):
   key_fields = key.split()
 
   fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
-  out = os.fdopen(fd, 'w')
   try:
-    f = open(file_name, 'r')
+    out = os.fdopen(fd, 'w')
+    try:
+      f = open(file_name, 'r')
+      try:
+        for line in f:
+          # Ignore whitespace changes while comparing lines
+          if line.split() != key_fields:
+            out.write(line)
+
+        out.flush()
+        os.rename(tmpname, file_name)
+      finally:
+        f.close()
+    finally:
+      out.close()
+  except:
+    RemoveFile(tmpname)
+    raise
+
+
+def SetEtcHostsEntry(file_name, ip, hostname, aliases):
+  """Sets the name of an IP address and hostname in /etc/hosts.
+
+  """
+  # Ensure aliases are unique
+  aliases = UniqueSequence([hostname] + aliases)[1:]
+
+  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
+  try:
+    out = os.fdopen(fd, 'w')
+    try:
+      f = open(file_name, 'r')
+      try:
+        written = False
+        for line in f:
+          fields = line.split()
+          if fields and not fields[0].startswith('#') and ip == fields[0]:
+            continue
+          out.write(line)
+
+        out.write("%s\t%s" % (ip, hostname))
+        if aliases:
+          out.write(" %s" % ' '.join(aliases))
+        out.write('\n')
+
+        out.flush()
+        os.fsync(out)
+        os.rename(tmpname, file_name)
+      finally:
+        f.close()
+    finally:
+      out.close()
+  except:
+    RemoveFile(tmpname)
+    raise
+
+
+def RemoveEtcHostsEntry(file_name, hostname):
+  """Removes a hostname from /etc/hosts.
+
+  IP addresses without names are removed from the file.
+  """
+  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
+  try:
+    out = os.fdopen(fd, 'w')
     try:
-      for line in f:
-        # Ignore whitespace changes while comparing lines
-        if line.split() != key_fields:
+      f = open(file_name, 'r')
+      try:
+        for line in f:
+          fields = line.split()
+          if len(fields) > 1 and not fields[0].startswith('#'):
+            names = fields[1:]
+            if hostname in names:
+              while hostname in names:
+                names.remove(hostname)
+              if names:
+                out.write("%s %s\n" % (fields[0], ' '.join(names)))
+              continue
+
           out.write(line)
 
-      out.flush()
-      os.rename(tmpname, file_name)
+        out.flush()
+        os.fsync(out)
+        os.rename(tmpname, file_name)
+      finally:
+        f.close()
     finally:
-      f.close()
-  finally:
-    out.close()
+      out.close()
+  except:
+    RemoveFile(tmpname)
+    raise
 
 
 def CreateBackup(file_name):
@@ -729,10 +874,20 @@ def CreateBackup(file_name):
     raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
                                 file_name)
 
-  # Warning: the following code contains a race condition when we create more
-  # than one backup of the same file in a second.
-  backup_name = file_name + '.backup-%d' % int(time.time())
-  shutil.copyfile(file_name, backup_name)
+  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
+  dir_name = os.path.dirname(file_name)
+
+  fsrc = open(file_name, 'rb')
+  try:
+    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
+    fdst = os.fdopen(fd, 'wb')
+    try:
+      shutil.copyfileobj(fsrc, fdst)
+    finally:
+      fdst.close()
+  finally:
+    fsrc.close()
+
   return backup_name
 
 
@@ -753,36 +908,160 @@ def ShellQuoteArgs(args):
   return ' '.join([ShellQuote(i) for i in args])
 
 
-def _ParseIpOutput(output):
-  """Parsing code for GetLocalIPAddresses().
 
-  This function is split out, so we can unit test it.
+def TcpPing(source, target, port, timeout=10, live_port_needed=False):
+  """Simple ping implementation using TCP connect(2).
+
+  Try to do a TCP connect(2) from the specified source IP to the specified
+  target IP and the specified target port. If live_port_needed is set to true,
+  requires the remote end to accept the connection. The timeout is specified
+  in seconds and defaults to 10 seconds
 
   """
-  re_ip = re.compile('^(\d+\.\d+\.\d+\.\d+)(?:/\d+)$')
+  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 
-  ips = []
-  for line in output.splitlines(False):
-    fields = line.split()
-    if len(line) < 4:
-      continue
-    m = re_ip.match(fields[3])
-    if m:
-      ips.append(m.group(1))
+  sucess = False
+
+  try:
+    sock.bind((source, 0))
+  except socket.error, (errcode, errstring):
+    if errcode == errno.EADDRNOTAVAIL:
+      success = False
 
-  return ips
+  sock.settimeout(timeout)
 
+  try:
+    sock.connect((target, port))
+    sock.close()
+    success = True
+  except socket.timeout:
+    success = False
+  except socket.error, (errcode, errstring):
+    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
 
-def GetLocalIPAddresses():
-  """Gets a list of all local IP addresses.
+  return success
 
-  Should this break one day, a small Python module written in C could
-  use the API call getifaddrs().
+
+def ListVisibleFiles(path):
+  """Returns a list of all visible files in a directory.
 
   """
-  result = RunCmd(["ip", "-family", "inet", "-oneline", "addr", "show"])
-  if result.failed:
-    raise errors.OpExecError("Command '%s' failed, error: %s,"
-      " output: %s" % (result.cmd, result.fail_reason, result.output))
+  files = [i for i in os.listdir(path) if not i.startswith(".")]
+  files.sort()
+  return files
+
+
+def GetHomeDir(user, default=None):
+  """Try to get the homedir of the given user.
+
+  The user can be passed either as a string (denoting the name) or as
+  an integer (denoting the user id). If the user is not found, the
+  'default' argument is returned, which defaults to None.
+
+  """
+  try:
+    if isinstance(user, basestring):
+      result = pwd.getpwnam(user)
+    elif isinstance(user, (int, long)):
+      result = pwd.getpwuid(user)
+    else:
+      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
+                                   type(user))
+  except KeyError:
+    return default
+  return result.pw_dir
+
+
+def NewUUID():
+  """Returns a random UUID.
+
+  """
+  f = open("/proc/sys/kernel/random/uuid", "r")
+  try:
+    return f.read(128).rstrip("\n")
+  finally:
+    f.close()
 
-  return _ParseIpOutput(result.output)
+
+def WriteFile(file_name, fn=None, data=None,
+              mode=None, uid=-1, gid=-1,
+              atime=None, mtime=None):
+  """(Over)write a file atomically.
+
+  The file_name and either fn (a function taking one argument, the
+  file descriptor, and which should write the data to it) or data (the
+  contents of the file) must be passed. The other arguments are
+  optional and allow setting the file mode, owner and group, and the
+  mtime/atime of the file.
+
+  If the function doesn't raise an exception, it has succeeded and the
+  target file has the new contents. If the file has raised an
+  exception, an existing target file should be unmodified and the
+  temporary file should be removed.
+
+  """
+  if not os.path.isabs(file_name):
+    raise errors.ProgrammerError("Path passed to WriteFile is not"
+                                 " absolute: '%s'" % file_name)
+
+  if [fn, data].count(None) != 1:
+    raise errors.ProgrammerError("fn or data required")
+
+  if [atime, mtime].count(None) == 1:
+    raise errors.ProgrammerError("Both atime and mtime must be either"
+                                 " set or None")
+
+
+  dir_name, base_name = os.path.split(file_name)
+  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
+  # here we need to make sure we remove the temp file, if any error
+  # leaves it in place
+  try:
+    if uid != -1 or gid != -1:
+      os.chown(new_name, uid, gid)
+    if mode:
+      os.chmod(new_name, mode)
+    if data is not None:
+      os.write(fd, data)
+    else:
+      fn(fd)
+    os.fsync(fd)
+    if atime is not None and mtime is not None:
+      os.utime(new_name, (atime, mtime))
+    os.rename(new_name, file_name)
+  finally:
+    os.close(fd)
+    RemoveFile(new_name)
+
+
+def all(seq, pred=bool):
+  "Returns True if pred(x) is True for every element in the iterable"
+  for elem in itertools.ifilterfalse(pred, seq):
+    return False
+  return True
+
+
+def any(seq, pred=bool):
+  "Returns True if pred(x) is True for at least one element in the iterable"
+  for elem in itertools.ifilter(pred, seq):
+    return True
+  return False
+
+
+def UniqueSequence(seq):
+  """Returns a list with unique elements.
+
+  Element order is preserved.
+  """
+  seen = set()
+  return [i for i in seq if i not in seen and not seen.add(i)]
+
+
+def IsValidMac(mac):
+  """Predicate to check if a MAC address is valid.
+
+  Checks wether the supplied MAC address is formally correct, only
+  accepts colon separated format.
+  """
+  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
+  return mac_check.match(mac) is not None