utils.py 55.2 KB
Newer Older
Iustin Pop's avatar
Iustin Pop committed
1
#
Iustin Pop's avatar
Iustin Pop committed
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21
#

# Copyright (C) 2006, 2007 Google Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.


22 23 24 25
"""Ganeti utility module.

This module holds functions that can be used in both daemons (all) and
the command line scripts.
26

Iustin Pop's avatar
Iustin Pop committed
27 28 29 30 31 32
"""


import sys
import os
import time
33
import subprocess
Iustin Pop's avatar
Iustin Pop committed
34 35 36 37
import re
import socket
import tempfile
import shutil
38
import errno
39
import pwd
Guido Trotter's avatar
Guido Trotter committed
40
import itertools
41 42
import select
import fcntl
43
import resource
44
import logging
Michael Hanselmann's avatar
Michael Hanselmann committed
45
import signal
46 47

from cStringIO import StringIO
Iustin Pop's avatar
Iustin Pop committed
48

49 50 51 52 53 54
try:
  from hashlib import sha1
except ImportError:
  import sha
  sha1 = sha.new

Iustin Pop's avatar
Iustin Pop committed
55
from ganeti import errors
Iustin Pop's avatar
Iustin Pop committed
56
from ganeti import constants
Iustin Pop's avatar
Iustin Pop committed
57

58

Iustin Pop's avatar
Iustin Pop committed
59 60 61
_locksheld = []
_re_shell_unquoted = re.compile('^[-.,=:/_+@A-Za-z0-9]+$')

62
debug = False
63
debug_locks = False
64 65

#: when set to True, L{RunCmd} is disabled
66
no_fork = False
67

68

Iustin Pop's avatar
Iustin Pop committed
69
class RunResult(object):
70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85
  """Holds the result of running external programs.

  @type exit_code: int
  @ivar exit_code: the exit code of the program, or None (if the program
      didn't exit())
  @type signal: int or None
  @ivar signal: the signal that caused the program to finish, or None
      (if the program wasn't terminated by a signal)
  @type stdout: str
  @ivar stdout: the standard output of the program
  @type stderr: str
  @ivar stderr: the standard error of the program
  @type failed: boolean
  @ivar failed: True in case the program was
      terminated by a signal or exited with a non-zero exit code
  @ivar fail_reason: a string detailing the termination reason
Iustin Pop's avatar
Iustin Pop committed
86 87 88 89 90 91

  """
  __slots__ = ["exit_code", "signal", "stdout", "stderr",
               "failed", "fail_reason", "cmd"]


Iustin Pop's avatar
Iustin Pop committed
92
  def __init__(self, exit_code, signal_, stdout, stderr, cmd):
Iustin Pop's avatar
Iustin Pop committed
93 94
    self.cmd = cmd
    self.exit_code = exit_code
Iustin Pop's avatar
Iustin Pop committed
95
    self.signal = signal_
Iustin Pop's avatar
Iustin Pop committed
96 97
    self.stdout = stdout
    self.stderr = stderr
Iustin Pop's avatar
Iustin Pop committed
98
    self.failed = (signal_ is not None or exit_code != 0)
Iustin Pop's avatar
Iustin Pop committed
99 100 101 102 103 104 105 106

    if self.signal is not None:
      self.fail_reason = "terminated by signal %s" % self.signal
    elif self.exit_code is not None:
      self.fail_reason = "exited with exit code %s" % self.exit_code
    else:
      self.fail_reason = "unable to determine termination reason"

107 108 109
    if self.failed:
      logging.debug("Command '%s' failed (%s); output: %s",
                    self.cmd, self.fail_reason, self.output)
110

Iustin Pop's avatar
Iustin Pop committed
111 112 113 114 115 116 117 118 119
  def _GetOutput(self):
    """Returns the combined stdout and stderr for easier usage.

    """
    return self.stdout + self.stderr

  output = property(_GetOutput, None, None, "Return full output")


120
def RunCmd(cmd, env=None, output=None, cwd='/'):
Iustin Pop's avatar
Iustin Pop committed
121 122 123 124 125
  """Execute a (shell) command.

  The command should not read from its standard input, as it will be
  closed.

126
  @type  cmd: string or list
127
  @param cmd: Command to run
128
  @type env: dict
129
  @param env: Additional environment
130
  @type output: str
131
  @param output: if desired, the output of the command can be
132 133
      saved in a file instead of the RunResult instance; this
      parameter denotes the file name (if not None)
134 135 136
  @type cwd: string
  @param cwd: if specified, will be used as the working
      directory for the command; the default will be /
137
  @rtype: L{RunResult}
138 139
  @return: RunResult instance
  @raise erors.ProgrammerError: if we call this when forks are disabled
Iustin Pop's avatar
Iustin Pop committed
140 141

  """
