Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ b0fc8c89

History | View | Annotate | Download (30.7 kB)

1
#
2
#
3

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

    
21

    
22
"""Ganeti small utilities
23

24
"""
25

    
26

    
27
import sys
28
import os
29
import sha
30
import time
31
import subprocess
32
import re
33
import socket
34
import tempfile
35
import shutil
36
import errno
37
import pwd
38
import itertools
39
import select
40
import fcntl
41
import resource
42

    
43
from cStringIO import StringIO
44

    
45
from ganeti import logger
46
from ganeti import errors
47
from ganeti import constants
48

    
49

    
50
_locksheld = []
51
_re_shell_unquoted = re.compile('^[-.,=:/_+@A-Za-z0-9]+$')
52

    
53
debug = False
54

    
55
class RunResult(object):
56
  """Simple class for holding the result of running external programs.
57

58
  Instance variables:
59
    exit_code: the exit code of the program, or None (if the program
60
               didn't exit())
61
    signal: numeric signal that caused the program to finish, or None
62
            (if the program wasn't terminated by a signal)
63
    stdout: the standard output of the program
64
    stderr: the standard error of the program
65
    failed: a Boolean value which is True in case the program was
66
            terminated by a signal or exited with a non-zero exit code
67
    fail_reason: a string detailing the termination reason
68

69
  """
70
  __slots__ = ["exit_code", "signal", "stdout", "stderr",
71
               "failed", "fail_reason", "cmd"]
72

    
73

    
74
  def __init__(self, exit_code, signal, stdout, stderr, cmd):
75
    self.cmd = cmd
76
    self.exit_code = exit_code
77
    self.signal = signal
78
    self.stdout = stdout
79
    self.stderr = stderr
80
    self.failed = (signal is not None or exit_code != 0)
81

    
82
    if self.signal is not None:
83
      self.fail_reason = "terminated by signal %s" % self.signal
84
    elif self.exit_code is not None:
85
      self.fail_reason = "exited with exit code %s" % self.exit_code
86
    else:
87
      self.fail_reason = "unable to determine termination reason"
88

    
89
    if debug and self.failed:
90
      logger.Debug("Command '%s' failed (%s); output: %s" %
91
                   (self.cmd, self.fail_reason, self.output))
92

    
93
  def _GetOutput(self):
94
    """Returns the combined stdout and stderr for easier usage.
95

96
    """
97
    return self.stdout + self.stderr
98

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

    
101

    
102
def _GetLockFile(subsystem):
103
  """Compute the file name for a given lock name."""
104
  return "%s/ganeti_lock_%s" % (constants.LOCK_DIR, subsystem)
105

    
106

    
107
def Lock(name, max_retries=None, debug=False, autoclean=True):
108
  """Lock a given subsystem.
109

110
  In case the lock is already held by an alive process, the function
111
  will sleep indefintely and poll with a one second interval.
112

113
  When the optional integer argument 'max_retries' is passed with a
114
  non-zero value, the function will sleep only for this number of
115
  times, and then it will will raise a LockError if the lock can't be
116
  acquired. Passing in a negative number will cause only one try to
117
  get the lock. Passing a positive number will make the function retry
118
  for approximately that number of seconds.
119

120
  """
121
  lockfile = _GetLockFile(name)
122

    
123
  if name in _locksheld:
124
    raise errors.LockError('Lock "%s" already held!' % (name,))
125

    
126
  errcount = 0
127
  cleanupcount = 0
128

    
129
  retries = 0
130
  while True:
131
    try:
132
      fd = os.open(lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_SYNC)
133
      break
134
    except OSError, creat_err:
135
      if creat_err.errno != errno.EEXIST:
136
        raise errors.LockError("Can't create the lock file. Error '%s'." %
137
                               str(creat_err))
138

    
139
      try:
140
        pf = open(lockfile, 'r')
141
      except IOError, open_err:
142
        errcount += 1
143
        if errcount >= 5:
144
          raise errors.LockError("Lock file exists but cannot be opened."
145
                                 " Error: '%s'." % str(open_err))
146
        time.sleep(1)
147
        continue
