utils.py 110 KB
Newer Older
Iustin Pop's avatar
Iustin Pop committed
1
#
Iustin Pop's avatar
Iustin Pop committed
2
3
#

4
# Copyright (C) 2006, 2007, 2010 Google Inc.
Iustin Pop's avatar
Iustin Pop committed
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
#
# 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
"""


import os
31
import sys
Iustin Pop's avatar
Iustin Pop committed
32
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
45
import logging.handlers
Michael Hanselmann's avatar
Michael Hanselmann committed
46
import signal
47
import OpenSSL
48
49
import datetime
import calendar
50
import hmac
51
import collections
52
53

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

Luca Bigliardi's avatar
Luca Bigliardi committed
55
try:
56
  # pylint: disable-msg=F0401
Luca Bigliardi's avatar
Luca Bigliardi committed
57
58
59
60
  import ctypes
except ImportError:
  ctypes = None

Iustin Pop's avatar
Iustin Pop committed
61
from ganeti import errors
Iustin Pop's avatar
Iustin Pop committed
62
from ganeti import constants
63
from ganeti import compat
Iustin Pop's avatar
Iustin Pop committed
64

65

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

69
debug_locks = False
70
71

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

74
75
_RANDOM_UUID_FILE = "/proc/sys/kernel/random/uuid"

76
77
78
79
80
81
82
HEX_CHAR_RE = r"[a-zA-Z0-9]"
VALID_X509_SIGNATURE_SALT = re.compile("^%s+$" % HEX_CHAR_RE, re.S)
X509_SIGNATURE = re.compile(r"^%s:\s*(?P<salt>%s+)/(?P<sign>%s+)$" %
                            (re.escape(constants.X509_CERT_SIGNATURE_HEADER),
                             HEX_CHAR_RE, HEX_CHAR_RE),
                            re.S | re.I)

83
84
_VALID_SERVICE_NAME_RE = re.compile("^[-_.a-zA-Z0-9]{1,128}$")

85
86
87
UUID_RE = re.compile('^[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-'
                     '[a-f0-9]{4}-[a-f0-9]{12}$')

88
89
90
91
# Certificate verification results
(CERT_WARNING,
 CERT_ERROR) = range(1, 3)

Luca Bigliardi's avatar
Luca Bigliardi committed
92
93
94
95
# Flags for mlockall() (from bits/mman.h)
_MCL_CURRENT = 1
_MCL_FUTURE = 2

Iustin Pop's avatar
Iustin Pop committed
96
97
98
#: MAC checker regexp
_MAC_CHECK = re.compile("^([0-9a-f]{2}:){5}[0-9a-f]{2}$", re.I)

René Nussbaumer's avatar
René Nussbaumer committed
99
100
101
102
(_TIMEOUT_NONE,
 _TIMEOUT_TERM,
 _TIMEOUT_KILL) = range(3)

103
104
105
106
107
108
109
110
111
#: Shell param checker regexp
_SHELLPARAM_REGEX = re.compile(r"^[-a-zA-Z0-9._+/:%@]+$")

#: Unit checker regexp
_PARSEUNIT_REGEX = re.compile(r"^([.\d]+)\s*([a-zA-Z]+)?$")

#: ASN1 time regexp
_ANS1_TIME_REGEX = re.compile(r"^(\d+)([-+]\d\d)(\d\d)$")

112

Iustin Pop's avatar
Iustin Pop committed
113
class RunResult(object):
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
  """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
130
131
132
133
134
135

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


René Nussbaumer's avatar
René Nussbaumer committed
136
137
  def __init__(self, exit_code, signal_, stdout, stderr, cmd, timeout_action,
               timeout):
Iustin Pop's avatar
Iustin Pop committed
138
139
    self.cmd = cmd
    self.exit_code = exit_code
Iustin Pop's avatar
Iustin Pop committed
140
    self.signal = signal_
Iustin Pop's avatar
Iustin Pop committed
141
142
    self.stdout = stdout
    self.stderr = stderr
Iustin Pop's avatar
Iustin Pop committed
143
    self.failed = (signal_ is not None or exit_code != 0)
Iustin Pop's avatar
Iustin Pop committed
144

René Nussbaumer's avatar
René Nussbaumer committed
145
    fail_msgs = []
Iustin Pop's avatar
Iustin Pop committed
146
    if self.signal is not None:
René Nussbaumer's avatar
René Nussbaumer committed
147
      fail_msgs.append("terminated by signal %s" % self.signal)
Iustin Pop's avatar
Iustin Pop committed
148
    elif self.exit_code is not None:
René Nussbaumer's avatar
René Nussbaumer committed
149
      fail_msgs.append("exited with exit code %s" % self.exit_code)
Iustin Pop's avatar
Iustin Pop committed
150
    else:
René Nussbaumer's avatar
René Nussbaumer committed
151
152
153
154
155
156
157
158
159
160
161
      fail_msgs.append("unable to determine termination reason")

    if timeout_action == _TIMEOUT_TERM:
      fail_msgs.append("terminated after timeout of %.2f seconds" % timeout)
    elif timeout_action == _TIMEOUT_KILL:
      fail_msgs.append(("force termination after timeout of %.2f seconds"
                        " and linger for another %.2f seconds") %
                       (timeout, constants.CHILD_LINGER_TIMEOUT))

    if fail_msgs and self.failed:
      self.fail_reason = CommaJoin(fail_msgs)