142 143 144
  if no_fork:
    raise errors.ProgrammerError("utils.RunCmd() called with fork() disabled")

Iustin Pop's avatar
Iustin Pop committed
145 146
  if isinstance(cmd, list):
    cmd = [str(val) for val in cmd]
147 148 149 150 151
    strcmd = " ".join(cmd)
    shell = False
  else:
    strcmd = cmd
    shell = True
152
  logging.debug("RunCmd '%s'", strcmd)
153 154 155 156 157 158

  cmd_env = os.environ.copy()
  cmd_env["LC_ALL"] = "C"
  if env is not None:
    cmd_env.update(env)

159
  if output is None:
160
    out, err, status = _RunCmdPipe(cmd, cmd_env, shell, cwd)
161
  else:
162
    status = _RunCmdFile(cmd, cmd_env, shell, output, cwd)
163 164 165 166 167 168 169 170 171 172 173
    out = err = ""

  if status >= 0:
    exitcode = status
    signal_ = None
  else:
    exitcode = None
    signal_ = -status

  return RunResult(exitcode, signal_, out, err, strcmd)

174

175
def _RunCmdPipe(cmd, env, via_shell, cwd):
176 177 178 179 180 181 182 183
  """Run a command and return its output.

  @type  cmd: string or list
  @param cmd: Command to run
  @type env: dict
  @param env: The environment to use
  @type via_shell: bool
  @param via_shell: if we should run via the shell
184 185
  @type cwd: string
  @param cwd: the working directory for the program
186 187 188 189
  @rtype: tuple
  @return: (out, err, status)

  """
190
  poller = select.poll()
191
  child = subprocess.Popen(cmd, shell=via_shell,
192 193 194
                           stderr=subprocess.PIPE,
                           stdout=subprocess.PIPE,
                           stdin=subprocess.PIPE,
195 196
                           close_fds=True, env=env,
                           cwd=cwd)
197 198

  child.stdin.close()
199 200 201 202 203 204 205 206 207 208 209 210 211
  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:
212 213 214 215 216 217 218 219 220 221 222 223
    try:
      pollresult = poller.poll()
    except EnvironmentError, eerr:
      if eerr.errno == errno.EINTR:
        continue
      raise
    except select.error, serr:
      if serr[0] == errno.EINTR:
        continue
      raise

    for fd, event in pollresult:
224 225 226 227 228 229 230 231 232 233 234 235 236 237 238
      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()
Iustin Pop's avatar
Iustin Pop committed
239 240

  status = child.wait()
241
  return out, err, status
Iustin Pop's avatar
Iustin Pop committed
242

243

244
def _RunCmdFile(cmd, env, via_shell, output, cwd):
245 246 247 248 249 250 251 252 253 254
  """Run a command and save its output to a file.

  @type  cmd: string or list
  @param cmd: Command to run
  @type env: dict
  @param env: The environment to use
  @type via_shell: bool
  @param via_shell: if we should run via the shell
  @type output: str
  @param output: the filename in which to save the output
255 256
  @type cwd: string
  @param cwd: the working directory for the program
257 258 259 260 261 262 263 264 265 266
  @rtype: int
  @return: the exit status

  """
  fh = open(output, "a")
  try:
    child = subprocess.Popen(cmd, shell=via_shell,
                             stderr=subprocess.STDOUT,
                             stdout=fh,
                             stdin=subprocess.PIPE,
267 268
                             close_fds=True, env=env,
                             cwd=cwd)
269 270 271 272 273 274

    child.stdin.close()
    status = child.wait()
  finally:
    fh.close()
  return status
Iustin Pop's avatar
Iustin Pop committed
275 276 277 278 279 280 281 282


def RemoveFile(filename):
  """Remove a file ignoring some errors.

  Remove a file, ignoring non-existing ones or directories. Other
  errors are passed.

283 284 285
  @type filename: str
  @param filename: the file to be removed

Iustin Pop's avatar
Iustin Pop committed
286 287 288 289
  """
  try:
    os.unlink(filename)
  except OSError, err:
290
    if err.errno not in (errno.ENOENT, errno.EISDIR):
Iustin Pop's avatar
Iustin Pop committed
291 292 293
      raise


294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319
def RenameFile(old, new, mkdir=False, mkdir_mode=0750):
  """Renames a file.

  @type old: string
  @param old: Original path
  @type new: string
  @param new: New path
  @type mkdir: bool
  @param mkdir: Whether to create target directory if it doesn't exist
  @type mkdir_mode: int
  @param mkdir_mode: Mode for newly created directories

  """
  try:
    return os.rename(old, new)
  except OSError, err:
    # In at least one use case of this function, the job queue, directory
    # creation is very rare. Checking for the directory before renaming is not
    # as efficient.
    if mkdir and err.errno == errno.ENOENT:
      # Create directory and try again
      os.makedirs(os.path.dirname(new), mkdir_mode)
      return os.rename(old, new)
    raise