148

    
149
      try:
150
        pid = int(pf.read())
151
      except ValueError:
152
        raise errors.LockError("Invalid pid string in %s" %
153
                               (lockfile,))
154

    
155
      if not IsProcessAlive(pid):
156
        if autoclean:
157
          cleanupcount += 1
158
          if cleanupcount >= 5:
159
            raise errors.LockError, ("Too many stale lock cleanups! Check"
160
                                     " what process is dying.")
161
          logger.Error('Stale lockfile %s for pid %d, autocleaned.' %
162
                       (lockfile, pid))
163
          RemoveFile(lockfile)
164
          continue
165
        else:
166
          raise errors.LockError("Stale lockfile %s for pid %d?" %
167
                                 (lockfile, pid))
168

    
169
      if max_retries and max_retries <= retries:
170
        raise errors.LockError("Can't acquire lock during the specified"
171
                               " time, aborting.")
172
      if retries == 5 and (debug or sys.stdin.isatty()):
173
        logger.ToStderr("Waiting for '%s' lock from pid %d..." % (name, pid))
174

    
175
      time.sleep(1)
176
      retries += 1
177
      continue
178

    
179
  os.write(fd, '%d\n' % (os.getpid(),))
180
  os.close(fd)
181

    
182
  _locksheld.append(name)
183

    
184

    
185
def Unlock(name):
186
  """Unlock a given subsystem.
187

188
  """
189
  lockfile = _GetLockFile(name)
190

    
191
  try:
192
    fd = os.open(lockfile, os.O_RDONLY)
193
  except OSError:
194
    raise errors.LockError('Lock "%s" not held.' % (name,))
195

    
196
  f = os.fdopen(fd, 'r')
197
  pid_str = f.read()
198

    
199
  try:
200
    pid = int(pid_str)
201
  except ValueError:
202
    raise errors.LockError('Unable to determine PID of locking process.')
203

    
204
  if pid != os.getpid():
205
    raise errors.LockError('Lock not held by me (%d != %d)' %
206
                           (os.getpid(), pid,))
207

    
208
  os.unlink(lockfile)
209
  _locksheld.remove(name)
210

    
211

    
212
def LockCleanup():
213
  """Remove all locks.
214

215
  """
216
  for lock in _locksheld:
217
    Unlock(lock)
218

    
219

    
220
def RunCmd(cmd):
221
  """Execute a (shell) command.
222

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

226
  @param cmd: Command to run
227
  @type  cmd: string or list
228
  @return: `RunResult` instance
229
  @rtype: RunResult
230

231
  """
232
  if isinstance(cmd, list):
233
    cmd = [str(val) for val in cmd]
234
    strcmd = " ".join(cmd)
235
    shell = False
236
  else:
237
    strcmd = cmd
238
    shell = True
239
  env = os.environ.copy()
240
  env["LC_ALL"] = "C"
241
  poller = select.poll()
242
  child = subprocess.Popen(cmd, shell=shell,
243
                           stderr=subprocess.PIPE,
244
                           stdout=subprocess.PIPE,
245
                           stdin=subprocess.PIPE,
246
                           close_fds=True, env=env)
247

    
248
  child.stdin.close()
249
  poller.register(child.stdout, select.POLLIN)
250
  poller.register(child.stderr, select.POLLIN)
251
  out = StringIO()
252
  err = StringIO()
253
  fdmap = {
254
    child.stdout.fileno(): (out, child.stdout),
255
    child.stderr.fileno(): (err, child.stderr),
256
    }
257
  for fd in fdmap:
258
    status = fcntl.fcntl(fd, fcntl.F_GETFL)
259
    fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
260

    
261
  while fdmap:
262
    for fd, event in poller.poll():
263
      if event & select.POLLIN or event & select.POLLPRI:
264
        data = fdmap[fd][1].read()
265
        # no data from read signifies EOF (the same as POLLHUP)
266
        if not data:
267
          poller.unregister(fd)
268
          del fdmap[fd]
269
          continue
270
        fdmap[fd][0].write(data)
271
      if (event & select.POLLNVAL or event & select.POLLHUP or
272
          event & select.POLLERR):
