utils.py 104 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
"""


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
import struct
import IN
54
55

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

Luca Bigliardi's avatar
Luca Bigliardi committed
57
58
59
60
61
try:
  import ctypes
except ImportError:
  ctypes = None

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

66

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

70
debug_locks = False
71
72

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

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

77
78
79
80
81
82
83
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)

84
85
86
87
88
89
90
91
92
93
94
95
# Structure definition for getsockopt(SOL_SOCKET, SO_PEERCRED, ...):
# struct ucred { pid_t pid; uid_t uid; gid_t gid; };
#
# The GNU C Library defines gid_t and uid_t to be "unsigned int" and
# pid_t to "int".
#
# IEEE Std 1003.1-2008:
# "nlink_t, uid_t, gid_t, and id_t shall be integer types"
# "blksize_t, pid_t, and ssize_t shall be signed integer types"
_STRUCT_UCRED = "iII"
_STRUCT_UCRED_SIZE = struct.calcsize(_STRUCT_UCRED)

96
97
98
99
# Certificate verification results
(CERT_WARNING,
 CERT_ERROR) = range(1, 3)

Luca Bigliardi's avatar
Luca Bigliardi committed
100
101
102
103
# Flags for mlockall() (from bits/mman.h)
_MCL_CURRENT = 1
_MCL_FUTURE = 2

104

Iustin Pop's avatar
Iustin Pop committed
105
class RunResult(object):
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
  """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
122
123
124
125
126
127

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


Iustin Pop's avatar
Iustin Pop committed
128
  def __init__(self, exit_code, signal_, stdout, stderr, cmd):
Iustin Pop's avatar
Iustin Pop committed
129
130
    self.cmd = cmd
    self.exit_code = exit_code
Iustin Pop's avatar
Iustin Pop committed
131
    self.signal = signal_
Iustin Pop's avatar
Iustin Pop committed
132
133
    self.stdout = stdout
    self.stderr = stderr
Iustin Pop's avatar
Iustin Pop committed
134
    self.failed = (signal_ is not None or exit_code != 0)
Iustin Pop's avatar
Iustin Pop committed
135
136
137
138
139
140
141
142

    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"

143
144
145
    if self.failed:
      logging.debug("Command '%s' failed (%s); output: %s",
                    self.cmd, self.fail_reason, self.output)
146

Iustin Pop's avatar
Iustin Pop committed
147
148
149
150
151
152
153
154
155
  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")


156
def _BuildCmdEnvironment(env, reset):
157
158
159
  """Builds the environment for an external program.

  """
160
161
162
163
164
165
  if reset:
    cmd_env = {}
  else:
    cmd_env = os.environ.copy()
    cmd_env["LC_ALL"] = "C"

166
167
  if env is not None:
    cmd_env.update(env)
168

169
170
171
  return cmd_env


172
def RunCmd(cmd, env=None, output=None, cwd="/", reset_env=False):
Iustin Pop's avatar
Iustin Pop committed
173
174
175
176
177
  """Execute a (shell) command.

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

178
  @type cmd: string or list
179
  @param cmd: Command to run
180
  @type env: dict
181
  @param env: Additional environment variables
182
  @type output: str
183
  @param output: if desired, the output of the command can be
184
185
      saved in a file instead of the RunResult instance; this
      parameter denotes the file name (if not None)
186
187
188
  @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
189
190
  @type reset_env: boolean
  @param reset_env: whether to reset or keep the default os environment
191
  @rtype: L{RunResult}
192
  @return: RunResult instance
Michael Hanselmann's avatar
Michael Hanselmann committed
193
  @raise errors.ProgrammerError: if we call this when forks are disabled
Iustin Pop's avatar
Iustin Pop committed
194
195

  """
196
197
198
  if no_fork:
    raise errors.ProgrammerError("utils.RunCmd() called with fork() disabled")

199
200
201
202
  if isinstance(cmd, basestring):
    strcmd = cmd
    shell = True
  else:
Iustin Pop's avatar
Iustin Pop committed
203
    cmd = [str(val) for val in cmd]
204
    strcmd = ShellQuoteArgs(cmd)
205
    shell = False
206
207
208

  if output:
    logging.debug("RunCmd %s, output file '%s'", strcmd, output)
209
  else:
210
    logging.debug("RunCmd %s", strcmd)
211

212
  cmd_env = _BuildCmdEnvironment(env, reset_env)
213

214
215
216
217
218
219
220
221
222
223
224
225
  try:
    if output is None:
      out, err, status = _RunCmdPipe(cmd, cmd_env, shell, cwd)
    else:
      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
226
227
228
229
230
231
232
233
234
235

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

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