Iustin Pop's avatar
Iustin Pop committed
162

163
164
165
    if self.failed:
      logging.debug("Command '%s' failed (%s); output: %s",
                    self.cmd, self.fail_reason, self.output)
166

Iustin Pop's avatar
Iustin Pop committed
167
168
169
170
171
172
173
174
175
  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")


176
def _BuildCmdEnvironment(env, reset):
177
178
179
  """Builds the environment for an external program.

  """
180
181
182
183
184
185
  if reset:
    cmd_env = {}
  else:
    cmd_env = os.environ.copy()
    cmd_env["LC_ALL"] = "C"

186
187
  if env is not None:
    cmd_env.update(env)
188

189
190
191
  return cmd_env


192
def RunCmd(cmd, env=None, output=None, cwd="/", reset_env=False,
René Nussbaumer's avatar
René Nussbaumer committed
193
           interactive=False, timeout=None):
Iustin Pop's avatar
Iustin Pop committed
194
195
196
197
198
  """Execute a (shell) command.

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

199
  @type cmd: string or list
200
  @param cmd: Command to run
201
  @type env: dict
202
  @param env: Additional environment variables
203
  @type output: str
204
  @param output: if desired, the output of the command can be
205
206
      saved in a file instead of the RunResult instance; this
      parameter denotes the file name (if not None)
207
208
209
  @type cwd: string
  @param cwd: if specified, will be used as the working
      directory for the command; the default will be /
Guido Trotter's avatar
Guido Trotter committed
210
211
  @type reset_env: boolean
  @param reset_env: whether to reset or keep the default os environment
212
213
214
  @type interactive: boolean
  @param interactive: weather we pipe stdin, stdout and stderr
                      (default behaviour) or run the command interactive
René Nussbaumer's avatar
René Nussbaumer committed
215
216
217
  @type timeout: int
  @param timeout: If not None, timeout in seconds until child process gets
                  killed
218
  @rtype: L{RunResult}
219
  @return: RunResult instance
Michael Hanselmann's avatar
Michael Hanselmann committed
220
  @raise errors.ProgrammerError: if we call this when forks are disabled
Iustin Pop's avatar
Iustin Pop committed
221
222

  """
223
224
225
  if no_fork:
    raise errors.ProgrammerError("utils.RunCmd() called with fork() disabled")

226
227
228
229
  if output and interactive:
    raise errors.ProgrammerError("Parameters 'output' and 'interactive' can"
                                 " not be provided at the same time")

230
231
232
233
  if isinstance(cmd, basestring):
    strcmd = cmd
    shell = True
  else:
Iustin Pop's avatar
Iustin Pop committed
234
    cmd = [str(val) for val in cmd]
235
    strcmd = ShellQuoteArgs(cmd)
236
    shell = False
237
238
239

  if output:
    logging.debug("RunCmd %s, output file '%s'", strcmd, output)
240
  else:
241
    logging.debug("RunCmd %s", strcmd)
242

243
  cmd_env = _BuildCmdEnvironment(env, reset_env)
244

245
246
  try:
    if output is None:
René Nussbaumer's avatar
René Nussbaumer committed
247
248
      out, err, status, timeout_action = _RunCmdPipe(cmd, cmd_env, shell, cwd,
                                                     interactive, timeout)
249
    else:
René Nussbaumer's avatar
René Nussbaumer committed
250
      timeout_action = _TIMEOUT_NONE
251
252
253
254
255
256
257
258
      status = _RunCmdFile(cmd, cmd_env, shell, output, cwd)
      out = err = ""
  except OSError, err:
    if err.errno == errno.ENOENT:
      raise errors.OpExecError("Can't execute '%s': not found (%s)" %
                               (strcmd, err))
    else:
      raise
259
260
261
262
263
264
265
266

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

René Nussbaumer's avatar
René Nussbaumer committed
267
  return RunResult(exitcode, signal_, out, err, strcmd, timeout_action, timeout)
268

269

270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
def SetupDaemonEnv(cwd="/", umask=077):
  """Setup a daemon's environment.

  This should be called between the first and second fork, due to
  setsid usage.

  @param cwd: the directory to which to chdir
  @param umask: the umask to setup

  """
  os.chdir(cwd)
  os.umask(umask)
  os.setsid()


285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
def SetupDaemonFDs(output_file, output_fd):
  """Setups up a daemon's file descriptors.

  @param output_file: if not None, the file to which to redirect
      stdout/stderr
  @param output_fd: if not None, the file descriptor for stdout/stderr

  """
  # check that at most one is defined
  assert [output_file, output_fd].count(None) >= 1

  # Open /dev/null (read-only, only for stdin)
  devnull_fd = os.open(os.devnull, os.O_RDONLY)

  if output_fd is not None:
    pass
  elif output_file is not None:
    # Open output file
    try:
      output_fd = os.open(output_file,
                          os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0600)
    except EnvironmentError, err:
      raise Exception("Opening output file failed: %s" % err)
  else:
    output_fd = os.open(os.devnull, os.O_WRONLY)

  # Redirect standard I/O
  os.dup2(devnull_fd, 0)
  os.dup2(output_fd, 1)
  os.dup2(output_fd, 2)


