Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ b74159ee

History | View | Annotate | Download (29.4 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
no_fork = False
55

    
56

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

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

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

    
75

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

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

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

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

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

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

    
103

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

    
108

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

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

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

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

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

    
128
  errcount = 0
129

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

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

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

    
156
      if not IsProcessAlive(pid):
157
        raise errors.LockError("Stale lockfile %s for pid %d?" %
158
                               (lockfile, pid))
159

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

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

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

    
173
  _locksheld.append(name)
174

    
175

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

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

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

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

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

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

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

    
202

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

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

    
210

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

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

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

220
  Returns: `RunResult` instance
221

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

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

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

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

    
270
  out = out.getvalue()
271
  err = err.getvalue()
272

    
273
  status = child.wait()
274
  if status >= 0:
275
    exitcode = status
276
    signal = None
277
  else:
278
    exitcode = None
279
    signal = -status
280

    
281
  return RunResult(exitcode, signal, out, err, strcmd)
282

    
283

    
284
def RunCmdUnlocked(cmd):
285
  """Execute a shell command without the 'cmd' lock.
286

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

291
  The argument and return values are the same as for the `RunCmd()`
292
  function.
293

294
  Args:
295
    cmd - command to run. (str)
296

297
  Returns:
298
    `RunResult`
299

300
  """
301
  Unlock('cmd')
302
  ret = RunCmd(cmd)
303
  Lock('cmd')
304

    
305
  return ret
306

    
307

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

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

314
  """
315
  try:
316
    os.unlink(filename)
317
  except OSError, err:
318
    if err.errno not in (errno.ENOENT, errno.EISDIR):
319
      raise
320

    
321

    
322
def _FingerprintFile(filename):
323
  """Compute the fingerprint of a file.
324

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

328
  Args:
329
    filename - Filename (str)
330

331
  """
332
  if not (os.path.exists(filename) and os.path.isfile(filename)):
333
    return None
334

    
335
  f = open(filename)
336

    
337
  fp = sha.sha()
338
  while True:
339
    data = f.read(4096)
340
    if not data:
341
      break
342

    
343
    fp.update(data)
344

    
345
  return fp.hexdigest()
346

    
347

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

351
  Args:
352
    files - array of filenames.  ( [str, ...] )
353

354
  Return value:
355
    dictionary of filename: fingerprint for the files that exist
356

357
  """
358
  ret = {}
359

    
360
  for filename in files:
361
    cksum = _FingerprintFile(filename)
362
    if cksum:
363
      ret[filename] = cksum
364

    
365
  return ret
366

    
367

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

371
  For the given dictionaries `target` and `template`, ensure target
372
  has all the keys from template. Missing keys are added with values
373
  from template.
374

375
  Args:
376
    target   - the dictionary to check
377
    template - template dictionary
378
    logname  - a caller-chosen string to identify the debug log
379
               entry; if None, no logging will be done
380

381
  Returns value:
382
    None
383

384
  """
385
  missing = []
386
  for k in template:
387
    if k not in target:
388
      missing.append(k)
389
      target[k] = template[k]
390

    
391
  if missing and logname:
392
    logger.Debug('%s missing keys %s' %
393
                 (logname, ', '.join(missing)))
394

    
395

    
396
def IsProcessAlive(pid):
397
  """Check if a given pid exists on the system.
398

399
  Returns: true or false, depending on if the pid exists or not
400

401
  Remarks: zombie processes treated as not alive
402

403
  """
404
  try:
405
    f = open("/proc/%d/status" % pid)
406
  except IOError, err:
407
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
408
      return False
409

    
410
  alive = True
411
  try:
412
    data = f.readlines()
413
    if len(data) > 1:
414
      state = data[1].split()
415
      if len(state) > 1 and state[1] == "Z":
416
        alive = False
417
  finally:
418
    f.close()
419

    
420
  return alive
421

    
422

    
423
def MatchNameComponent(key, name_list):
424
  """Try to match a name against a list.
425

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

432
  Args:
433
    key: the name to be searched
434
    name_list: the list of strings against which to search the key
435

436
  Returns:
437
    None if there is no match *or* if there are multiple matches
438
    otherwise the element from the list which matches
439

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

    
447

    
448
class HostInfo:
449
  """Class implementing resolver and hostname functionality
450

451
  """
452
  def __init__(self, name=None):
453
    """Initialize the host name object.
454

455
    If the name argument is not passed, it will use this system's
456
    name.
457

458
    """
459
    if name is None:
460
      name = self.SysName()
461

    
462
    self.query = name
463
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
464
    self.ip = self.ipaddrs[0]
465

    
466
  def ShortName(self):
467
    """Returns the hostname without domain.
468

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

    
472
  @staticmethod
473
  def SysName():
474
    """Return the current system's name.
475

476
    This is simply a wrapper over socket.gethostname()
477

478
    """
479
    return socket.gethostname()
480

    
481
  @staticmethod
482
  def LookupHostname(hostname):
483
    """Look up hostname
484

485
    Args:
486
      hostname: hostname to look up
487

488
    Returns:
489
      a tuple (name, aliases, ipaddrs) as returned by socket.gethostbyname_ex
490
      in case of errors in resolving, we raise a ResolverError
491

492
    """
493
    try:
494
      result = socket.gethostbyname_ex(hostname)
495
    except socket.gaierror, err:
496
      # hostname not found in DNS
497
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
498

    
499
    return result
500

    
501

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

505
  Returns:
506
     Dictionary with keys volume name and values the size of the volume
507

508
  """
509
  command = "vgs --noheadings --units m --nosuffix -o name,size"
510
  result = RunCmd(command)
511
  retval = {}
512
  if result.failed:
513
    return retval
514

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

    
523
    retval[name] = size
524

    
525
  return retval
526

    
527

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

531
  Returns:
532
     True if it does, false otherwise.
533

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

    
537

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

541
  Given a list of names ['a1', 'a10', 'a11', 'a2'] this function will
542
  sort the list in the logical order ['a1', 'a2', 'a10', 'a11'].
543

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

548
  Return value
549
    - a copy of the list sorted according to our algorithm
550

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

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

    
571

    
572
def CheckDaemonAlive(pid_file, process_string):
573
  """Check wether the specified daemon is alive.
574

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

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

591
  """
592
  try:
593
    pid_file = file(pid_file, 'r')
594
    try:
595
      pid = int(pid_file.readline())
596
    finally:
597
      pid_file.close()
598

    
599
    cmdline_file_path = "/proc/%s/cmdline" % (pid)
600
    cmdline_file = open(cmdline_file_path, 'r')
601
    try:
602
      cmdline = cmdline_file.readline()
603
    finally:
604
      cmdline_file.close()
605

    
606
    if not process_string in cmdline:
607
      return False
608

    
609
    stat_file_path =  "/proc/%s/stat" % (pid)
610
    stat_file = open(stat_file_path, 'r')
611
    try:
612
      process_state = stat_file.readline().split()[2]
613
    finally:
614
      stat_file.close()
615

    
616
    if process_state == 'Z':
617
      return False
618

    
619
  except (IndexError, IOError, ValueError):
620
    return False
621

    
622
  return True
623

    
624

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

628
  This function tries to apply function `fn` to `val`. If no
629
  ValueError or TypeError exceptions are raised, it will return the
630
  result, else it will return the original value. Any other exceptions
631
  are propagated to the caller.
632

633
  """
634
  try:
635
    nv = fn(val)
636
  except (ValueError, TypeError), err:
637
    nv = val
638
  return nv
639

    
640

    
641
def IsValidIP(ip):
642
  """Verifies the syntax of an IP address.
643

644
  This function checks if the ip address passes is valid or not based
645
  on syntax (not ip range, class calculations or anything).
646

647
  """
648
  unit = "(0|[1-9]\d{0,2})"
649
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
650

    
651

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

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

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

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

    
665

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

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

674
  """
675
  for word in args:
676
    if not IsValidShellParam(word):
677
      raise errors.ProgrammerError("Shell argument '%s' contains"
678
                                   " invalid characters" % word)
679
  return template % args
680

    
681

    
682
def FormatUnit(value):
683
  """Formats an incoming number of MiB with the appropriate unit.
684

685
  Value needs to be passed as a numeric type. Return value is always a string.
686

687
  """
688
  if value < 1024:
689
    return "%dM" % round(value, 0)
690

    
691
  elif value < (1024 * 1024):
692
    return "%0.1fG" % round(float(value) / 1024, 1)
693

    
694
  else:
695
    return "%0.1fT" % round(float(value) / 1024 / 1024, 1)
696

    
697

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

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

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

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

    
711
  unit = m.groups()[1]
712
  if unit:
713
    lcunit = unit.lower()
714
  else:
715
    lcunit = 'm'
716

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

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

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

    
727
  else:
728
    raise errors.UnitParseError("Unknown unit: %s" % unit)
729

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

    
734
  # Round up to the next multiple of 4
735
  value = int(value)
736
  if value % 4:
737
    value += 4 - value % 4
738

    
739
  return value
740

    
741

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

745
  Args:
746
    file_name: Path to authorized_keys file
747
    key: String containing key
748
  """
749
  key_fields = key.split()
750

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

    
768

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

772
  Args:
773
    file_name: Path to authorized_keys file
774
    key: String containing key
775
  """
776
  key_fields = key.split()
777

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

    
789
        out.flush()
790
        os.rename(tmpname, file_name)
791
      finally:
792
        f.close()
793
    finally:
794
      out.close()
795
  except:
796
    RemoveFile(tmpname)
797
    raise
798

    
799

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

803
  """
804
  # Ensure aliases are unique
805
  aliases = UniqueSequence([hostname] + aliases)[1:]
806

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

    
820
        out.write("%s\t%s" % (ip, hostname))
821
        if aliases:
822
          out.write(" %s" % ' '.join(aliases))
823
        out.write('\n')
824

    
825
        out.flush()
826
        os.fsync(out)
827
        os.rename(tmpname, file_name)
828
      finally:
829
        f.close()
830
    finally:
831
      out.close()
832
  except:
833
    RemoveFile(tmpname)
834
    raise
835

    
836

    
837
def RemoveEtcHostsEntry(file_name, hostname):
838
  """Removes a hostname from /etc/hosts.
839

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

    
859
          out.write(line)
860

    
861
        out.flush()
862
        os.fsync(out)
863
        os.rename(tmpname, file_name)
864
      finally:
865
        f.close()
866
    finally:
867
      out.close()
868
  except:
869
    RemoveFile(tmpname)
870
    raise
871

    
872

    
873
def CreateBackup(file_name):
874
  """Creates a backup of a file.
875

876
  Returns: the path to the newly created backup file.
877

878
  """
879
  if not os.path.isfile(file_name):
880
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
881
                                file_name)
