Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ 70f4497c

History | View | Annotate | Download (29.3 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

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

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

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

    
74

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

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

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

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

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

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

    
102

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

    
107

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

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

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

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

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

    
127
  errcount = 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
        raise errors.LockError("Stale lockfile %s for pid %d?" %
157
                               (lockfile, pid))
158

    
159
      if max_retries and max_retries <= retries:
160
        raise errors.LockError("Can't acquire lock during the specified"
161
                               " time, aborting.")
162
      if retries == 5 and (debug or sys.stdin.isatty()):
163
        logger.ToStderr("Waiting for '%s' lock from pid %d..." % (name, pid))
164

    
165
      time.sleep(1)
166
      retries += 1
167
      continue
168

    
169
  os.write(fd, '%d\n' % (os.getpid(),))
170
  os.close(fd)
171

    
172
  _locksheld.append(name)
173

    
174

    
175
def Unlock(name):
176
  """Unlock a given subsystem.
177

178
  """
179
  lockfile = _GetLockFile(name)
180

    
181
  try:
182
    fd = os.open(lockfile, os.O_RDONLY)
183
  except OSError:
184
    raise errors.LockError('Lock "%s" not held.' % (name,))
185

    
186
  f = os.fdopen(fd, 'r')
187
  pid_str = f.read()
188

    
189
  try:
190
    pid = int(pid_str)
191
  except ValueError:
192
    raise errors.LockError('Unable to determine PID of locking process.')
193

    
194
  if pid != os.getpid():
195
    raise errors.LockError('Lock not held by me (%d != %d)' %
196
                           (os.getpid(), pid,))
197

    
198
  os.unlink(lockfile)
199
  _locksheld.remove(name)
200

    
201

    
202
def LockCleanup():
203
  """Remove all locks.
204

205
  """
206
  for lock in _locksheld:
207
    Unlock(lock)
208

    
209

    
210
def RunCmd(cmd):
211
  """Execute a (shell) command.
212

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

216
  Args:
217
    cmd: command to run. (str)
218

219
  Returns: `RunResult` instance
220

221
  """
222
  if isinstance(cmd, list):
223
    cmd = [str(val) for val in cmd]
224
    strcmd = " ".join(cmd)
225
    shell = False
226
  else:
227
    strcmd = cmd
228
    shell = True
229
  env = os.environ.copy()
230
  env["LC_ALL"] = "C"
231
  poller = select.poll()
232
  child = subprocess.Popen(cmd, shell=shell,
233
                           stderr=subprocess.PIPE,
234
                           stdout=subprocess.PIPE,
235
                           stdin=subprocess.PIPE,
236
                           close_fds=True, env=env)
237

    
238
  child.stdin.close()
239
  poller.register(child.stdout, select.POLLIN)
240
  poller.register(child.stderr, select.POLLIN)
241
  out = StringIO()
242
  err = StringIO()
243
  fdmap = {
244
    child.stdout.fileno(): (out, child.stdout),
245
    child.stderr.fileno(): (err, child.stderr),
246
    }
247
  for fd in fdmap:
248
    status = fcntl.fcntl(fd, fcntl.F_GETFL)
249
    fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
250

    
251
  while fdmap:
252
    for fd, event in poller.poll():
253
      if event & select.POLLIN or event & select.POLLPRI:
254
        data = fdmap[fd][1].read()
255
        # no data from read signifies EOF (the same as POLLHUP)
256
        if not data:
257
          poller.unregister(fd)
258
          del fdmap[fd]
259
          continue
260
        fdmap[fd][0].write(data)
261
      if (event & select.POLLNVAL or event & select.POLLHUP or
262
          event & select.POLLERR):
263
        poller.unregister(fd)
264
        del fdmap[fd]
265

    
266
  out = out.getvalue()
267
  err = err.getvalue()
268

    
269
  status = child.wait()
270
  if status >= 0:
271
    exitcode = status
272
    signal = None
273
  else:
274
    exitcode = None
275
    signal = -status
276

    
277
  return RunResult(exitcode, signal, out, err, strcmd)
278

    
279

    
280
def RunCmdUnlocked(cmd):
281
  """Execute a shell command without the 'cmd' lock.
282

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

287
  The argument and return values are the same as for the `RunCmd()`
288
  function.
289

290
  Args:
291
    cmd - command to run. (str)
292

293
  Returns:
294
    `RunResult`
295

296
  """
297
  Unlock('cmd')
298
  ret = RunCmd(cmd)
299
  Lock('cmd')
300

    
301
  return ret
302

    
303

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

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

310
  """
311
  try:
312
    os.unlink(filename)
313
  except OSError, err:
314
    if err.errno not in (errno.ENOENT, errno.EISDIR):
315
      raise
316

    
317

    
318
def _FingerprintFile(filename):
319
  """Compute the fingerprint of a file.
320

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

324
  Args:
325
    filename - Filename (str)
326

327
  """
328
  if not (os.path.exists(filename) and os.path.isfile(filename)):
329
    return None
330

    
331
  f = open(filename)
332

    
333
  fp = sha.sha()
334
  while True:
335
    data = f.read(4096)
336
    if not data:
337
      break
338

    
339
    fp.update(data)
340

    
341
  return fp.hexdigest()
342

    
343

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

347
  Args:
348
    files - array of filenames.  ( [str, ...] )
349

350
  Return value:
351
    dictionary of filename: fingerprint for the files that exist
352

353
  """
354
  ret = {}
355

    
356
  for filename in files:
357
    cksum = _FingerprintFile(filename)
358
    if cksum:
359
      ret[filename] = cksum
360

    
361
  return ret
362

    
363

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

367
  For the given dictionaries `target` and `template`, ensure target
368
  has all the keys from template. Missing keys are added with values
369
  from template.
370

371
  Args:
372
    target   - the dictionary to check
373
    template - template dictionary
374
    logname  - a caller-chosen string to identify the debug log
375
               entry; if None, no logging will be done
376

377
  Returns value:
378
    None
379

380
  """
381
  missing = []
382
  for k in template:
383
    if k not in target:
384
      missing.append(k)
385
      target[k] = template[k]
386

    
387
  if missing and logname:
388
    logger.Debug('%s missing keys %s' %
389
                 (logname, ', '.join(missing)))
390

    
391

    
392
def IsProcessAlive(pid):
393
  """Check if a given pid exists on the system.
394

395
  Returns: true or false, depending on if the pid exists or not
396

397
  Remarks: zombie processes treated as not alive
398

399
  """
400
  try:
401
    f = open("/proc/%d/status" % pid)
402
  except IOError, err:
403
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
404
      return False
405

    
406
  alive = True
407
  try:
408
    data = f.readlines()
409
    if len(data) > 1:
410
      state = data[1].split()
411
      if len(state) > 1 and state[1] == "Z":
412
        alive = False
413
  finally:
414
    f.close()
415

    
416
  return alive
417

    
418

    
419
def MatchNameComponent(key, name_list):
420
  """Try to match a name against a list.
421

422
  This function will try to match a name like test1 against a list
423
  like ['test1.example.com', 'test2.example.com', ...]. Against this
424
  list, 'test1' as well as 'test1.example' will match, but not
425
  'test1.ex'. A multiple match will be considered as no match at all
426
  (e.g. 'test1' against ['test1.example.com', 'test1.example.org']).
427

428
  Args:
429
    key: the name to be searched
430
    name_list: the list of strings against which to search the key
431

432
  Returns:
433
    None if there is no match *or* if there are multiple matches
434
    otherwise the element from the list which matches
435

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

    
443

    
444
class HostInfo:
445
  """Class implementing resolver and hostname functionality
446

447
  """
448
  def __init__(self, name=None):
449
    """Initialize the host name object.
450

451
    If the name argument is not passed, it will use this system's
452
    name.
453

454
    """
455
    if name is None:
456
      name = self.SysName()
457

    
458
    self.query = name
459
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
460
    self.ip = self.ipaddrs[0]
461

    
462
  def ShortName(self):
463
    """Returns the hostname without domain.
464

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

    
468
  @staticmethod
469
  def SysName():
470
    """Return the current system's name.
471

472
    This is simply a wrapper over socket.gethostname()
473

474
    """
475
    return socket.gethostname()
476

    
477
  @staticmethod
478
  def LookupHostname(hostname):
479
    """Look up hostname
480

481
    Args:
482
      hostname: hostname to look up
483

484
    Returns:
485
      a tuple (name, aliases, ipaddrs) as returned by socket.gethostbyname_ex
486
      in case of errors in resolving, we raise a ResolverError
487

488
    """
489
    try:
490
      result = socket.gethostbyname_ex(hostname)
491
    except socket.gaierror, err:
492
      # hostname not found in DNS
493
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
494

    
495
    return result
496

    
497

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

501
  Returns:
502
     Dictionary with keys volume name and values the size of the volume
503

504
  """
505
  command = "vgs --noheadings --units m --nosuffix -o name,size"
506
  result = RunCmd(command)
507
  retval = {}
508
  if result.failed:
509
    return retval
510

    
511
  for line in result.stdout.splitlines():
512
    try:
513
      name, size = line.split()
514
      size = int(float(size))
515
    except (IndexError, ValueError), err:
516
      logger.Error("Invalid output from vgs (%s): %s" % (err, line))
517
      continue
518

    
519
    retval[name] = size
520

    
521
  return retval
522

    
523

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

527
  Returns:
528
     True if it does, false otherwise.
529

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

    
533

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

537
  Given a list of names ['a1', 'a10', 'a11', 'a2'] this function will
538
  sort the list in the logical order ['a1', 'a2', 'a10', 'a11'].
539

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

544
  Return value
545
    - a copy of the list sorted according to our algorithm
546

547
  """
548
  _SORTER_BASE = "(\D+|\d+)"
549
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
550
                                                  _SORTER_BASE, _SORTER_BASE,
551
                                                  _SORTER_BASE, _SORTER_BASE,
552
                                                  _SORTER_BASE, _SORTER_BASE)
553
  _SORTER_RE = re.compile(_SORTER_FULL)
554
  _SORTER_NODIGIT = re.compile("^\D*$")
555
  def _TryInt(val):
556
    """Attempts to convert a variable to integer."""
557
    if val is None or _SORTER_NODIGIT.match(val):
558
      return val
559
    rval = int(val)
560
    return rval
561

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

    
567

    
568
def CheckDaemonAlive(pid_file, process_string):
569
  """Check wether the specified daemon is alive.
570

571
  Args:
572
   - pid_file: file to read the daemon pid from, the file is
573
               expected to contain only a single line containing
574
               only the PID
575
   - process_string: a substring that we expect to find in
576
                     the command line of the daemon process
577

578
  Returns:
579
   - True if the daemon is judged to be alive (that is:
580
      - the PID file exists, is readable and contains a number
581
      - a process of the specified PID is running
582
      - that process contains the specified string in its
583
        command line
584
      - the process is not in state Z (zombie))
585
   - False otherwise
586

587
  """
588
  try:
589
    pid_file = file(pid_file, 'r')
590
    try:
591
      pid = int(pid_file.readline())
592
    finally:
593
      pid_file.close()
594

    
595
    cmdline_file_path = "/proc/%s/cmdline" % (pid)
596
    cmdline_file = open(cmdline_file_path, 'r')
597
    try:
598
      cmdline = cmdline_file.readline()
599
    finally:
600
      cmdline_file.close()
601

    
602
    if not process_string in cmdline:
603
      return False
604

    
605
    stat_file_path =  "/proc/%s/stat" % (pid)
606
    stat_file = open(stat_file_path, 'r')
607
    try:
608
      process_state = stat_file.readline().split()[2]
609
    finally:
610
      stat_file.close()
611

    
612
    if process_state == 'Z':
613
      return False
614

    
615
  except (IndexError, IOError, ValueError):
616
    return False
617

    
618
  return True
619

    
620

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

624
  This function tries to apply function `fn` to `val`. If no
625
  ValueError or TypeError exceptions are raised, it will return the
626
  result, else it will return the original value. Any other exceptions
627
  are propagated to the caller.
628

629
  """
630
  try:
631
    nv = fn(val)
632
  except (ValueError, TypeError), err:
633
    nv = val
634
  return nv
635

    
636

    
637
def IsValidIP(ip):
638
  """Verifies the syntax of an IP address.
639

640
  This function checks if the ip address passes is valid or not based
641
  on syntax (not ip range, class calculations or anything).
642

643
  """
644
  unit = "(0|[1-9]\d{0,2})"
645
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
646

    
647

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

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

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

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

    
661

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

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

670
  """
671
  for word in args:
672
    if not IsValidShellParam(word):
673
      raise errors.ProgrammerError("Shell argument '%s' contains"
674
                                   " invalid characters" % word)
675
  return template % args
676

    
677

    
678
def FormatUnit(value):
679
  """Formats an incoming number of MiB with the appropriate unit.
680

681
  Value needs to be passed as a numeric type. Return value is always a string.
682

683
  """
684
  if value < 1024:
685
    return "%dM" % round(value, 0)
686

    
687
  elif value < (1024 * 1024):
688
    return "%0.1fG" % round(float(value) / 1024, 1)
689

    
690
  else:
691
    return "%0.1fT" % round(float(value) / 1024 / 1024, 1)
692

    
693

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

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

700
  """
701
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', input_string)
702
  if not m:
703
    raise errors.UnitParseError("Invalid format")
704

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

    
707
  unit = m.groups()[1]
708
  if unit:
709
    lcunit = unit.lower()
710
  else:
711
    lcunit = 'm'
712

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

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

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

    
723
  else:
724
    raise errors.UnitParseError("Unknown unit: %s" % unit)
725

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

    
730
  # Round up to the next multiple of 4
731
  value = int(value)
732
  if value % 4:
733
    value += 4 - value % 4
734

    
735
  return value
736

    
737

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

741
  Args:
742
    file_name: Path to authorized_keys file
743
    key: String containing key
744
  """
745
  key_fields = key.split()
746

    
747
  f = open(file_name, 'a+')
748
  try:
749
    nl = True
750
    for line in f:
751
      # Ignore whitespace changes
752
      if line.split() == key_fields:
753
        break
754
      nl = line.endswith('\n')
755
    else:
756
      if not nl:
757
        f.write("\n")
758
      f.write(key.rstrip('\r\n'))
759
      f.write("\n")
760
      f.flush()
761
  finally:
762
    f.close()
763

    
764

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

768
  Args:
769
    file_name: Path to authorized_keys file
770
    key: String containing key
771
  """
772
  key_fields = key.split()
773

    
774
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
775
  try:
776
    out = os.fdopen(fd, 'w')
777
    try:
778
      f = open(file_name, 'r')
779
      try:
780
        for line in f:
781
          # Ignore whitespace changes while comparing lines
782
          if line.split() != key_fields:
783
            out.write(line)
784

    
785
        out.flush()
786
        os.rename(tmpname, file_name)
787
      finally:
788
        f.close()
789
    finally:
790
      out.close()
791
  except:
792
    RemoveFile(tmpname)
793
    raise
794

    
795

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

799
  """
800
  # Ensure aliases are unique
801
  aliases = UniqueSequence([hostname] + aliases)[1:]
802

    
803
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
804
  try:
805
    out = os.fdopen(fd, 'w')
806
    try:
807
      f = open(file_name, 'r')
808
      try:
809
        written = False
810
        for line in f:
811
          fields = line.split()
812
          if fields and not fields[0].startswith('#') and ip == fields[0]:
813
            continue
814
          out.write(line)
815

    
816
        out.write("%s\t%s" % (ip, hostname))
817
        if aliases:
818
          out.write(" %s" % ' '.join(aliases))
819
        out.write('\n')
820

    
821
        out.flush()
822
        os.fsync(out)
823
        os.rename(tmpname, file_name)
824
      finally:
825
        f.close()
826
    finally:
827
      out.close()
828
  except:
829
    RemoveFile(tmpname)
830
    raise
831

    
832

    
833
def RemoveEtcHostsEntry(file_name, hostname):
834
  """Removes a hostname from /etc/hosts.
835

836
  IP addresses without names are removed from the file.
837
  """
838
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
839
  try:
840
    out = os.fdopen(fd, 'w')
841
    try:
842
      f = open(file_name, 'r')
843
      try:
844
        for line in f:
845
          fields = line.split()
846
          if len(fields) > 1 and not fields[0].startswith('#'):
847
            names = fields[1:]
848
            if hostname in names:
849
              while hostname in names:
850
                names.remove(hostname)
851
              if names:
852
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
853
              continue
854

    
855
          out.write(line)
856

    
857
        out.flush()
858
        os.fsync(out)
859
        os.rename(tmpname, file_name)
860
      finally:
861
        f.close()
862
    finally:
863
      out.close()
864
  except:
865
    RemoveFile(tmpname)
866
    raise
867

    
868

    
869
def CreateBackup(file_name):
870
  """Creates a backup of a file.
871

872
  Returns: the path to the newly created backup file.
873

874
  """
875
  if not os.path.isfile(file_name):
876
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
877
                                file_name)
878

    
879
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
880
  dir_name = os.path.dirname(file_name)
881

    
882
  fsrc = open(file_name, 'rb')
883
  try:
884
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
885
    fdst = os.fdopen(fd, 'wb')
886
    try:
887
      shutil.copyfileobj(fsrc, fdst)
888
    finally:
889
      fdst.close()
890
  finally:
891
    fsrc.close()
892

    
893
  return backup_name
894

    
895

    
896
def ShellQuote(value):
897
  """Quotes shell argument according to POSIX.
898

899
  """
900
  if _re_shell_unquoted.match(value):
901
    return value
902
  else:
903
    return "'%s'" % value.replace("'", "'\\''")
904

    
905

    
906
def ShellQuoteArgs(args):
907
  """Quotes all given shell arguments and concatenates using spaces.
908

909
  """
910
  return ' '.join([ShellQuote(i) for i in args])
911

    
912

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

916
  Try to do a TCP connect(2) from an optional source IP to the
917
  specified target IP and the specified target port. If the optional
918
  parameter live_port_needed is set to true, requires the remote end
919
  to accept the connection. The timeout is specified in seconds and
920
  defaults to 10 seconds. If the source optional argument is not
921
  passed, the source address selection is left to the kernel,
922
  otherwise we try to connect using the passed address (failures to
923
  bind other than EADDRNOTAVAIL will be ignored).
924

925
  """
926
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
927

    
928
  sucess = False
929

    
930
  if source is not None:
931
    try:
932
      sock.bind((source, 0))
933
    except socket.error, (errcode, errstring):
934
      if errcode == errno.EADDRNOTAVAIL:
935
        success = False
936

    
937
  sock.settimeout(timeout)
938

    
939
  try:
940
    sock.connect((target, port))
941
    sock.close()
942
    success = True
943
  except socket.timeout:
944
    success = False
945
  except socket.error, (errcode, errstring):
946
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
947

    
948
  return success
949

    
950

    
951
def ListVisibleFiles(path):
952
  """Returns a list of all visible files in a directory.
953

954
  """
955
  files = [i for i in os.listdir(path) if not i.startswith(".")]
956
  files.sort()
957
  return files
958

    
959

    
960
def GetHomeDir(user, default=None):
961
  """Try to get the homedir of the given user.
962

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

967
  """
968
  try:
969
    if isinstance(user, basestring):
970
      result = pwd.getpwnam(user)
971
    elif isinstance(user, (int, long)):
972
      result = pwd.getpwuid(user)
973
    else:
974
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
975
                                   type(user))