317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
def StartDaemon(cmd, env=None, cwd="/", output=None, output_fd=None,
                pidfile=None):
  """Start a daemon process after forking twice.

  @type cmd: string or list
  @param cmd: Command to run
  @type env: dict
  @param env: Additional environment variables
  @type cwd: string
  @param cwd: Working directory for the program
  @type output: string
  @param output: Path to file in which to save the output
  @type output_fd: int
  @param output_fd: File descriptor for output
  @type pidfile: string
  @param pidfile: Process ID file
  @rtype: int
  @return: Daemon process ID
  @raise errors.ProgrammerError: if we call this when forks are disabled

  """
  if no_fork:
    raise errors.ProgrammerError("utils.StartDaemon() called with fork()"
                                 " disabled")

  if output and not (bool(output) ^ (output_fd is not None)):
    raise errors.ProgrammerError("Only one of 'output' and 'output_fd' can be"
                                 " specified")

  if isinstance(cmd, basestring):
    cmd = ["/bin/sh", "-c", cmd]

  strcmd = ShellQuoteArgs(cmd)

  if output:
    logging.debug("StartDaemon %s, output file '%s'", strcmd, output)
  else:
    logging.debug("StartDaemon %s", strcmd)

356
  cmd_env = _BuildCmdEnvironment(env, False)
357
358
359
360
361
362
363
364
365
366
367
368
369
370

  # Create pipe for sending PID back
  (pidpipe_read, pidpipe_write) = os.pipe()
  try:
    try:
      # Create pipe for sending error messages
      (errpipe_read, errpipe_write) = os.pipe()
      try:
        try:
          # First fork
          pid = os.fork()
          if pid == 0:
            try:
              # Child process, won't return
371
372
373
374
              _StartDaemonChild(errpipe_read, errpipe_write,
                                pidpipe_read, pidpipe_write,
                                cmd, cmd_env, cwd,
                                output, output_fd, pidfile)
375
376
            finally:
              # Well, maybe child process failed
377
              os._exit(1) # pylint: disable-msg=W0212
378
379
380
        finally:
          _CloseFDNoErr(errpipe_write)

381
382
        # Wait for daemon to be started (or an error message to
        # arrive) and read up to 100 KB as an error message
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
        errormsg = RetryOnSignal(os.read, errpipe_read, 100 * 1024)
      finally:
        _CloseFDNoErr(errpipe_read)
    finally:
      _CloseFDNoErr(pidpipe_write)

    # Read up to 128 bytes for PID
    pidtext = RetryOnSignal(os.read, pidpipe_read, 128)
  finally:
    _CloseFDNoErr(pidpipe_read)

  # Try to avoid zombies by waiting for child process
  try:
    os.waitpid(pid, 0)
  except OSError:
    pass

  if errormsg:
    raise errors.OpExecError("Error when starting daemon process: %r" %
                             errormsg)

  try:
    return int(pidtext)
  except (ValueError, TypeError), err:
    raise errors.OpExecError("Error while trying to parse PID %r: %s" %
                             (pidtext, err))


411
412
413
414
def _StartDaemonChild(errpipe_read, errpipe_write,
                      pidpipe_read, pidpipe_write,
                      args, env, cwd,
                      output, fd_output, pidfile):
415
416
417
418
419
420
421
422
423
  """Child process for starting daemon.

  """
  try:
    # Close parent's side
    _CloseFDNoErr(errpipe_read)
    _CloseFDNoErr(pidpipe_read)

    # First child process
424
    SetupDaemonEnv()
425
426
427
428
429
430
431

    # And fork for the second time
    pid = os.fork()
    if pid != 0:
      # Exit first child process
      os._exit(0) # pylint: disable-msg=W0212

432
433
    # Make sure pipe is closed on execv* (and thereby notifies
    # original process)
434
435
436
437
438
439
440
    SetCloseOnExecFlag(errpipe_write, True)

    # List of file descriptors to be left open
    noclose_fds = [errpipe_write]

    # Open PID file
    if pidfile:
441
      fd_pidfile = WritePidFile(pidfile)
442
443
444
445
446
447
448
449

      # Keeping the file open to hold the lock
      noclose_fds.append(fd_pidfile)

      SetCloseOnExecFlag(fd_pidfile, False)
    else:
      fd_pidfile = None

450
    SetupDaemonFDs(output, fd_output)
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467

    # Send daemon PID to parent
    RetryOnSignal(os.write, pidpipe_write, str(os.getpid()))

    # Close all file descriptors except stdio and error message pipe
    CloseFDs(noclose_fds=noclose_fds)

    # Change working directory
    os.chdir(cwd)

    if env is None:
      os.execvp(args[0], args)
    else:
      os.execvpe(args[0], args, env)
  except: # pylint: disable-msg=W0702
    try:
      # Report errors to original process
468
      WriteErrorToFD(errpipe_write, str(sys.exc_info()[1]))
469
470
471
472
473
474
475
    except: # pylint: disable-msg=W0702
      # Ignore errors in error handling
      pass

  os._exit(1) # pylint: disable-msg=W0212