Iustin Pop's avatar
Iustin Pop committed
320 321 322 323 324 325
def _FingerprintFile(filename):
  """Compute the fingerprint of a file.

  If the file does not exist, a None will be returned
  instead.

326 327 328 329 330
  @type filename: str
  @param filename: the filename to checksum
  @rtype: str
  @return: the hex digest of the sha checksum of the contents
      of the file
Iustin Pop's avatar
Iustin Pop committed
331 332 333 334 335 336 337

  """
  if not (os.path.exists(filename) and os.path.isfile(filename)):
    return None

  f = open(filename)

338
  fp = sha1()
Iustin Pop's avatar
Iustin Pop committed
339 340 341 342 343 344 345 346 347 348 349 350 351
  while True:
    data = f.read(4096)
    if not data:
      break

    fp.update(data)

  return fp.hexdigest()


def FingerprintFiles(files):
  """Compute fingerprints for a list of files.

352 353 354 355 356
  @type files: list
  @param files: the list of filename to fingerprint
  @rtype: dict
  @return: a dictionary filename: fingerprint, holding only
      existing files
Iustin Pop's avatar
Iustin Pop committed
357 358 359 360 361 362 363 364 365 366 367 368 369 370 371

  """
  ret = {}

  for filename in files:
    cksum = _FingerprintFile(filename)
    if cksum:
      ret[filename] = cksum

  return ret


def CheckDict(target, template, logname=None):
  """Ensure a dictionary has a required set of keys.

372 373 374
  For the given dictionaries I{target} and I{template}, ensure
  I{target} has all the keys from I{template}. Missing keys are added
  with values from template.
Iustin Pop's avatar
Iustin Pop committed
375

376 377 378 379 380 381 382
  @type target: dict
  @param target: the dictionary to update
  @type template: dict
  @param template: the dictionary holding the default values
  @type logname: str or None
  @param logname: if not None, causes the missing keys to be
      logged with this name
Iustin Pop's avatar
Iustin Pop committed
383 384 385 386 387 388 389 390 391

  """
  missing = []
  for k in template:
    if k not in target:
      missing.append(k)
      target[k] = template[k]

  if missing and logname:
392
    logging.warning('%s missing keys %s', logname, ', '.join(missing))
Iustin Pop's avatar
Iustin Pop committed
393 394


395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457
def ForceDictType(target, key_types, allowed_values=None):
  """Force the values of a dict to have certain types.

  @type target: dict
  @param target: the dict to update
  @type key_types: dict
  @param key_types: dict mapping target dict keys to types
                    in constants.ENFORCEABLE_TYPES
  @type allowed_values: list
  @keyword allowed_values: list of specially allowed values

  """
  if allowed_values is None:
    allowed_values = []

  for key in target:
    if key not in key_types:
      msg = "Unknown key '%s'" % key
      raise errors.TypeEnforcementError(msg)

    if target[key] in allowed_values:
      continue

    type = key_types[key]
    if type not in constants.ENFORCEABLE_TYPES:
      msg = "'%s' has non-enforceable type %s" % (key, type)
      raise errors.ProgrammerError(msg)

    if type == constants.VTYPE_STRING:
      if not isinstance(target[key], basestring):
        if isinstance(target[key], bool) and not target[key]:
          target[key] = ''
        else:
          msg = "'%s' (value %s) is not a valid string" % (key, target[key])
          raise errors.TypeEnforcementError(msg)
    elif type == constants.VTYPE_BOOL:
      if isinstance(target[key], basestring) and target[key]:
        if target[key].lower() == constants.VALUE_FALSE:
          target[key] = False
        elif target[key].lower() == constants.VALUE_TRUE:
          target[key] = True
        else:
          msg = "'%s' (value %s) is not a valid boolean" % (key, target[key])
          raise errors.TypeEnforcementError(msg)
      elif target[key]:
        target[key] = True
      else:
        target[key] = False
    elif type == constants.VTYPE_SIZE:
      try:
        target[key] = ParseUnit(target[key])
      except errors.UnitParseError, err:
        msg = "'%s' (value %s) is not a valid size. error: %s" % \
              (key, target[key], err)
        raise errors.TypeEnforcementError(msg)
    elif type == constants.VTYPE_INT:
      try:
        target[key] = int(target[key])
      except (ValueError, TypeError):
        msg = "'%s' (value %s) is not a valid integer" % (key, target[key])
        raise errors.TypeEnforcementError(msg)