976
  except KeyError:
977
    return default
978
  return result.pw_dir
979

    
980

    
981
def NewUUID():
982
  """Returns a random UUID.
983

984
  """
985
  f = open("/proc/sys/kernel/random/uuid", "r")
986
  try:
987
    return f.read(128).rstrip("\n")
988
  finally:
989
    f.close()
990

    
991

    
992
def WriteFile(file_name, fn=None, data=None,
993
              mode=None, uid=-1, gid=-1,
994
              atime=None, mtime=None,
995
              check_abspath=True, dry_run=False, backup=False):
996
  """(Over)write a file atomically.
997

998
  The file_name and either fn (a function taking one argument, the
999
  file descriptor, and which should write the data to it) or data (the
1000
  contents of the file) must be passed. The other arguments are
1001
  optional and allow setting the file mode, owner and group, and the
1002
  mtime/atime of the file.
1003

1004
  If the function doesn't raise an exception, it has succeeded and the
1005
  target file has the new contents. If the file has raised an
1006
  exception, an existing target file should be unmodified and the
1007
  temporary file should be removed.
1008

1009
  """
1010
  if check_abspath and not os.path.isabs(file_name):
1011
    raise errors.ProgrammerError("Path passed to WriteFile is not"
1012
                                 " absolute: '%s'" % file_name)