273
        poller.unregister(fd)
274
        del fdmap[fd]
275

    
276
  out = out.getvalue()
277
  err = err.getvalue()
278

    
279
  status = child.wait()
280
  if status >= 0:
281
    exitcode = status
282
    signal = None
283
  else:
284
    exitcode = None
285
    signal = -status
286

    
287
  return RunResult(exitcode, signal, out, err, strcmd)
288

    
289

    
290
def RunCmdUnlocked(cmd):
291
  """Execute a shell command without the 'cmd' lock.
292

293
  This variant of `RunCmd()` drops the 'cmd' lock before running the
294
  command and re-aquires it afterwards, thus it can be used to call
295
  other ganeti commands.
296

297
  The argument and return values are the same as for the `RunCmd()`
298
  function.
299

300
  Args:
301
    cmd - command to run. (str)
302

303
  Returns:
304
    `RunResult`
305

306
  """
307
  Unlock('cmd')
308
  ret = RunCmd(cmd)
309
  Lock('cmd')
310

    
311
  return ret
312

    
313

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

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

320
  """
321
  try:
322
    os.unlink(filename)
323
  except OSError, err:
324
    if err.errno not in (errno.ENOENT, errno.EISDIR):
325
      raise
326

    
327

    
328
def _FingerprintFile(filename):
329
  """Compute the fingerprint of a file.
330

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

334
  Args:
335
    filename - Filename (str)
336

337
  """
338
  if not (os.path.exists(filename) and os.path.isfile(filename)):
339
    return None
340

    
341
  f = open(filename)
342

    
343
  fp = sha.sha()
344
  while True:
345
    data = f.read(4096)
346
    if not data:
347
      break
348

    
349
    fp.update(data)
350

    
351
  return fp.hexdigest()
352

    
353

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

357
  Args:
358
    files - array of filenames.  ( [str, ...] )
359

360
  Return value:
361
    dictionary of filename: fingerprint for the files that exist
362

363
  """
364
  ret = {}
365

    
366
  for filename in files:
367
    cksum = _FingerprintFile(filename)
368
    if cksum:
369
      ret[filename] = cksum
370

    
371
  return ret
372

    
373

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

377
  For the given dictionaries `target` and `template`, ensure target
378
  has all the keys from template. Missing keys are added with values
379
  from template.
380

381
  Args:
382
    target   - the dictionary to check
383
    template - template dictionary
384
    logname  - a caller-chosen string to identify the debug log
385
               entry; if None, no logging will be done
386

387
  Returns value:
388
    None
389

390
  """
391
  missing = []
392
  for k in template:
393
    if k not in target:
394
      missing.append(k)
395
      target[k] = template[k]
396

    
397
  if missing and logname:
398
    logger.Debug('%s missing keys %s' %
399
                 (logname, ', '.join(missing)))
400

    
401

    
402
def IsProcessAlive(pid):
403
  """Check if a given pid exists on the system.
404

405
  Returns: true or false, depending on if the pid exists or not
406

407
  Remarks: zombie processes treated as not alive
408

409
  """
410
  try:
411
    f = open("/proc/%d/status" % pid)
412
  except IOError, err:
413
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
414
      return False
415

    
416
  alive = True
417
  try:
418
    data = f.readlines()
419
    if len(data) > 1:
420
      state = data[1].split()
421
      if len(state) > 1 and state[1] == "Z":
422
        alive = False
423
  finally:
424
    f.close()
425

    
426
  return alive
427

    
428

    
429
def MatchNameComponent(key, name_list):
430
  """Try to match a name against a list.
431

432
  This function will try to match a name like test1 against a list
433
  like ['test1.example.com', 'test2.example.com', ...]. Against this
434
  list, 'test1' as well as 'test1.example' will match, but not
435
  'test1.ex'. A multiple match will be considered as no match at all
436
  (e.g. 'test1' against ['test1.example.com', 'test1.example.org']).
437

438
  Args:
439
    key: the name to be searched
440
    name_list: the list of strings against which to search the key
441

442
  Returns:
443
    None if there is no match *or* if there are multiple matches