882

    
883
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
884
  dir_name = os.path.dirname(file_name)
885

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

    
897
  return backup_name
898

    
899

    
900
def ShellQuote(value):
901
  """Quotes shell argument according to POSIX.
902

903
  """
904
  if _re_shell_unquoted.match(value):
905
    return value
906
  else:
907
    return "'%s'" % value.replace("'", "'\\''")
908

    
909

    
910
def ShellQuoteArgs(args):
911
  """Quotes all given shell arguments and concatenates using spaces.
912

913
  """
914
  return ' '.join([ShellQuote(i) for i in args])
915

    
916

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

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

929
  """
930
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
931

    
932
  sucess = False
933

    
934
  if source is not None:
935
    try:
936
      sock.bind((source, 0))
937
    except socket.error, (errcode, errstring):
938
      if errcode == errno.EADDRNOTAVAIL:
939
        success = False
940

    
941
  sock.settimeout(timeout)
942

    
943
  try:
944
    sock.connect((target, port))
945
    sock.close()
946
    success = True
947
  except socket.timeout:
948
    success = False
949
  except socket.error, (errcode, errstring):
950
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
951

    
952
  return success
953

    
954

    
955
def ListVisibleFiles(path):
956
  """Returns a list of all visible files in a directory.
957

958
  """
959
  files = [i for i in os.listdir(path) if not i.startswith(".")]
960
  files.sort()
961
  return files
962

    
963

    
964
def GetHomeDir(user, default=None):
965
  """Try to get the homedir of the given user.