476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
def WriteErrorToFD(fd, err):
  """Possibly write an error message to a fd.

  @type fd: None or int (file descriptor)
  @param fd: if not None, the error will be written to this fd
  @param err: string, the error message

  """
  if fd is None:
    return

  if not err:
    err = "<unknown error>"

  RetryOnSignal(os.write, fd, err)


René Nussbaumer's avatar
René Nussbaumer committed
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
def _CheckIfAlive(child):
  """Raises L{RetryAgain} if child is still alive.

  @raises RetryAgain: If child is still alive

  """
  if child.poll() is None:
    raise RetryAgain()


def _WaitForProcess(child, timeout):
  """Waits for the child to terminate or until we reach timeout.

  """
  try:
    Retry(_CheckIfAlive, (1.0, 1.2, 5.0), max(0, timeout), args=[child])
  except RetryTimeout:
    pass


def _RunCmdPipe(cmd, env, via_shell, cwd, interactive, timeout,
                _linger_timeout=constants.CHILD_LINGER_TIMEOUT):
515
516
517
518
519
520
521
522
  """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
523
524
  @type cwd: string
  @param cwd: the working directory for the program
525
526
  @type interactive: boolean
  @param interactive: Run command interactive (without piping)
René Nussbaumer's avatar
René Nussbaumer committed
527
528
  @type timeout: int
  @param timeout: Timeout after the programm gets terminated
529
530
531
532
  @rtype: tuple
  @return: (out, err, status)

  """
533
  poller = select.poll()
534
535
536
537
538
539
540
541

  stderr = subprocess.PIPE
  stdout = subprocess.PIPE
  stdin = subprocess.PIPE

  if interactive:
    stderr = stdout = stdin = None

542
  child = subprocess.Popen(cmd, shell=via_shell,
543
544
545
                           stderr=stderr,
                           stdout=stdout,
                           stdin=stdin,
546
547
                           close_fds=True, env=env,
                           cwd=cwd)
548

549
550
  out = StringIO()
  err = StringIO()
René Nussbaumer's avatar
René Nussbaumer committed
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565

  linger_timeout = None

  if timeout is None:
    poll_timeout = None
  else:
    poll_timeout = RunningTimeout(timeout, True).Remaining

  msg_timeout = ("Command %s (%d) run into execution timeout, terminating" %
                 (cmd, child.pid))
  msg_linger = ("Command %s (%d) run into linger timeout, killing" %
                (cmd, child.pid))

  timeout_action = _TIMEOUT_NONE

566
567
568
569
570
571
572
573
574
575
576
577
  if not interactive:
    child.stdin.close()
    poller.register(child.stdout, select.POLLIN)
    poller.register(child.stderr, select.POLLIN)
    fdmap = {
      child.stdout.fileno(): (out, child.stdout),
      child.stderr.fileno(): (err, child.stderr),
      }
    for fd in fdmap:
      SetNonblockFlag(fd, True)

    while fdmap:
René Nussbaumer's avatar
René Nussbaumer committed
578
      if poll_timeout:
579
580
        pt = poll_timeout() * 1000
        if pt < 0:
René Nussbaumer's avatar
René Nussbaumer committed
581
582
583
584
585
586
          if linger_timeout is None:
            logging.warning(msg_timeout)
            if child.poll() is None:
              timeout_action = _TIMEOUT_TERM
              IgnoreProcessNotFound(os.kill, child.pid, signal.SIGTERM)
            linger_timeout = RunningTimeout(_linger_timeout, True).Remaining
587
588
          pt = linger_timeout() * 1000
          if pt < 0:
René Nussbaumer's avatar
René Nussbaumer committed
589
590
591
592
593
            break
      else:
        pt = None

      pollresult = RetryOnSignal(poller.poll, pt)
594
595
596
597
598
599
600
601
602
603
604
605

      for fd, event in pollresult:
        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):
606
607
608
          poller.unregister(fd)
          del fdmap[fd]

René Nussbaumer's avatar
René Nussbaumer committed
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
  if timeout is not None:
    assert callable(poll_timeout)

    # We have no I/O left but it might still run
    if child.poll() is None:
      _WaitForProcess(child, poll_timeout())

    # Terminate if still alive after timeout
    if child.poll() is None:
      if linger_timeout is None:
        logging.warning(msg_timeout)
        timeout_action = _TIMEOUT_TERM
        IgnoreProcessNotFound(os.kill, child.pid, signal.SIGTERM)
        lt = _linger_timeout
      else:
        lt = linger_timeout()
      _WaitForProcess(child, lt)

    # Okay, still alive after timeout and linger timeout? Kill it!
    if child.poll() is None:
      timeout_action = _TIMEOUT_KILL
      logging.warning(msg_linger)
      IgnoreProcessNotFound(os.kill, child.pid, signal.SIGKILL)

633
634
  out = out.getvalue()
  err = err.getvalue()
Iustin Pop's avatar
Iustin Pop committed
635
636

  status = child.wait()
René Nussbaumer's avatar
René Nussbaumer committed
637
  return out, err, status, timeout_action
Iustin Pop's avatar
Iustin Pop committed
638

639