1013

    
1014
  if [fn, data].count(None) != 1:
1015
    raise errors.ProgrammerError("fn or data required")
1016

    
1017
  if [atime, mtime].count(None) == 1:
1018
    raise errors.ProgrammerError("Both atime and mtime must be either"
1019
                                 " set or None")
1020

    
1021
  if backup and not dry_run and os.path.isfile(file_name):
1022
    CreateBackup(file_name)
1023

    
1024
  dir_name, base_name = os.path.split(file_name)
1025
  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
1026
  # here we need to make sure we remove the temp file, if any error
1027
  # leaves it in place
1028
  try:
1029
    if uid != -1 or gid != -1:
1030
      os.chown(new_name, uid, gid)
1031
    if mode:
1032
      os.chmod(new_name, mode)
1033
    if data is not None:
1034
      os.write(fd, data)
1035
    else:
1036
      fn(fd)
1037
    os.fsync(fd)
1038
    if atime is not None and mtime is not None:
1039
      os.utime(new_name, (atime, mtime))
1040
    if not dry_run:
1041
      os.rename(new_name, file_name)
1042
  finally:
1043
    os.close(fd)
1044
    RemoveFile(new_name)
1045

    
1046

    
1047
def all(seq, pred=bool):
1048
  "Returns True if pred(x) is True for every element in the iterable"