444
    otherwise the element from the list which matches
445

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

    
453

    
454
class HostInfo:
455
  """Class implementing resolver and hostname functionality
456

457
  """
458
  def __init__(self, name=None):
459
    """Initialize the host name object.
460

461
    If the name argument is not passed, it will use this system's
462
    name.
463

464
    """
465
    if name is None:
466
      name = self.SysName()
467

    
468
    self.query = name
469
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
470
    self.ip = self.ipaddrs[0]
471

    
472
  def ShortName(self):
473
    """Returns the hostname without domain.
474

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

    
478
  @staticmethod
479
  def SysName():
480
    """Return the current system's name.
481

482
    This is simply a wrapper over socket.gethostname()
483

484
    """
485
    return socket.gethostname()
486

    
487
  @staticmethod
488
  def LookupHostname(hostname):
489
    """Look up hostname
490

491
    Args:
492
      hostname: hostname to look up
493

494
    Returns:
495
      a tuple (name, aliases, ipaddrs) as returned by socket.gethostbyname_ex
496
      in case of errors in resolving, we raise a ResolverError
497

498
    """
499
    try:
500
      result = socket.gethostbyname_ex(hostname)
501
    except socket.gaierror, err:
502
      # hostname not found in DNS
503
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
504

    
505
    return result
506

    
507

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

511
  Returns:
512
     Dictionary with keys volume name and values the size of the volume
513

514
  """
515
  command = "vgs --noheadings --units m --nosuffix -o name,size"
516
  result = RunCmd(command)
517
  retval = {}
518
  if result.failed:
519
    return retval
520

    
521
  for line in result.stdout.splitlines():
522
    try:
523
      name, size = line.split()
524
      size = int(float(size))
525
    except (IndexError, ValueError), err:
526
      logger.Error("Invalid output from vgs (%s): %s" % (err, line))
527
      continue
528

    
529
    retval[name] = size
530

    
531
  return retval
532

    
533

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

537
  Returns:
538
     True if it does, false otherwise.
539

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

    
543

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

547
  Given a list of names ['a1', 'a10', 'a11', 'a2'] this function will
548
  sort the list in the logical order ['a1', 'a2', 'a10', 'a11'].
549

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

554
  Return value
555
    - a copy of the list sorted according to our algorithm
556

557
  """
558
  _SORTER_BASE = "(\D+|\d+)"
559
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
560
                                                  _SORTER_BASE, _SORTER_BASE,
561
                                                  _SORTER_BASE, _SORTER_BASE,
562
                                                  _SORTER_BASE, _SORTER_BASE)
563
  _SORTER_RE = re.compile(_SORTER_FULL)
564
  _SORTER_NODIGIT = re.compile("^\D*$")
565
  def _TryInt(val):
566
    """Attempts to convert a variable to integer."""
567
    if val is None or _SORTER_NODIGIT.match(val):
568
      return val
569
    rval = int(val)
570
    return rval
571

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

    
577

    
578
def CheckDaemonAlive(pid_file, process_string):
579
  """Check wether the specified daemon is alive.
580

581
  Args:
582
   - pid_file: file to read the daemon pid from, the file is
583
               expected to contain only a single line containing
584
               only the PID
585
   - process_string: a substring that we expect to find in
586
                     the command line of the daemon process
587

588
  Returns:
589
   - True if the daemon is judged to be alive (that is:
590
      - the PID file exists, is readable and contains a number
591
      - a process of the specified PID is running
592
      - that process contains the specified string in its
593
        command line
594
      - the process is not in state Z (zombie))
595
   - False otherwise
596