640
def _RunCmdFile(cmd, env, via_shell, output, cwd):
641
642
643
644
645
646
647
648
649
650
  """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
651
652
  @type cwd: string
  @param cwd: the working directory for the program
653
654
655
656
657
658
659
660
661
662
  @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,
663
664
                             close_fds=True, env=env,
                             cwd=cwd)
665
666
667
668
669
670

    child.stdin.close()
    status = child.wait()
  finally:
    fh.close()
  return status
Iustin Pop's avatar
Iustin Pop committed
671
672


673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
def SetCloseOnExecFlag(fd, enable):
  """Sets or unsets the close-on-exec flag on a file descriptor.

  @type fd: int
  @param fd: File descriptor
  @type enable: bool
  @param enable: Whether to set or unset it.

  """
  flags = fcntl.fcntl(fd, fcntl.F_GETFD)

  if enable:
    flags |= fcntl.FD_CLOEXEC
  else:
    flags &= ~fcntl.FD_CLOEXEC

  fcntl.fcntl(fd, fcntl.F_SETFD, flags)


692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
def SetNonblockFlag(fd, enable):
  """Sets or unsets the O_NONBLOCK flag on on a file descriptor.

  @type fd: int
  @param fd: File descriptor
  @type enable: bool
  @param enable: Whether to set or unset it

  """
  flags = fcntl.fcntl(fd, fcntl.F_GETFL)

  if enable:
    flags |= os.O_NONBLOCK
  else:
    flags &= ~os.O_NONBLOCK

  fcntl.fcntl(fd, fcntl.F_SETFL, flags)


711
712
713
714
715
716
717
def RetryOnSignal(fn, *args, **kwargs):
  """Calls a function again if it failed due to EINTR.

  """
  while True:
    try:
      return fn(*args, **kwargs)
718
    except EnvironmentError, err:
719
720
      if err.errno != errno.EINTR:
        raise
721
722
723
    except (socket.error, select.error), err:
      # In python 2.6 and above select.error is an IOError, so it's handled
      # above, in 2.5 and below it's not, and it's handled here.
724
725
726
727
      if not (err.args and err.args[0] == errno.EINTR):
        raise


728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
def RunParts(dir_name, env=None, reset_env=False):
  """Run Scripts or programs in a directory

  @type dir_name: string
  @param dir_name: absolute path to a directory
  @type env: dict
  @param env: The environment to use
  @type reset_env: boolean
  @param reset_env: whether to reset or keep the default os environment
  @rtype: list of tuples
  @return: list of (name, (one of RUNDIR_STATUS), RunResult)

  """
  rr = []

  try:
    dir_contents = ListVisibleFiles(dir_name)
  except OSError, err:
    logging.warning("RunParts: skipping %s (cannot list: %s)", dir_name, err)
    return rr

  for relname in sorted(dir_contents):
750
    fname = PathJoin(dir_name, relname)
751
752
753
754
755
756
757
758
759
760
761
762
763
764
    if not (os.path.isfile(fname) and os.access(fname, os.X_OK) and
            constants.EXT_PLUGIN_MASK.match(relname) is not None):
      rr.append((relname, constants.RUNPARTS_SKIP, None))
    else:
      try:
        result = RunCmd([fname], env=env, reset_env=reset_env)
      except Exception, err: # pylint: disable-msg=W0703
        rr.append((relname, constants.RUNPARTS_ERR, str(err)))
      else:
        rr.append((relname, constants.RUNPARTS_RUN, result))

  return rr


Iustin Pop's avatar
Iustin Pop committed
765
766
767
768
769
770
def RemoveFile(filename):
  """Remove a file ignoring some errors.

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

771
772
773
  @type filename: str
  @param filename: the file to be removed

Iustin Pop's avatar
Iustin Pop committed
774
775
776
777
  """
  try:
    os.unlink(filename)
  except OSError, err:
778
    if err.errno not in (errno.ENOENT, errno.EISDIR):
Iustin Pop's avatar
Iustin Pop committed
779
780
      raise

Balazs Lecz's avatar
Balazs Lecz committed
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798

def RemoveDir(dirname):
  """Remove an empty directory.

  Remove a directory, ignoring non-existing ones.
  Other errors are passed. This includes the case,
  where the directory is not empty, so it can't be removed.

  @type dirname: str
  @param dirname: the empty directory to be removed

  """
  try:
    os.rmdir(dirname)
  except OSError, err:
    if err.errno != errno.ENOENT:
      raise

Iustin Pop's avatar
Iustin Pop committed
799

800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
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
821
      Makedirs(os.path.dirname(new), mode=mkdir_mode)
822

823
      return os.rename(old, new)
824

825
826
827
    raise


828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
def Makedirs(path, mode=0750):
  """Super-mkdir; create a leaf directory and all intermediate ones.

  This is a wrapper around C{os.makedirs} adding error handling not implemented
  before Python 2.5.

  """
  try:
    os.makedirs(path, mode)
  except OSError, err:
    # Ignore EEXIST. This is only handled in os.makedirs as included in
    # Python 2.5 and above.
    if err.errno != errno.EEXIST or not os.path.exists(path):
      raise


844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
def ResetTempfileModule():
  """Resets the random name generator of the tempfile module.

  This function should be called after C{os.fork} in the child process to
  ensure it creates a newly seeded random generator. Otherwise it would
  generate the same random parts as the parent process. If several processes
  race for the creation of a temporary file, this could lead to one not getting
  a temporary name.

  """
  # pylint: disable-msg=W0212
  if hasattr(tempfile, "_once_lock") and hasattr(tempfile, "_name_sequence"):
    tempfile._once_lock.acquire()
    try:
      # Reset random name generator
      tempfile._name_sequence = None
    finally:
      tempfile._once_lock.release()
  else:
    logging.critical("The tempfile module misses at least one of the"
                     " '_once_lock' and '_name_sequence' attributes")


Iustin Pop's avatar
Iustin Pop committed
867
868
869
870
871
872
def _FingerprintFile(filename):
  """Compute the fingerprint of a file.

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