236

237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
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)

276
  cmd_env = _BuildCmdEnvironment(env, False)
277
278
279
280
281
282
283
284
285
286
287
288
289
290

  # 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
291
292
293
294
              _StartDaemonChild(errpipe_read, errpipe_write,
                                pidpipe_read, pidpipe_write,
                                cmd, cmd_env, cwd,
                                output, output_fd, pidfile)
295
296
            finally:
              # Well, maybe child process failed
297
              os._exit(1) # pylint: disable-msg=W0212
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
        finally:
          _CloseFDNoErr(errpipe_write)

        # Wait for daemon to be started (or an error message to arrive) and read
        # up to 100 KB as an error message
        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))


331
332
333
334
def _StartDaemonChild(errpipe_read, errpipe_write,
                      pidpipe_read, pidpipe_write,
                      args, env, cwd,
                      output, fd_output, pidfile):
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
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
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
  """Child process for starting daemon.

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

    # First child process
    os.chdir("/")
    os.umask(077)
    os.setsid()

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

    # Make sure pipe is closed on execv* (and thereby notifies original process)
    SetCloseOnExecFlag(errpipe_write, True)

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

    # Open PID file
    if pidfile:
      try:
        # TODO: Atomic replace with another locked file instead of writing into
        # it after creating
        fd_pidfile = os.open(pidfile, os.O_WRONLY | os.O_CREAT, 0600)

        # Lock the PID file (and fail if not possible to do so). Any code
        # wanting to send a signal to the daemon should try to lock the PID
        # file before reading it. If acquiring the lock succeeds, the daemon is
        # no longer running and the signal should not be sent.
        LockFile(fd_pidfile)

        os.write(fd_pidfile, "%d\n" % os.getpid())
      except Exception, err:
        raise Exception("Creating and locking PID file failed: %s" % err)

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

      SetCloseOnExecFlag(fd_pidfile, False)
    else:
      fd_pidfile = None

    # Open /dev/null
    fd_devnull = os.open(os.devnull, os.O_RDWR)

    assert not output or (bool(output) ^ (fd_output is not None))

    if fd_output is not None:
      pass
    elif output:
      # Open output file
      try:
        # TODO: Implement flag to set append=yes/no
        fd_output = os.open(output, os.O_WRONLY | os.O_CREAT, 0600)
      except EnvironmentError, err:
        raise Exception("Opening output file failed: %s" % err)
    else:
      fd_output = fd_devnull

    # Redirect standard I/O
    os.dup2(fd_devnull, 0)
    os.dup2(fd_output, 1)
    os.dup2(fd_output, 2)

    # 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
      buf = str(sys.exc_info()[1])

      RetryOnSignal(os.write, errpipe_write, buf)
    except: # pylint: disable-msg=W0702
      # Ignore errors in error handling
      pass

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


432
def _RunCmdPipe(cmd, env, via_shell, cwd):
433
434
435
436
437
438
439
440
  """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
441
442
  @type cwd: string
  @param cwd: the working directory for the program
443
444
445
446
  @rtype: tuple
  @return: (out, err, status)

  """
447
  poller = select.poll()
448
  child = subprocess.Popen(cmd, shell=via_shell,
449
450
451
                           stderr=subprocess.PIPE,
                           stdout=subprocess.PIPE,
                           stdin=subprocess.PIPE,
452
453
                           close_fds=True, env=env,
                           cwd=cwd)
454
455

  child.stdin.close()
456
457
458
459
460
461
462
463
464
  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:
465
    SetNonblockFlag(fd, True)
466
467

  while fdmap:
468
    pollresult = RetryOnSignal(poller.poll)
469
470

    for fd, event in pollresult:
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
      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
486
487

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

490

491
def _RunCmdFile(cmd, env, via_shell, output, cwd):
492
493
494
495
496
497
498
499
500
501
  """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
502
503
  @type cwd: string
  @param cwd: the working directory for the program
504
505
506
507
508
509
510
511
512
513
  @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,
514
515
                             close_fds=True, env=env,
                             cwd=cwd)
516
517
518
519
520
521

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


524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
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)


543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
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)


562
563
564
565
566
567
568
def RetryOnSignal(fn, *args, **kwargs):
  """Calls a function again if it failed due to EINTR.

  """
  while True:
    try:
      return fn(*args, **kwargs)
569
    except EnvironmentError, err:
570
571
      if err.errno != errno.EINTR:
        raise
572
573
574
    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.
575
576
577
578
      if not (err.args and err.args[0] == errno.EINTR):
        raise


579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
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):
601
    fname = PathJoin(dir_name, relname)