597
  """
598
  try:
599
    pid_file = file(pid_file, 'r')
600
    try:
601
      pid = int(pid_file.readline())
602
    finally:
603
      pid_file.close()
604

    
605
    cmdline_file_path = "/proc/%s/cmdline" % (pid)
606
    cmdline_file = open(cmdline_file_path, 'r')
607
    try:
608
      cmdline = cmdline_file.readline()
609
    finally:
610
      cmdline_file.close()
611

    
612
    if not process_string in cmdline:
613
      return False
614

    
615
    stat_file_path =  "/proc/%s/stat" % (pid)
616
    stat_file = open(stat_file_path, 'r')
617
    try:
618
      process_state = stat_file.readline().split()[2]
619
    finally:
620
      stat_file.close()
621

    
622
    if process_state == 'Z':
623
      return False
624

    
625
  except (IndexError, IOError, ValueError):
626
    return False
627

    
628
  return True
629

    
630

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

634
  This function tries to apply function `fn` to `val`. If no
635
  ValueError or TypeError exceptions are raised, it will return the
636
  result, else it will return the original value. Any other exceptions
637
  are propagated to the caller.
638

639
  """
640
  try:
641
    nv = fn(val)
642
  except (ValueError, TypeError), err:
643
    nv = val
644
  return nv
645

    
646

    
647
def IsValidIP(ip):
648
  """Verifies the syntax of an IP address.
649

650
  This function checks if the ip address passes is valid or not based
651
  on syntax (not ip range, class calculations or anything).
652

653
  """
654
  unit = "(0|[1-9]\d{0,2})"
655
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
656

    
657

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

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

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

668
  """
669
  return bool(re.match("^[-a-zA-Z0-9._+/:%@]+$", word))
670

    
671

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

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

680
  """
681
  for word in args:
682
    if not IsValidShellParam(word):
683
      raise errors.ProgrammerError("Shell argument '%s' contains"
684
                                   " invalid characters" % word)
685
  return template % args
686

    
687

    
688
def FormatUnit(value):
689
  """Formats an incoming number of MiB with the appropriate unit.
690

691
  Value needs to be passed as a numeric type. Return value is always a string.
692

693
  """
694
  if value < 1024:
695
    return "%dM" % round(value, 0)
696

    
697
  elif value < (1024 * 1024):
698
    return "%0.1fG" % round(float(value) / 1024, 1)
699

    
700
  else:
701
    return "%0.1fT" % round(float(value) / 1024 / 1024, 1)
702

    
703

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

707
  Input must be in the format NUMBER+ [DOT NUMBER+] SPACE* [UNIT]. If no unit
708
  is specified, it defaults to MiB. Return value is always an int in MiB.
709

710
  """
711
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', input_string)
712
  if not m:
713
    raise errors.UnitParseError("Invalid format")
714

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

    
717
  unit = m.groups()[1]
718
  if unit:
719
    lcunit = unit.lower()
720
  else:
721
    lcunit = 'm'
722

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

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

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

    
733
  else:
734
    raise errors.UnitParseError("Unknown unit: %s" % unit)
735

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

    
740
  # Round up to the next multiple of 4
741
  value = int(value)
742
  if value % 4:
743
    value += 4 - value % 4
744

    
745
  return value
746

    
747

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

751
  Args:
752
    file_name: Path to authorized_keys file
753
    key: String containing key
754
  """
755
  key_fields = key.split()
756

    
757
  f = open(file_name, 'a+')
758
  try:
759
    nl = True
760
    for line in f:
761
      # Ignore whitespace changes
762
      if line.split() == key_fields:
763
        break
764
      nl = line.endswith('\n')
765
    else:
766
      if not nl:
767
        f.write("\n")
768
      f.write(key.rstrip('\r\n'))
769
      f.write("\n")
770
      f.flush()
771
  finally:
772
    f.close()
773

    
774

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

778
  Args:
779
    file_name: Path to authorized_keys file
780
    key: String containing key
781
  """
782
  key_fields = key.split()
783

    
784
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
785
  try:
786
    out = os.fdopen(fd, 'w')
787
    try:
788
      f = open(file_name, 'r')
789
      try:
790
        for line in f:
791
          # Ignore whitespace changes while comparing lines
792
          if line.split() != key_fields:
793
            out.write(line)
794

    
795
        out.flush()
796
        os.rename(tmpname, file_name)
797
      finally:
798
        f.close()
799
    finally:
800
      out.close()
801
  except:
802
    RemoveFile(tmpname)
803
    raise
804

    
805

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

809
  """
810
  # Ensure aliases are unique