966

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

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

    
984

    
985
def NewUUID():
986
  """Returns a random UUID.
987

988
  """
989
  f = open("/proc/sys/kernel/random/uuid", "r")
990
  try:
991
    return f.read(128).rstrip("\n")
992
  finally:
993
    f.close()
994

    
995

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

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

1008
  If the function doesn't raise an exception, it has succeeded and the
1009
  target file has the new contents. If the file has raised an
1010
  exception, an existing target file should be unmodified and the
1011
  temporary file should be removed.
1012

1013
  """
1014
  if check_abspath and not os.path.isabs(file_name):
1015
    raise errors.ProgrammerError("Path passed to WriteFile is not"
1016
                                 " absolute: '%s'" % file_name)
1017

    
1018
  if [fn, data].count(None) != 1:
1019
    raise errors.ProgrammerError("fn or data required")
1020

    
1021
  if [atime, mtime].count(None) == 1:
1022
    raise errors.ProgrammerError("Both atime and mtime must be either"
1023
                                 " set or None")
1024

    
1025
  if backup and not dry_run and os.path.isfile(file_name):
1026
    CreateBackup(file_name)
1027

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

    
1050

    
1051
def all(seq, pred=bool):
1052
  "Returns True if pred(x) is True for every element in the iterable"
1053
  for elem in itertools.ifilterfalse(pred, seq):
1054
    return False
1055
  return True
1056

    
1057

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

    
1064

    
1065
def UniqueSequence(seq):
1066
  """Returns a list with unique elements.
1067

1068
  Element order is preserved.
1069
  """
1070
  seen = set()
1071
  return [i for i in seq if i not in seen and not seen.add(i)]
1072

    
1073

    
1074
def IsValidMac(mac):
1075
  """Predicate to check if a MAC address is valid.
1076

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

    
1083

    
1084
def TestDelay(duration):
1085
  """Sleep for a fixed amount of time.
1086

1087
  """
1088
  if duration < 0:
1089
    return False
1090
  time.sleep(duration)
1091
  return True
1092

    
1093

    
1094
def Daemonize(logfile, noclose_fds=None):
1095
  """Daemonize the current process.
1096

1097
  This detaches the current process from the controlling terminal and
1098
  runs it in the background as a daemon.
1099

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

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

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

    
1146

    
1147
def FindFile(name, search_path, test=os.path.exists):
1148
  """Look for a filesystem object in a given path.
1149

1150
  This is an abstract method to search for filesystem object (files,
1151
  dirs) under a given search path.
1152

1153
  Args:
1154
    - name: the name to look for
1155
    - search_path: list of directory names
1156
    - test: the test which the full path must satisfy
1157
      (defaults to os.path.exists)
1158

1159
  Returns:
1160
    - full path to the item if found
1161
    - None otherwise
1162

1163
  """
1164
  for dir_name in search_path:
1165
    item_name = os.path.sep.join([dir_name, name])
1166
    if test(item_name):
1167
      return item_name
1168
  return None