Iustin Pop's avatar
Iustin Pop committed
458 459 460
def IsProcessAlive(pid):
  """Check if a given pid exists on the system.

461 462
  @note: zombie status is not handled, so zombie processes
      will be returned as alive
463 464 465 466
  @type pid: int
  @param pid: the process ID to check
  @rtype: boolean
  @return: True if the process exists
Iustin Pop's avatar
Iustin Pop committed
467 468

  """
469 470 471
  if pid <= 0:
    return False

Iustin Pop's avatar
Iustin Pop committed
472
  try:
473 474 475
    os.stat("/proc/%d/status" % pid)
    return True
  except EnvironmentError, err:
476
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
Iustin Pop's avatar
Iustin Pop committed
477
      return False
478
    raise
Iustin Pop's avatar
Iustin Pop committed
479 480


481
def ReadPidFile(pidfile):
482
  """Read a pid from a file.
483

484 485 486
  @type  pidfile: string
  @param pidfile: path to the file containing the pid
  @rtype: int
487
  @return: The process id, if the file exists and contains a valid PID,
488
           otherwise 0
489 490 491 492

  """
  try:
    pf = open(pidfile, 'r')
493 494 495 496
  except EnvironmentError, err:
    if err.errno != errno.ENOENT:
      logging.exception("Can't read pid file?!")
    return 0
497 498 499

  try:
    pid = int(pf.read())
500
  except ValueError, err:
501
    logging.info("Can't parse pid file contents", exc_info=True)
502
    return 0
503

504
  return pid
505 506


Iustin Pop's avatar
Iustin Pop committed
507 508 509 510
def MatchNameComponent(key, name_list):
  """Try to match a name against a list.

  This function will try to match a name like test1 against a list
511 512 513 514 515
  like C{['test1.example.com', 'test2.example.com', ...]}. Against
  this list, I{'test1'} as well as I{'test1.example'} will match, but
  not I{'test1.ex'}. A multiple match will be considered as no match
  at all (e.g. I{'test1'} against C{['test1.example.com',
  'test1.example.org']}).
Iustin Pop's avatar
Iustin Pop committed
516

517 518 519 520
  @type key: str
  @param key: the name to be searched
  @type name_list: list
  @param name_list: the list of strings against which to search the key
Iustin Pop's avatar
Iustin Pop committed
521

522 523 524
  @rtype: None or str
  @return: None if there is no match I{or} if there are multiple matches,
      otherwise the element from the list which matches
Iustin Pop's avatar
Iustin Pop committed
525 526 527 528 529 530 531 532 533

  """
  mo = re.compile("^%s(\..*)?$" % re.escape(key))
  names_filtered = [name for name in name_list if mo.match(name) is not None]
  if len(names_filtered) != 1:
    return None
  return names_filtered[0]


534
class HostInfo:
535
  """Class implementing resolver and hostname functionality
536 537

  """
538
  def __init__(self, name=None):
539 540
    """Initialize the host name object.

541 542
    If the name argument is not passed, it will use this system's
    name.
543 544

    """
545 546 547 548 549
    if name is None:
      name = self.SysName()

    self.query = name
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
550 551
    self.ip = self.ipaddrs[0]

552 553 554 555 556 557
  def ShortName(self):
    """Returns the hostname without domain.

    """
    return self.name.split('.')[0]

558 559 560
  @staticmethod
  def SysName():
    """Return the current system's name.
561

562
    This is simply a wrapper over C{socket.gethostname()}.
Iustin Pop's avatar
Iustin Pop committed
563

564 565
    """
    return socket.gethostname()
Iustin Pop's avatar
Iustin Pop committed
566

567 568 569
  @staticmethod
  def LookupHostname(hostname):
    """Look up hostname
Iustin Pop's avatar
Iustin Pop committed
570

571 572
    @type hostname: str
    @param hostname: hostname to look up
573

574 575 576 577
    @rtype: tuple
    @return: a tuple (name, aliases, ipaddrs) as returned by
        C{socket.gethostbyname_ex}
    @raise errors.ResolverError: in case of errors in resolving
578 579 580 581 582 583 584

    """
    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])
Iustin Pop's avatar
Iustin Pop committed
585

586
    return result
Iustin Pop's avatar
Iustin Pop committed
587 588 589 590 591


def ListVolumeGroups():
  """List volume groups and their size

592 593 594 595
  @rtype: dict
  @return:
       Dictionary with keys volume name and values
       the size of the volume
Iustin Pop's avatar
Iustin Pop committed
596 597 598 599 600 601 602 603 604 605 606 607 608

  """
  command = "vgs --noheadings --units m --nosuffix -o name,size"
  result = RunCmd(command)
  retval = {}
  if result.failed:
    return retval

  for line in result.stdout.splitlines():
    try:
      name, size = line.split()
      size = int(float(size))
    except (IndexError, ValueError), err:
609
      logging.error("Invalid output from vgs (%s): %s", err, line)
Iustin Pop's avatar
Iustin Pop committed
610 611 612 613 614 615 616 617 618 619
      continue

    retval[name] = size

  return retval


def BridgeExists(bridge):
  """Check whether the given bridge exists in the system

620 621 622 623
  @type bridge: str
  @param bridge: the bridge name to check
  @rtype: boolean
  @return: True if it does
Iustin Pop's avatar
Iustin Pop committed
624 625 626 627 628 629 630 631

  """
  return os.path.isdir("/sys/class/net/%s/bridge" % bridge)


def NiceSort(name_list):
  """Sort a list of strings based on digit and non-digit groupings.

632 633 634
  Given a list of names C{['a1', 'a10', 'a11', 'a2']} this function
  will sort the list in the logical order C{['a1', 'a2', 'a10',
  'a11']}.
Iustin Pop's avatar
Iustin Pop committed
635 636 637 638 639

  The sort algorithm breaks each name in groups of either only-digits
  or no-digits. Only the first eight such groups are considered, and
  after that we just use what's left of the string.

640 641 642 643
  @type name_list: list
  @param name_list: the names to be sorted
  @rtype: list
  @return: a copy of the name list sorted with our algorithm
Iustin Pop's avatar
Iustin Pop committed
644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668

  """
  _SORTER_BASE = "(\D+|\d+)"
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
                                                  _SORTER_BASE, _SORTER_BASE,
                                                  _SORTER_BASE, _SORTER_BASE,
                                                  _SORTER_BASE, _SORTER_BASE)
  _SORTER_RE = re.compile(_SORTER_FULL)
  _SORTER_NODIGIT = re.compile("^\D*$")
  def _TryInt(val):
    """Attempts to convert a variable to integer."""
    if val is None or _SORTER_NODIGIT.match(val):
      return val
    rval = int(val)
    return rval

  to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
             for name in name_list]
  to_sort.sort()
  return [tup[1] for tup in to_sort]


def TryConvert(fn, val):
  """Try to convert a value ignoring errors.

669 670 671 672 673 674 675 676 677 678
  This function tries to apply function I{fn} to I{val}. If no
  C{ValueError} or C{TypeError} exceptions are raised, it will return
  the result, else it will return the original value. Any other
  exceptions are propagated to the caller.

  @type fn: callable
  @param fn: function to apply to the value
  @param val: the value to be converted
  @return: The converted value if the conversion was successful,
      otherwise the original value.
Iustin Pop's avatar
Iustin Pop committed
679 680 681 682 683 684 685 686 687 688

  """
  try:
    nv = fn(val)
  except (ValueError, TypeError), err:
    nv = val
  return nv


def IsValidIP(ip):
689
  """Verifies the syntax of an IPv4 address.
Iustin Pop's avatar
Iustin Pop committed
690

691 692 693 694 695 696 697 698
  This function checks if the IPv4 address passes is valid or not based
  on syntax (not IP range, class calculations, etc.).

  @type ip: str
  @param ip: the address to be checked
  @rtype: a regular expression match object
  @return: a regular epression match object, or None if the
      address is not valid
Iustin Pop's avatar
Iustin Pop committed
699 700 701

  """
  unit = "(0|[1-9]\d{0,2})"
702
  #TODO: convert and return only boolean
Iustin Pop's avatar
Iustin Pop committed
703 704 705 706 707 708 709 710 711 712 713 714 715
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)


def IsValidShellParam(word):
  """Verifies is the given word is safe from the shell's p.o.v.

  This means that we can pass this to a command via the shell and be
  sure that it doesn't alter the command line and is passed as such to
  the actual command.

  Note that we are overly restrictive here, in order to be on the safe
  side.

716 717 718 719 720
  @type word: str
  @param word: the word to check
  @rtype: boolean
  @return: True if the word is 'safe'

Iustin Pop's avatar
Iustin Pop committed
721 722 723 724 725 726 727 728 729 730 731 732
  """
  return bool(re.match("^[-a-zA-Z0-9._+/:%@]+$", word))


def BuildShellCmd(template, *args):
  """Build a safe shell command line from the given arguments.

  This function will check all arguments in the args list so that they
  are valid shell parameters (i.e. they don't contain shell
  metacharaters). If everything is ok, it will return the result of
  template % args.

733 734 735 736 737 738
  @type template: str
  @param template: the string holding the template for the
      string formatting
  @rtype: str
  @return: the expanded command line

Iustin Pop's avatar
Iustin Pop committed
739 740 741
  """
  for word in args:
    if not IsValidShellParam(word):