811
  aliases = UniqueSequence([hostname] + aliases)[1:]
812

    
813
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
814
  try:
815
    out = os.fdopen(fd, 'w')
816
    try:
817
      f = open(file_name, 'r')
818
      try:
819
        written = False
820
        for line in f:
821
          fields = line.split()
822
          if fields and not fields[0].startswith('#') and ip == fields[0]:
823
            continue
824
          out.write(line)
825

    
826
        out.write("%s\t%s" % (ip, hostname))
827
        if aliases:
828
          out.write(" %s" % ' '.join(aliases))
829
        out.write('\n')
830

    
831
        out.flush()
832
        os.fsync(out)
833
        os.rename(tmpname, file_name)
834
      finally:
835
        f.close()
836
    finally:
837
      out.close()
838
  except:
839
    RemoveFile(tmpname)
840
    raise
841

    
842

    
843
def RemoveEtcHostsEntry(file_name, hostname):
844
  """Removes a hostname from /etc/hosts.
845

846
  IP addresses without names are removed from the file.
847
  """
848
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
849
  try:
850
    out = os.fdopen(fd, 'w')
851
    try:
852
      f = open(file_name, 'r')
853
      try:
854
        for line in f:
855
          fields = line.split()
856
          if len(fields) > 1 and not fields[0].startswith('#'):
857
            names = fields[1:]
858
            if hostname in names:
859
              while hostname in names:
860
                names.remove(hostname)
861
              if names:
862
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
863
              continue
864

    
865
          out.write(line)
866

    
867
        out.flush()
868
        os.fsync(out)
869
        os.rename(tmpname, file_name)
870
      finally:
871
        f.close()
872
    finally:
873
      out.close()
874
  except:
875
    RemoveFile(tmpname)
876
    raise
877

    
878

    
879
def CreateBackup(file_name):
880
  """Creates a backup of a file.
881

882
  Returns: the path to the newly created backup file.
883

884
  """
885
  if not os.path.isfile(file_name):
886
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
887
                                file_name)
888

    
889
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
890
  dir_name = os.path.dirname(file_name)
891

    
892
  fsrc = open(file_name, 'rb')
893
  try:
894
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
895
    fdst = os.fdopen(fd, 'wb')
896
    try:
897
      shutil.copyfileobj(fsrc, fdst)
898
    finally:
899
      fdst.close()
900
  finally:
901
    fsrc.close()
902

    
903
  return backup_name
904

    
905

    
906
def ShellQuote(value):
907
  """Quotes shell argument according to POSIX.
908

909
  """
910
  if _re_shell_unquoted.match(value):
911
    return value
912
  else:
913
    return "'%s'" % value.replace("'", "'\\''")
914

    
915

    
916
def ShellQuoteArgs(args):
917
  """Quotes all given shell arguments and concatenates using spaces.
918

919
  """
920
  return ' '.join([ShellQuote(i) for i in args])
921

    
922

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

926
  Try to do a TCP connect(2) from an optional source IP to the
927
  specified target IP and the specified target port. If the optional
928
  parameter live_port_needed is set to true, requires the remote end
929
  to accept the connection. The timeout is specified in seconds and
930
  defaults to 10 seconds. If the source optional argument is not
931
  passed, the source address selection is left to the kernel,
932
  otherwise we try to connect using the passed address (failures to
933
  bind other than EADDRNOTAVAIL will be ignored).
934

935
  """
936
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
937

    
938
  sucess = False
939

    
940
  if source is not None:
941
    try:
942
      sock.bind((source, 0))
943
    except socket.error, (errcode, errstring):
944
      if errcode == errno.EADDRNOTAVAIL:
945
        success = False
946

    
947
  sock.settimeout(timeout)
948

    
949
  try:
950
    sock.connect((target, port))
951
    sock.close()
952
    success = True
953
  except socket.timeout:
954
    success = False
955
  except socket.error, (errcode, errstring):
956
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
957

    
958
  return success
959

    
960

    
961
def ListVisibleFiles(path):
962
  """Returns a list of all visible files in a directory.
963