602
603
604
605
606
607
608
609
610
611
612
613
614
615
    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


616
617
618
619
620
621
622
623
624
625
626
627
628
def GetSocketCredentials(sock):
  """Returns the credentials of the foreign process connected to a socket.

  @param sock: Unix socket
  @rtype: tuple; (number, number, number)
  @return: The PID, UID and GID of the connected foreign process.

  """
  peercred = sock.getsockopt(socket.SOL_SOCKET, IN.SO_PEERCRED,
                             _STRUCT_UCRED_SIZE)
  return struct.unpack(_STRUCT_UCRED, peercred)


Iustin Pop's avatar
Iustin Pop committed
629
630
631
632
633
634
def RemoveFile(filename):
  """Remove a file ignoring some errors.

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

635
636
637
  @type filename: str
  @param filename: the file to be removed

Iustin Pop's avatar
Iustin Pop committed
638
639
640
641
  """
  try:
    os.unlink(filename)
  except OSError, err:
642
    if err.errno not in (errno.ENOENT, errno.EISDIR):
Iustin Pop's avatar
Iustin Pop committed
643
644
      raise

Balazs Lecz's avatar
Balazs Lecz committed
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662

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
663

664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
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
685
      Makedirs(os.path.dirname(new), mode=mkdir_mode)
686

687
      return os.rename(old, new)
688

689
690
691
    raise


692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
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


708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
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
731
732
733
734
735
736
def _FingerprintFile(filename):
  """Compute the fingerprint of a file.

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

737
738
739
740
741
  @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
742
743
744
745
746
747
748

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

  f = open(filename)

749
  fp = compat.sha1_hash()
Iustin Pop's avatar
Iustin Pop committed
750
751
752
753
754
755
756
757
758
759
760
761
762
  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.

763
764
765
766
767
  @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
768
769
770
771
772
773
774
775
776
777
778
779

  """
  ret = {}

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

  return ret


780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
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 = []

795
796
797
798
  if not isinstance(target, dict):
    msg = "Expected dictionary, got '%s'" % target
    raise errors.TypeEnforcementError(msg)

799
800
801
802
803
804
805
806
  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
807
808
809
    ktype = key_types[key]
    if ktype not in constants.ENFORCEABLE_TYPES:
      msg = "'%s' has non-enforceable type %s" % (key, ktype)
810
811
      raise errors.ProgrammerError(msg)

Iustin Pop's avatar
Iustin Pop committed
812
    if ktype == constants.VTYPE_STRING:
813
814
815
816
817
818
      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)
Iustin Pop's avatar
Iustin Pop committed
819
    elif ktype == constants.VTYPE_BOOL:
820
821
822
823
824
825
826
827
828
829
830
831
      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
832
    elif ktype == constants.VTYPE_SIZE:
833
834
835
836
837
838
      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
839
    elif ktype == constants.VTYPE_INT:
840
841
842
843
844
845
846
      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)


847
848
849
850
851
852
853
854
855
856
857
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
858
859
860
def IsProcessAlive(pid):
  """Check if a given pid exists on the system.

861
862
  @note: zombie status is not handled, so zombie processes
      will be returned as alive
863
864
865
866
  @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
867
868

  """
869
870
871
872
873
874
875
876
877
878
879
880
  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"
881
882
883
  if pid <= 0:
    return False

884
885
  # /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
886
  try:
887
888
    return Retry(_TryStat, (0.01, 1.5, 0.1), 0.5,
                 args=[_GetProcStatusPath(pid)])
889
890
  except RetryTimeout, err:
    err.RaiseInner()
Iustin Pop's avatar
Iustin Pop committed
891
892


893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
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)


977
def ReadPidFile(pidfile):
978
  """Read a pid from a file.
979

980
981
982
  @type  pidfile: string
  @param pidfile: path to the file containing the pid
  @rtype: int
983
  @return: The process id, if the file exists and contains a valid PID,
984
           otherwise 0
985
986
987

  """
  try:
988
    raw_data = ReadOneLineFile(pidfile)
989
990
  except EnvironmentError, err:
    if err.errno != errno.ENOENT:
991
      logging.exception("Can't read pid file")
992
    return 0
993
994

  try:
995
    pid = int(raw_data)
996
  except (TypeError, ValueError), err:
997
    logging.info("Can't parse pid file contents", exc_info=True)
998
    return 0
999

1000
  return pid
1001
1002


1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
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