742 743
      raise errors.ProgrammerError("Shell argument '%s' contains"
                                   " invalid characters" % word)
Iustin Pop's avatar
Iustin Pop committed
744 745 746
  return template % args


747
def FormatUnit(value, units):
Iustin Pop's avatar
Iustin Pop committed
748 749
  """Formats an incoming number of MiB with the appropriate unit.

750 751
  @type value: int
  @param value: integer representing the value in MiB (1048576)
752 753 754 755 756 757
  @type units: char
  @param units: the type of formatting we should do:
      - 'h' for automatic scaling
      - 'm' for MiBs
      - 'g' for GiBs
      - 't' for TiBs
758 759
  @rtype: str
  @return: the formatted value (with suffix)
Iustin Pop's avatar
Iustin Pop committed
760 761

  """
762 763
  if units not in ('m', 'g', 't', 'h'):
    raise errors.ProgrammerError("Invalid unit specified '%s'" % str(units))
Iustin Pop's avatar
Iustin Pop committed
764

765 766 767 768 769 770 771 772 773 774 775
  suffix = ''

  if units == 'm' or (units == 'h' and value < 1024):
    if units == 'h':
      suffix = 'M'
    return "%d%s" % (round(value, 0), suffix)

  elif units == 'g' or (units == 'h' and value < (1024 * 1024)):
    if units == 'h':
      suffix = 'G'
    return "%0.1f%s" % (round(float(value) / 1024, 1), suffix)
Iustin Pop's avatar
Iustin Pop committed
776 777

  else:
778 779 780
    if units == 'h':
      suffix = 'T'
    return "%0.1f%s" % (round(float(value) / 1024 / 1024, 1), suffix)
Iustin Pop's avatar
Iustin Pop committed
781 782 783 784 785


def ParseUnit(input_string):
  """Tries to extract number and scale from the given string.

786 787 788
  Input must be in the format C{NUMBER+ [DOT NUMBER+] SPACE*
  [UNIT]}. If no unit is specified, it defaults to MiB. Return value
  is always an int in MiB.
Iustin Pop's avatar
Iustin Pop committed
789 790

  """
791
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', str(input_string))
Iustin Pop's avatar
Iustin Pop committed
792
  if not m:
793
    raise errors.UnitParseError("Invalid format")
Iustin Pop's avatar
Iustin Pop committed
794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813

  value = float(m.groups()[0])

  unit = m.groups()[1]
  if unit:
    lcunit = unit.lower()
  else:
    lcunit = 'm'

  if lcunit in ('m', 'mb', 'mib'):
    # Value already in MiB
    pass

  elif lcunit in ('g', 'gb', 'gib'):
    value *= 1024

  elif lcunit in ('t', 'tb', 'tib'):
    value *= 1024 * 1024

  else:
814
    raise errors.UnitParseError("Unknown unit: %s" % unit)
Iustin Pop's avatar
Iustin Pop committed
815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830

  # Make sure we round up
  if int(value) < value:
    value += 1

  # Round up to the next multiple of 4
  value = int(value)
  if value % 4:
    value += 4 - value % 4

  return value


def AddAuthorizedKey(file_name, key):
  """Adds an SSH public key to an authorized_keys file.

831 832 833 834 835
  @type file_name: str
  @param file_name: path to authorized_keys file
  @type key: str
  @param key: string containing key

Iustin Pop's avatar
Iustin Pop committed
836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859
  """
  key_fields = key.split()

  f = open(file_name, 'a+')
  try:
    nl = True
    for line in f:
      # Ignore whitespace changes
      if line.split() == key_fields:
        break
      nl = line.endswith('\n')
    else:
      if not nl:
        f.write("\n")
      f.write(key.rstrip('\r\n'))
      f.write("\n")
      f.flush()
  finally:
    f.close()


def RemoveAuthorizedKey(file_name, key):
  """Removes an SSH public key from an authorized_keys file.

860 861 862 863 864
  @type file_name: str
  @param file_name: path to authorized_keys file
  @type key: str
  @param key: string containing key

Iustin Pop's avatar
Iustin Pop committed
865 866 867 868 869
  """
  key_fields = key.split()

  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
  try:
870
    out = os.fdopen(fd, 'w')
Iustin Pop's avatar
Iustin Pop committed
871
    try:
872 873 874 875 876 877
      f = open(file_name, 'r')
      try:
        for line in f:
          # Ignore whitespace changes while comparing lines
          if line.split() != key_fields:
            out.write(line)
878 879 880 881 882 883 884 885 886 887 888 889

        out.flush()
        os.rename(tmpname, file_name)
      finally:
        f.close()
    finally:
      out.close()
  except:
    RemoveFile(tmpname)
    raise