873
874
875
876
877
  @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
878
879
880
881
882
883
884

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

  f = open(filename)

885
  fp = compat.sha1_hash()
Iustin Pop's avatar
Iustin Pop committed
886
887
888
889
890
891
892
893
894
895
896
897
898
  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.

899
900
901
902
903
  @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
904
905
906
907
908
909
910
911
912
913
914
915

  """
  ret = {}

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

  return ret


916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
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 = []

931
932
933
934
  if not isinstance(target, dict):
    msg = "Expected dictionary, got '%s'" % target
    raise errors.TypeEnforcementError(msg)

935
936
937
938
939
940
941
942
  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

Iustin Pop's avatar
Iustin Pop committed
943
944
945
    ktype = key_types[key]
    if ktype not in constants.ENFORCEABLE_TYPES:
      msg = "'%s' has non-enforceable type %s" % (key, ktype)
946
947
      raise errors.ProgrammerError(msg)

948
949
950
951
    if ktype in (constants.VTYPE_STRING, constants.VTYPE_MAYBE_STRING):
      if target[key] is None and ktype == constants.VTYPE_MAYBE_STRING:
        pass
      elif not isinstance(target[key], basestring):
952
953
954
955
956
        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)
Iustin Pop's avatar
Iustin Pop committed
957
    elif ktype == constants.VTYPE_BOOL:
958
959
960
961
962
963
964
965
966
967
968
969
      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
Iustin Pop's avatar
Iustin Pop committed
970
    elif ktype == constants.VTYPE_SIZE:
971
972
973
974
975
976
      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)
Iustin Pop's avatar
Iustin Pop committed
977
    elif ktype == constants.VTYPE_INT:
978
979
980
981
982
983
984
      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)


985
986
987
988
989
990
991
992
993
994
995
def _GetProcStatusPath(pid):
  """Returns the path for a PID's proc status file.

  @type pid: int
  @param pid: Process ID
  @rtype: string

  """
  return "/proc/%d/status" % pid


Iustin Pop's avatar
Iustin Pop committed
996
997
998
def IsProcessAlive(pid):
  """Check if a given pid exists on the system.

999
1000
  @note: zombie status is not handled, so zombie processes
      will be returned as alive
1001
1002
1003
1004
  @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
1005
1006

  """
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
  def _TryStat(name):
    try:
      os.stat(name)
      return True
    except EnvironmentError, err:
      if err.errno in (errno.ENOENT, errno.ENOTDIR):
        return False
      elif err.errno == errno.EINVAL:
        raise RetryAgain(err)
      raise

  assert isinstance(pid, int), "pid must be an integer"
1019
1020
1021
  if pid <= 0:
    return False

1022
1023
  # /proc in a multiprocessor environment can have strange behaviors.
  # Retry the os.stat a few times until we get a good result.
Iustin Pop's avatar
Iustin Pop committed
1024
  try:
1025
1026
    return Retry(_TryStat, (0.01, 1.5, 0.1), 0.5,
                 args=[_GetProcStatusPath(pid)])
1027
1028
  except RetryTimeout, err:
    err.RaiseInner()
Iustin Pop's avatar
Iustin Pop committed
1029
1030


1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
def _ParseSigsetT(sigset):
  """Parse a rendered sigset_t value.

  This is the opposite of the Linux kernel's fs/proc/array.c:render_sigset_t
  function.

  @type sigset: string
  @param sigset: Rendered signal set from /proc/$pid/status
  @rtype: set
  @return: Set of all enabled signal numbers

  """
  result = set()

  signum = 0
  for ch in reversed(sigset):
    chv = int(ch, 16)

    # The following could be done in a loop, but it's easier to read and
    # understand in the unrolled form
    if chv & 1:
      result.add(signum + 1)
    if chv & 2:
      result.add(signum + 2)
    if chv & 4:
      result.add(signum + 3)
    if chv & 8:
      result.add(signum + 4)

    signum += 4

  return result


def _GetProcStatusField(pstatus, field):
  """Retrieves a field from the contents of a proc status file.

  @type pstatus: string
  @param pstatus: Contents of /proc/$pid/status
  @type field: string
  @param field: Name of field whose value should be returned
  @rtype: string

  """
  for line in pstatus.splitlines():
    parts = line.split(":", 1)

    if len(parts) < 2 or parts[0] != field:
      continue

    return parts[1].strip()

  return None


def IsProcessHandlingSignal(pid, signum, status_path=None):
  """Checks whether a process is handling a signal.

  @type pid: int
  @param pid: Process ID
  @type signum: int
  @param signum: Signal number
  @rtype: bool

  """
  if status_path is None:
    status_path = _GetProcStatusPath(pid)

  try:
    proc_status = ReadFile(status_path)
  except EnvironmentError, err:
    # In at least one case, reading /proc/$pid/status failed with ESRCH.
    if err.errno in (errno.ENOENT, errno.ENOTDIR, errno.EINVAL, errno.ESRCH):
      return False
    raise

  sigcgt = _GetProcStatusField(proc_status, "SigCgt")
  if sigcgt is None:
    raise RuntimeError("%s is missing 'SigCgt' field" % status_path)

  # Now check whether signal is handled
  return signum in _ParseSigsetT(sigcgt)


1115
def ReadPidFile(pidfile):
1116
  """Read a pid from a file.