1049
  for elem in itertools.ifilterfalse(pred, seq):
1050
    return False
1051
  return True
1052

    
1053

    
1054
def any(seq, pred=bool):
1055
  "Returns True if pred(x) is True for at least one element in the iterable"
1056
  for elem in itertools.ifilter(pred, seq):
1057
    return True
1058
  return False
1059

    
1060

    
1061
def UniqueSequence(seq):
1062
  """Returns a list with unique elements.
1063

1064
  Element order is preserved.
1065
  """
1066
  seen = set()
1067
  return [i for i in seq if i not in seen and not seen.add(i)]
1068

    
1069

    
1070
def IsValidMac(mac):
1071
  """Predicate to check if a MAC address is valid.
1072

1073
  Checks wether the supplied MAC address is formally correct, only
1074
  accepts colon separated format.
1075
  """
1076
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
1077
  return mac_check.match(mac) is not None
1078

    
1079

    
1080
def TestDelay(duration):
1081
  """Sleep for a fixed amount of time.
1082

1083
  """
1084
  if duration < 0:
1085
    return False
1086
  time.sleep(duration)
1087
  return True
1088

    
1089

    
1090
def Daemonize(logfile, noclose_fds=None):
1091
  """Daemonize the current process.
1092

1093
  This detaches the current process from the controlling terminal and
1094
  runs it in the background as a daemon.
1095

1096
  """