890 891
def SetEtcHostsEntry(file_name, ip, hostname, aliases):
  """Sets the name of an IP address and hostname in /etc/hosts.
892

893 894 895 896 897 898 899 900 901
  @type file_name: str
  @param file_name: path to the file to modify (usually C{/etc/hosts})
  @type ip: str
  @param ip: the IP address
  @type hostname: str
  @param hostname: the hostname to be added
  @type aliases: list
  @param aliases: the list of aliases to add for the hostname

902
  """
903
  # FIXME: use WriteFile + fn rather than duplicating its efforts
904 905 906
  # Ensure aliases are unique
  aliases = UniqueSequence([hostname] + aliases)[1:]

907
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
908
  try:
909 910 911 912 913 914
    out = os.fdopen(fd, 'w')
    try:
      f = open(file_name, 'r')
      try:
        for line in f:
          fields = line.split()
915
          if fields and not fields[0].startswith('#') and ip == fields[0]:
916 917 918
            continue
          out.write(line)

919
        out.write("%s\t%s" % (ip, hostname))
920 921 922 923 924
        if aliases:
          out.write(" %s" % ' '.join(aliases))
        out.write('\n')

        out.flush()
925
        os.fsync(out)
926
        os.chmod(tmpname, 0644)
927 928 929 930 931 932 933 934
        os.rename(tmpname, file_name)
      finally:
        f.close()
    finally:
      out.close()
  except:
    RemoveFile(tmpname)
    raise
935 936


937 938 939
def AddHostToEtcHosts(hostname):
  """Wrapper around SetEtcHostsEntry.

940 941 942 943
  @type hostname: str
  @param hostname: a hostname that will be resolved and added to
      L{constants.ETC_HOSTS}

944 945 946 947 948
  """
  hi = HostInfo(name=hostname)
  SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])


949
def RemoveEtcHostsEntry(file_name, hostname):
950
  """Removes a hostname from /etc/hosts.
951

952
  IP addresses without names are removed from the file.
953 954 955 956 957 958

  @type file_name: str
  @param file_name: path to the file to modify (usually C{/etc/hosts})
  @type hostname: str
  @param hostname: the hostname to be removed

959
  """
960
  # FIXME: use WriteFile + fn rather than duplicating its efforts
961 962 963 964 965 966 967 968 969 970 971 972 973 974
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
  try:
    out = os.fdopen(fd, 'w')
    try:
      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:
975
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
976 977 978
              continue

          out.write(line)
979 980

        out.flush()
981
        os.fsync(out)
982
        os.chmod(tmpname, 0644)
983 984 985
        os.rename(tmpname, file_name)
      finally:
        f.close()
Iustin Pop's avatar
Iustin Pop committed
986
    finally:
987 988 989 990
      out.close()
  except:
    RemoveFile(tmpname)
    raise
Iustin Pop's avatar
Iustin Pop committed
991 992


993 994 995
def RemoveHostFromEtcHosts(hostname):
  """Wrapper around RemoveEtcHostsEntry.

996 997 998 999 1000
  @type hostname: str
  @param hostname: hostname that will be resolved and its
      full and shot name will be removed from
      L{constants.ETC_HOSTS}

1001 1002 1003 1004 1005 1006
  """
  hi = HostInfo(name=hostname)
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.name)
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.ShortName())


Iustin Pop's avatar
Iustin Pop committed
1007 1008 1009
def CreateBackup(file_name):
  """Creates a backup of a file.

1010 1011 1012 1013 1014
  @type file_name: str
  @param file_name: file to be backed up
  @rtype: str
  @return: the path to the newly created backup
  @raise errors.ProgrammerError: for invalid file names
Iustin Pop's avatar
Iustin Pop committed
1015 1016 1017

  """
  if not os.path.isfile(file_name):
1018 1019
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
                                file_name)
Iustin Pop's avatar
Iustin Pop committed
1020

1021
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
Iustin Pop's avatar
Iustin Pop committed
1022
  dir_name = os.path.dirname(file_name)
1023 1024 1025

  fsrc = open(file_name, 'rb')
  try:
Iustin Pop's avatar
Iustin Pop committed
1026
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
1027 1028 1029 1030 1031 1032 1033 1034
    fdst = os.fdopen(fd, 'wb')
    try:
      shutil.copyfileobj(fsrc, fdst)
    finally:
      fdst.close()
  finally:
    fsrc.close()

Iustin Pop's avatar
Iustin Pop committed
1035 1036 1037 1038 1039
  return backup_name