964
  """
965
  files = [i for i in os.listdir(path) if not i.startswith(".")]
966
  files.sort()
967
  return files
968

    
969

    
970
def GetHomeDir(user, default=None):
971
  """Try to get the homedir of the given user.
972

973
  The user can be passed either as a string (denoting the name) or as
974
  an integer (denoting the user id). If the user is not found, the
975
  'default' argument is returned, which defaults to None.
976

977
  """
978
  try:
979
    if isinstance(user, basestring):
980
      result = pwd.getpwnam(user)
981
    elif isinstance(user, (int, long)):
982
      result = pwd.getpwuid(user)
983
    else:
984
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
985
                                   type(user))
986
  except KeyError:
987
    return default
988
  return result.pw_dir
989

    
990

    
991
def NewUUID():
992
  """Returns a random UUID.
993

994
  """
995
  f = open("/proc/sys/kernel/random/uuid", "r")
996
  try:
997
    return f.read(128).rstrip("\n")
998
  finally:
999
    f.close()
1000

    
1001

    
1002
def WriteFile(file_name, fn=None, data=None,
1003
              mode=None, uid=-1, gid=-1,
1004
              atime=None, mtime=None, close=True,
1005
              prewrite=None, postwrite=None):
1006
  """(Over)write a file atomically.
1007

1008
  The file_name and either fn (a function taking one argument, the
1009
  file descriptor, and which should write the data to it) or data (the
1010
  contents of the file) must be passed. The other arguments are
1011
  optional and allow setting the file mode, owner and group, and the
1012
  mtime/atime of the file.
1013

1014
  If the function doesn't raise an exception, it has succeeded and the
1015
  target file has the new contents. If the file has raised an
1016
  exception, an existing target file should be unmodified and the
1017
  temporary file should be removed.
1018

1019
  Args:
1020
    file_name: New filename
1021
    fn: Content writing function, called with file descriptor as parameter
1022
    data: Content as string
1023
    mode: File mode
1024
    uid: Owner
1025
    gid: Group
1026
    atime: Access time
1027
    mtime: Modification time
1028
    close: Whether to close file after writing it
1029
    prewrite: Function object called before writing content
1030
    postwrite: Function object called after writing content
1031

1032
  Returns:
1033
    None if "close" parameter evaluates to True, otherwise file descriptor.
1034

1035
  """
1036
  if not os.path.isabs(file_name):
1037
    raise errors.ProgrammerError("Path passed to WriteFile is not"
1038
                                 " absolute: '%s'" % file_name)
1039

    
1040
  if [fn, data].count(None) != 1:
1041
    raise errors.ProgrammerError("fn or data required")
1042

    
1043
  if [atime, mtime].count(None) == 1:
1044
    raise errors.ProgrammerError("Both atime and mtime must be either"
1045
                                 " set or None")
1046

    
1047
  dir_name, base_name = os.path.split(file_name)
1048
  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
1049
  # here we need to make sure we remove the temp file, if any error
1050
  # leaves it in place
1051
  try:
1052
    if uid != -1 or gid != -1:
1053
      os.chown(new_name, uid, gid)
1054
    if mode:
1055
      os.chmod(new_name, mode)
1056
    if callable(prewrite):
1057
      prewrite(fd)
1058
    if data is not None:
1059
      os.write(fd, data)
1060
    else:
1061
      fn(fd)
1062
    if callable(postwrite):
1063
      postwrite(fd)
1064
    os.fsync(fd)
1065
    if atime is not None and mtime is not None:
1066
      os.utime(new_name, (atime, mtime))
1067
    os.rename(new_name, file_name)
1068
  finally:
1069
    if close:
1070
      os.close(fd)
1071
      result = None
1072
    else:
1073
      result = fd
1074
    RemoveFile(new_name)
1075

    
1076
  return result
1077

    
1078

    
1079
def all(seq, pred=bool):
1080
  "Returns True if pred(x) is True for every element in the iterable"
1081
  for elem in itertools.ifilterfalse(pred, seq):
1082
    return False
1083
  return True
1084

    
1085

    
1086
def any(seq, pred=bool):
1087
  "Returns True if pred(x) is True for at least one element in the iterable"
1088
  for elem in itertools.ifilter(pred, seq):
1089
    return True
1090
  return False
1091

    
1092

    
1093
def UniqueSequence(seq):
1094
  """Returns a list with unique elements.