1097
  UMASK = 077
1098
  WORKDIR = "/"
1099
  # Default maximum for the number of available file descriptors.
1100
  if 'SC_OPEN_MAX' in os.sysconf_names:
1101
    try:
1102
      MAXFD = os.sysconf('SC_OPEN_MAX')
1103
      if MAXFD < 0:
1104
        MAXFD = 1024
1105
    except OSError:
1106
      MAXFD = 1024
1107
  else:
1108
    MAXFD = 1024
1109

    
1110
  # this might fail
1111
  pid = os.fork()
1112
  if (pid == 0):  # The first child.
1113
    os.setsid()
1114
    # this might fail
1115
    pid = os.fork() # Fork a second child.
1116
    if (pid == 0):  # The second child.
1117
      os.chdir(WORKDIR)
1118
      os.umask(UMASK)
1119
    else:
1120
      # exit() or _exit()?  See below.
1121
      os._exit(0) # Exit parent (the first child) of the second child.
1122
  else:
1123
    os._exit(0) # Exit parent of the first child.
1124
  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1125
  if (maxfd == resource.RLIM_INFINITY):
1126
    maxfd = MAXFD
1127

    
1128
  # Iterate through and close all file descriptors.
1129
  for fd in range(0, maxfd):
1130
    if noclose_fds and fd in noclose_fds:
1131
      continue
1132
    try:
1133
      os.close(fd)
1134
    except OSError: # ERROR, fd wasn't open to begin with (ignored)
1135
      pass
1136
  os.open(logfile, os.O_RDWR|os.O_CREAT|os.O_APPEND, 0600)
1137
  # Duplicate standard input to standard output and standard error.
1138
  os.dup2(0, 1)     # standard output (1)
1139
  os.dup2(0, 2)     # standard error (2)
1140
  return 0
1141

    
1142

    
1143
def FindFile(name, search_path, test=os.path.exists):
1144
  """Look for a filesystem object in a given path.
1145

1146
  This is an abstract method to search for filesystem object (files,
1147
  dirs) under a given search path.
1148

1149
  Args:
1150
    - name: the name to look for
1151
    - search_path: list of directory names
1152
    - test: the test which the full path must satisfy
1153
      (defaults to os.path.exists)
1154

1155
  Returns:
1156
    - full path to the item if found
1157
    - None otherwise
1158

1159
  """
1160
  for dir_name in search_path:
1161
    item_name = os.path.sep.join([dir_name, name])
1162
    if test(item_name):
1163
      return item_name
1164
  return None