def ShellQuote(value):
  """Quotes shell argument according to POSIX.
1040

1041 1042 1043 1044 1045
  @type value: str
  @param value: the argument to be quoted
  @rtype: str
  @return: the quoted value

Iustin Pop's avatar
Iustin Pop committed
1046 1047 1048 1049 1050 1051 1052 1053
  """
  if _re_shell_unquoted.match(value):
    return value
  else:
    return "'%s'" % value.replace("'", "'\\''")


def ShellQuoteArgs(args):
1054 1055 1056 1057 1058 1059
  """Quotes a list of shell arguments.

  @type args: list
  @param args: list of arguments to be quoted
  @rtype: str
  @return: the quoted arguments concatenaned with spaces
Iustin Pop's avatar
Iustin Pop committed
1060 1061 1062

  """
  return ' '.join([ShellQuote(i) for i in args])
1063 1064


1065
def TcpPing(target, port, timeout=10, live_port_needed=False, source=None):
1066 1067
  """Simple ping implementation using TCP connect(2).

1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083
  Check if the given IP is reachable by doing attempting a TCP connect
  to it.

  @type target: str
  @param target: the IP or hostname to ping
  @type port: int
  @param port: the port to connect to
  @type timeout: int
  @param timeout: the timeout on the connection attemp
  @type live_port_needed: boolean
  @param live_port_needed: whether a closed port will cause the
      function to return failure, as if there was a timeout
  @type source: str or None
  @param source: if specified, will cause the connect to be made
      from this specific source address; failures to bind other
      than C{EADDRNOTAVAIL} will be ignored
1084 1085 1086 1087

  """
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

1088
  success = False
1089

1090 1091 1092 1093 1094 1095
  if source is not None:
    try:
      sock.bind((source, 0))
    except socket.error, (errcode, errstring):
      if errcode == errno.EADDRNOTAVAIL:
        success = False
1096 1097 1098 1099 1100 1101 1102 1103 1104 1105

  sock.settimeout(timeout)

  try:
    sock.connect((target, port))
    sock.close()
    success = True
  except socket.timeout:
    success = False
  except socket.error, (errcode, errstring):
1106
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
1107 1108

  return success
1109 1110


1111 1112 1113
def OwnIpAddress(address):
  """Check if the current host has the the given IP address.

1114
  Currently this is done by TCP-pinging the address from the loopback
1115 1116 1117 1118 1119
  address.

  @type address: string
  @param address: the addres to check
  @rtype: bool
1120
  @return: True if we own the address
1121 1122 1123 1124 1125 1126

  """
  return TcpPing(address, constants.DEFAULT_NODED_PORT,
                 source=constants.LOCALHOST_IP_ADDRESS)


1127
def ListVisibleFiles(path):
1128 1129 1130 1131 1132 1133
  """Returns a list of visible files in a directory.

  @type path: str
  @param path: the directory to enumerate
  @rtype: list
  @return: the list of all files not starting with a dot
1134 1135

  """
1136 1137 1138
  files = [i for i in os.listdir(path) if not i.startswith(".")]
  files.sort()
  return files
1139 1140


1141 1142 1143 1144 1145 1146
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.
1147 1148 1149

  """
  try:
1150 1151 1152 1153 1154 1155 1156
    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))
1157 1158 1159
  except KeyError:
    return default
  return result.pw_dir
1160 1161


1162
def NewUUID():
1163 1164
  """Returns a random UUID.

1165 1166 1167 1168
  @note: This is a Linux-specific method as it uses the /proc
      filesystem.
  @rtype: str

1169 1170 1171 1172 1173 1174
  """
  f = open("/proc/sys/kernel/random/uuid", "r")
  try:
    return f.read(128).rstrip("\n")
  finally:
    f.close()
Iustin Pop's avatar
Iustin Pop committed
1175 1176


1177 1178 1179 1180 1181 1182
def GenerateSecret():
  """Generates a random secret.

  This will generate a pseudo-random secret, and return its sha digest
  (so that it can be used where an ASCII string is needed).

1183 1184 1185
  @rtype: str
  @return: a sha1 hexdigest of a block of 64 random bytes

1186
  """
1187
  return sha1(os.urandom(64)).hexdigest()
1188 1189


1190 1191 1192 1193 1194 1195 1196 1197 1198
def EnsureDirs(dirs):
  """Make required directories, if they don't exist.

  @param dirs: list of tuples (dir_name, dir_mode)
  @type dirs: list of (string, integer)

  """
  for dir_name, dir_mode in dirs:
    try:
Iustin Pop's avatar
Iustin Pop committed
1199
      os.mkdir(dir_name, dir_mode)
1200 1201 1202
    except EnvironmentError, err:
      if err.errno != errno.EEXIST:
        raise errors.GenericError("Cannot create needed directory"