1117

1118
1119
1120
  @type  pidfile: string
  @param pidfile: path to the file containing the pid
  @rtype: int
1121
  @return: The process id, if the file exists and contains a valid PID,
1122
           otherwise 0
1123
1124
1125

  """
  try:
1126
    raw_data = ReadOneLineFile(pidfile)
1127
1128
  except EnvironmentError, err:
    if err.errno != errno.ENOENT:
1129
      logging.exception("Can't read pid file")
1130
    return 0
1131
1132

  try:
1133
    pid = int(raw_data)
1134
  except (TypeError, ValueError), err:
1135
    logging.info("Can't parse pid file contents", exc_info=True)
1136
    return 0
1137

1138
  return pid
1139
1140


1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
def ReadLockedPidFile(path):
  """Reads a locked PID file.

  This can be used together with L{StartDaemon}.

  @type path: string
  @param path: Path to PID file
  @return: PID as integer or, if file was unlocked or couldn't be opened, None

  """
  try:
    fd = os.open(path, os.O_RDONLY)
  except EnvironmentError, err:
    if err.errno == errno.ENOENT:
      # PID file doesn't exist
      return None
    raise

  try:
    try:
      # Try to acquire lock
      LockFile(fd)
    except errors.LockError:
      # Couldn't lock, daemon is running
      return int(os.read(fd, 100))
  finally:
    os.close(fd)

  return None


1172
def MatchNameComponent(key, name_list, case_sensitive=True):
Iustin Pop's avatar
Iustin Pop committed
1173
1174
1175
  """Try to match a name against a list.

  This function will try to match a name like test1 against a list
1176
1177
1178
1179
  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',
1180
1181
  'test1.example.org']}), except when the key fully matches an entry
  (e.g. I{'test1'} against C{['test1', 'test1.example.com']}).
Iustin Pop's avatar
Iustin Pop committed
1182

1183
1184
1185
1186
  @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
1187
1188
  @type case_sensitive: boolean
  @param case_sensitive: whether to provide a case-sensitive match
Iustin Pop's avatar
Iustin Pop committed
1189

1190
1191
1192
  @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
1193
1194

  """
1195
1196
  if key in name_list:
    return key
1197
1198
1199
1200

  re_flags = 0
  if not case_sensitive:
    re_flags |= re.IGNORECASE
1201
    key = key.upper()
1202
1203
1204
1205
1206
1207
  mo = re.compile("^%s(\..*)?$" % re.escape(key), re_flags)
  names_filtered = []
  string_matches = []
  for name in name_list:
    if mo.match(name) is not None:
      names_filtered.append(name)
1208
      if not case_sensitive and key == name.upper():
1209
1210
1211
1212
1213
1214
1215
        string_matches.append(name)

  if len(string_matches) == 1:
    return string_matches[0]
  if len(names_filtered) == 1:
    return names_filtered[0]
  return None
Iustin Pop's avatar
Iustin Pop committed
1216
1217


1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
def ValidateServiceName(name):
  """Validate the given service name.

  @type name: number or string
  @param name: Service name or port specification

  """
  try:
    numport = int(name)
  except (ValueError, TypeError):
    # Non-numeric service name
    valid = _VALID_SERVICE_NAME_RE.match(name)
  else:
    # Numeric port (protocols other than TCP or UDP might need adjustments
    # here)
    valid = (numport >= 0 and numport < (1 << 16))

  if not valid:
    raise errors.OpPrereqError("Invalid service name '%s'" % name,
                               errors.ECODE_INVAL)

  return name


Iustin Pop's avatar
Iustin Pop committed
1242
1243
1244
def ListVolumeGroups():
  """List volume groups and their size

1245
1246
1247
1248
  @rtype: dict
  @return:
       Dictionary with keys volume name and values
       the size of the volume
Iustin Pop's avatar
Iustin Pop committed
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261

  """
  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:
1262
      logging.error("Invalid output from vgs (%s): %s", err, line)
Iustin Pop's avatar
Iustin Pop committed
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
      continue

    retval[name] = size

  return retval


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

1273
1274
1275
1276
  @type bridge: str
  @param bridge: the bridge name to check
  @rtype: boolean
  @return: True if it does
Iustin Pop's avatar
Iustin Pop committed
1277
1278
1279
1280
1281
1282
1283
1284

  """
  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.

1285
1286
1287
  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
1288
1289
1290
1291
1292

  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.

1293
1294
1295
1296
  @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
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321

  """
  _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.

1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
  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
1332
1333
1334
1335

  """
  try:
    nv = fn(val)
Michael Hanselmann's avatar
Michael Hanselmann committed
1336
  except (ValueError, TypeError):
Iustin Pop's avatar
Iustin Pop committed
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
    nv = val
  return nv


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.