1095

1096
  Element order is preserved.
1097
  """
1098
  seen = set()
1099
  return [i for i in seq if i not in seen and not seen.add(i)]
1100

    
1101

    
1102
def IsValidMac(mac):
1103
  """Predicate to check if a MAC address is valid.
1104

1105
  Checks wether the supplied MAC address is formally correct, only
1106
  accepts colon separated format.
1107
  """
1108
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
1109
  return mac_check.match(mac) is not None
1110

    
1111

    
1112
def TestDelay(duration):
1113
  """Sleep for a fixed amount of time.
1114

1115
  """
1116
  if duration < 0:
1117
    return False
1118
  time.sleep(duration)
1119
  return True
1120

    
1121

    
1122
def Daemonize(logfile, noclose_fds=None):
1123
  """Daemonize the current process.
1124

1125
  This detaches the current process from the controlling terminal and
1126
  runs it in the background as a daemon.
1127

1128
  """
1129
  UMASK = 077
1130
  WORKDIR = "/"
1131
  # Default maximum for the number of available file descriptors.
1132
  if 'SC_OPEN_MAX' in os.sysconf_names:
1133
    try:
1134
      MAXFD = os.sysconf('SC_OPEN_MAX')
1135
      if MAXFD < 0:
1136
        MAXFD = 1024
1137
    except OSError:
1138
      MAXFD = 1024
1139
  else:
1140
    MAXFD = 1024
1141

    
1142
  # this might fail
1143
  pid = os.fork()
1144
  if (pid == 0):  # The first child.
1145
    os.setsid()
1146
    # this might fail
1147
    pid = os.fork() # Fork a second child.
1148
    if (pid == 0):  # The second child.
1149
      os.chdir(WORKDIR)
1150
      os.umask(UMASK)
1151
    else:
1152
      # exit() or _exit()?  See below.
1153
      os._exit(0) # Exit parent (the first child) of the second child.
1154
  else:
1155
    os._exit(0) # Exit parent of the first child.
1156
  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1157
  if (maxfd == resource.RLIM_INFINITY):
1158
    maxfd = MAXFD
1159

    
1160
  # Iterate through and close all file descriptors.
1161
  for fd in range(0, maxfd):
1162
    if noclose_fds and fd in noclose_fds:
1163
      assert fd > 2, "Cannot keep open standard fd %d" % fd
1164
      continue
1165
    try:
1166
      os.close(fd)
1167
    except OSError: # ERROR, fd wasn't open to begin with (ignored)
1168
      pass
1169
  os.open(logfile, os.O_RDWR|os.O_CREAT|os.O_APPEND, 0600)
1170
  # Duplicate standard input to standard output and standard error.
1171
  os.dup2(0, 1)     # standard output (1)
1172
  os.dup2(0, 2)     # standard error (2)
1173
  return 0
1174

    
1175

    
1176
def FindFile(name, search_path, test=os.path.exists):
1177
  """Look for a filesystem object in a given path.
1178

1179
  This is an abstract method to search for filesystem object (files,
1180
  dirs) under a given search path.
1181

1182
  Args:
1183
    - name: the name to look for
1184
    - search_path: list of directory names
1185
    - test: the test which the full path must satisfy
1186
      (defaults to os.path.exists)
1187

1188
  Returns:
1189
    - full path to the item if found
1190
    - None otherwise
1191

1192
  """
1193
  # validate the filename mask
1194
  if constants.EXT_PLUGIN_MASK.match(name) is None:
1195
    logger.Error("Invalid value passed for external script name: '%s'" %
1196
                 name)
1197
    return None
1198

    
1199
  for dir_name in search_path:
1200
    item_name = os.path.sep.join([dir_name, name])
1201
    # check the user test and that we're indeed resolving to the given
1202
    # basename
1203
    if test(item_name) and os.path.basename(item_name) == name:
1204
      return item_name
1205
  return None