1034
def MatchNameComponent(key, name_list, case_sensitive=True):
Iustin Pop's avatar
Iustin Pop committed
1035
1036
1037
  """Try to match a name against a list.

  This function will try to match a name like test1 against a list
1038
1039
1040
1041
  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',
1042
1043
  '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
1044

1045
1046
1047
1048
  @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
1049
1050
  @type case_sensitive: boolean
  @param case_sensitive: whether to provide a case-sensitive match
Iustin Pop's avatar
Iustin Pop committed
1051

1052
1053
1054
  @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
1055
1056

  """
1057
1058
  if key in name_list:
    return key
1059
1060
1061
1062

  re_flags = 0
  if not case_sensitive:
    re_flags |= re.IGNORECASE
1063
    key = key.upper()
1064
1065
1066
1067
1068
1069
  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)
1070
      if not case_sensitive and key == name.upper():
1071
1072
1073
1074
1075
1076
1077
        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
1078
1079


1080
class HostInfo:
1081
  """Class implementing resolver and hostname functionality
1082
1083

  """
1084
1085
  _VALID_NAME_RE = re.compile("^[a-z0-9._-]{1,255}$")

1086
  def __init__(self, name=None):
1087
1088
    """Initialize the host name object.

1089
1090
    If the name argument is not passed, it will use this system's
    name.
1091
1092

    """
1093
1094
1095
1096
1097
    if name is None:
      name = self.SysName()

    self.query = name
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
1098
1099
    self.ip = self.ipaddrs[0]

1100
1101
1102
1103
1104
1105
  def ShortName(self):
    """Returns the hostname without domain.

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

1106
1107
1108
  @staticmethod
  def SysName():
    """Return the current system's name.
1109

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

1112
1113
    """
    return socket.gethostname()
Iustin Pop's avatar
Iustin Pop committed
1114

1115
1116
1117
  @staticmethod
  def LookupHostname(hostname):
    """Look up hostname
Iustin Pop's avatar
Iustin Pop committed
1118

1119
1120
    @type hostname: str
    @param hostname: hostname to look up
1121

1122
1123
1124
1125
    @rtype: tuple
    @return: a tuple (name, aliases, ipaddrs) as returned by
        C{socket.gethostbyname_ex}
    @raise errors.ResolverError: in case of errors in resolving
1126
1127
1128
1129
1130
1131
1132

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

1134
    return result
Iustin Pop's avatar
Iustin Pop committed
1135

1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
  @classmethod
  def NormalizeName(cls, hostname):
    """Validate and normalize the given hostname.

    @attention: the validation is a bit more relaxed than the standards
        require; most importantly, we allow underscores in names
    @raise errors.OpPrereqError: when the name is not valid

    """
    hostname = hostname.lower()
    if (not cls._VALID_NAME_RE.match(hostname) or
        # double-dots, meaning empty label
        ".." in hostname or
        # empty initial label
        hostname.startswith(".")):
      raise errors.OpPrereqError("Invalid hostname '%s'" % hostname,
                                 errors.ECODE_INVAL)
    if hostname.endswith("."):
      hostname = hostname.rstrip(".")
    return hostname

Iustin Pop's avatar
Iustin Pop committed
1157

1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
def GetHostInfo(name=None):
  """Lookup host name and raise an OpPrereqError for failures"""

  try:
    return HostInfo(name)
  except errors.ResolverError, err:
    raise errors.OpPrereqError("The given name (%s) does not resolve: %s" %
                               (err[0], err[2]), errors.ECODE_RESOLVER)


Iustin Pop's avatar
Iustin Pop committed
1168
1169
1170
def ListVolumeGroups():
  """List volume groups and their size

1171
1172
1173
1174
  @rtype: dict
  @return:
       Dictionary with keys volume name and values
       the size of the volume
Iustin Pop's avatar
Iustin Pop committed
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187

  """
  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:
1188
      logging.error("Invalid output from vgs (%s): %s", err, line)
Iustin Pop's avatar
Iustin Pop committed
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
      continue

    retval[name] = size

  return retval


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

1199
1200
1201
1202
  @type bridge: str
  @param bridge: the bridge name to check
  @rtype: boolean
  @return: True if it does
Iustin Pop's avatar
Iustin Pop committed
1203
1204
1205
1206
1207
1208
1209
1210

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

1211
1212
1213
  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
1214
1215
1216
1217
1218

  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.

1219
1220
1221
1222
  @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
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247

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

1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
  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
1258
1259
1260
1261

  """
  try:
    nv = fn(val)
Michael Hanselmann's avatar
Michael Hanselmann committed
1262
  except (ValueError, TypeError):
Iustin Pop's avatar
Iustin Pop committed
1263
1264
1265
1266
1267
    nv = val
  return nv


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

1270
1271
1272
1273
1274
1275
  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
Michael Hanselmann's avatar
Michael Hanselmann committed
1276
  @return: a regular expression match object, or None if the
1277
      address is not valid
Iustin Pop's avatar
Iustin Pop committed
1278
1279
1280

  """
  unit = "(0|[1-9]\d{0,2})"