1351
1352
1353
1354
1355
  @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
1356
  """
1357
  return bool(_SHELLPARAM_REGEX.match(word))
Iustin Pop's avatar
Iustin Pop committed
1358
1359
1360
1361
1362
1363
1364


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
Michael Hanselmann's avatar
Michael Hanselmann committed
1365
  metacharacters). If everything is ok, it will return the result of
Iustin Pop's avatar
Iustin Pop committed
1366
1367
  template % args.

1368
1369
1370
1371
1372
1373
  @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
1374
1375
1376
  """
  for word in args:
    if not IsValidShellParam(word):
1377
1378
      raise errors.ProgrammerError("Shell argument '%s' contains"
                                   " invalid characters" % word)
Iustin Pop's avatar
Iustin Pop committed
1379
1380
1381
  return template % args


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

1385
1386
  @type value: int
  @param value: integer representing the value in MiB (1048576)
1387
1388
1389
1390
1391
1392
  @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
1393
1394
  @rtype: str
  @return: the formatted value (with suffix)
Iustin Pop's avatar
Iustin Pop committed
1395
1396

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

1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
  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
1411
1412

  else:
1413
1414
1415
    if units == 'h':
      suffix = 'T'
    return "%0.1f%s" % (round(float(value) / 1024 / 1024, 1), suffix)
Iustin Pop's avatar
Iustin Pop committed
1416
1417
1418
1419
1420


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

1421
1422
1423
  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
1424
1425

  """
1426
  m = _PARSEUNIT_REGEX.match(str(input_string))
Iustin Pop's avatar
Iustin Pop committed
1427
  if not m:
1428
    raise errors.UnitParseError("Invalid format")
Iustin Pop's avatar
Iustin Pop committed
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448

  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:
1449
    raise errors.UnitParseError("Unknown unit: %s" % unit)
Iustin Pop's avatar
Iustin Pop committed
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462

  # 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


1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
def ParseCpuMask(cpu_mask):
  """Parse a CPU mask definition and return the list of CPU IDs.

  CPU mask format: comma-separated list of CPU IDs
  or dash-separated ID ranges
  Example: "0-2,5" -> "0,1,2,5"

  @type cpu_mask: str
  @param cpu_mask: CPU mask definition
  @rtype: list of int
  @return: list of CPU IDs

  """
  if not cpu_mask:
    return []
  cpu_list = []
  for range_def in cpu_mask.split(","):
    boundaries = range_def.split("-")
    n_elements = len(boundaries)
    if n_elements > 2:
      raise errors.ParseError("Invalid CPU ID range definition"
                              " (only one hyphen allowed): %s" % range_def)
    try:
      lower = int(boundaries[0])
    except (ValueError, TypeError), err:
      raise errors.ParseError("Invalid CPU ID value for lower boundary of"
                              " CPU ID range: %s" % str(err))
    try:
      higher = int(boundaries[-1])
    except (ValueError, TypeError), err:
      raise errors.ParseError("Invalid CPU ID value for higher boundary of"
                              " CPU ID range: %s" % str(err))
    if lower > higher:
      raise errors.ParseError("Invalid CPU ID range definition"
                              " (%d > %d): %s" % (lower, higher, range_def))
    cpu_list.extend(range(lower, higher + 1))
  return cpu_list


1502
def AddAuthorizedKey(file_obj, key):
Iustin Pop's avatar
Iustin Pop committed
1503
1504
  """Adds an SSH public key to an authorized_keys file.

1505
1506
  @type file_obj: str or file handle
  @param file_obj: path to authorized_keys file
1507
1508
1509
  @type key: str
  @param key: string containing key

Iustin Pop's avatar
Iustin Pop committed
1510
1511
1512
  """
  key_fields = key.split()

1513
1514
1515
1516
1517
  if isinstance(file_obj, basestring):
    f = open(file_obj, 'a+')
  else:
    f = file_obj

Iustin Pop's avatar
Iustin Pop committed
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
  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.

1538
1539
1540
1541
1542
  @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
1543
1544
1545
1546
1547
  """
  key_fields = key.split()

  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
  try:
1548
    out = os.fdopen(fd, 'w')
Iustin Pop's avatar
Iustin Pop committed
1549
    try:
1550
1551
1552
1553
1554
1555
      f = open(file_name, 'r')
      try:
        for line in f:
          # Ignore whitespace changes while comparing lines
          if line.split() != key_fields:
            out.write(line)
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567

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


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

1571
1572
1573
1574
1575
1576
1577
1578
1579
  @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

1580
  """
1581
1582
1583
  # Ensure aliases are unique
  aliases = UniqueSequence([hostname] + aliases)[1:]

1584
1585
1586
1587
  def _WriteEtcHosts(fd):
    # Duplicating file descriptor because os.fdopen's result will automatically
    # close the descriptor, but we would still like to have its functionality.
    out = os.fdopen(os.dup(fd), "w")
1588
    try:
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
      for line in ReadFile(file_name).splitlines(True):
        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()
1600
1601
    finally:
      out.close()
1602
1603

  WriteFile(file_name, fn=_WriteEtcHosts, mode=0644)
1604
1605


1606
def AddHostToEtcHosts(hostname, ip):
Michael Hanselmann's avatar