1281
  #TODO: convert and return only boolean
Iustin Pop's avatar
Iustin Pop committed
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
  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.

1295
1296
1297
1298
1299
  @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
1300
1301
1302
1303
1304
1305
1306
1307
1308
  """
  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
Michael Hanselmann's avatar
Michael Hanselmann committed
1309
  metacharacters). If everything is ok, it will return the result of
Iustin Pop's avatar
Iustin Pop committed
1310
1311
  template % args.

1312
1313
1314
1315
1316
1317
  @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
1318
1319
1320
  """
  for word in args:
    if not IsValidShellParam(word):
1321
1322
      raise errors.ProgrammerError("Shell argument '%s' contains"
                                   " invalid characters" % word)
Iustin Pop's avatar
Iustin Pop committed
1323
1324
1325
  return template % args


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

1329
1330
  @type value: int
  @param value: integer representing the value in MiB (1048576)
1331
1332
1333
1334
1335
1336
  @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
1337
1338
  @rtype: str
  @return: the formatted value (with suffix)
Iustin Pop's avatar
Iustin Pop committed
1339
1340

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

1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
  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
1355
1356

  else:
1357
1358
1359
    if units == 'h':
      suffix = 'T'
    return "%0.1f%s" % (round(float(value) / 1024 / 1024, 1), suffix)
Iustin Pop's avatar
Iustin Pop committed
1360
1361
1362
1363
1364


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

1365
1366
1367
  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
1368
1369

  """
1370
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', str(input_string))
Iustin Pop's avatar
Iustin Pop committed
1371
  if not m:
1372
    raise errors.UnitParseError("Invalid format")
Iustin Pop's avatar
Iustin Pop committed
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392

  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:
1393
    raise errors.UnitParseError("Unknown unit: %s" % unit)
Iustin Pop's avatar
Iustin Pop committed
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409

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

1410
1411
1412
1413
1414
  @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
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
  """
  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.

1439
1440
1441
1442
1443
  @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
1444
1445
1446
1447
1448
  """
  key_fields = key.split()

  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
  try:
1449
    out = os.fdopen(fd, 'w')
Iustin Pop's avatar
Iustin Pop committed
1450
    try:
1451
1452
1453
1454
1455
1456
      f = open(file_name, 'r')
      try:
        for line in f:
          # Ignore whitespace changes while comparing lines
          if line.split() != key_fields:
            out.write(line)
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468

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


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

1472
1473
1474
1475
1476
1477
1478
1479
1480
  @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

1481
  """
1482
  # FIXME: use WriteFile + fn rather than duplicating its efforts
1483
1484
1485
  # Ensure aliases are unique
  aliases = UniqueSequence([hostname] + aliases)[1:]

1486
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
1487
  try:
1488
1489
1490
1491
1492
1493
    out = os.fdopen(fd, 'w')
    try:
      f = open(file_name, 'r')
      try:
        for line in f:
          fields = line.split()
1494
          if fields and not fields[0].startswith('#') and ip == fields[0]:
1495
1496
1497
            continue
          out.write(line)

1498
        out.write("%s\t%s" % (ip, hostname))
1499
1500
1501
1502
1503
        if aliases:
          out.write(" %s" % ' '.join(aliases))
        out.write('\n')

        out.flush()
1504
        os.fsync(out)
1505
        os.chmod(tmpname, 0644)
1506
1507
1508
1509
1510
1511
1512
1513
        os.rename(tmpname, file_name)
      finally:
        f.close()
    finally:
      out.close()
  except:
    RemoveFile(tmpname)
    raise
1514
1515


1516
1517
1518
def AddHostToEtcHosts(hostname):
  """Wrapper around SetEtcHostsEntry.

1519
1520
1521
1522
  @type hostname: str
  @param hostname: a hostname that will be resolved and added to
      L{constants.ETC_HOSTS}

1523
1524
1525
1526
1527
  """
  hi = HostInfo(name=hostname)
  SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])


1528
def RemoveEtcHostsEntry(file_name, hostname):
1529
  """Removes a hostname from /etc/hosts.
1530

1531
  IP addresses without names are removed from the file.
1532
1533
1534
1535
1536
1537

  @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

1538
  """
1539
  # FIXME: use WriteFile + fn rather than duplicating its efforts
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
  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:
1554
                out.write("%s %s\n" % (fields[0], ' '.join(names)))