Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ b15d625f

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

    
42
from cStringIO import StringIO
43

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

    
48

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

    
52
debug = False
53

    
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):
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

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

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

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

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

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

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

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

    
171
  _locksheld.append(name)
172

    
173

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

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

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

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

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

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

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

    
200

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

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

    
208

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

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

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

218
  Returns: `RunResult` instance
219

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

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

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

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

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

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

    
278

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

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

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

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

292
  Returns:
293
    `RunResult`
294

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

    
300
  return ret
301

    
302

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

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

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

    
316

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

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

323
  Args:
324
    filename - Filename (str)
325

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

    
330
  f = open(filename)
331

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

    
338
    fp.update(data)
339

    
340
  return fp.hexdigest()
341

    
342

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

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

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

352
  """
353
  ret = {}
354

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

    
360
  return ret
361

    
362

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

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

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

376
  Returns value:
377
    None
378

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

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

    
390

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

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

396
  Remarks: zombie processes treated as not alive
397

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

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

    
415
  return alive
416

    
417

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

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

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

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

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

    
442

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

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

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

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

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

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

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

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

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

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

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

480
    Args:
481
      hostname: hostname to look up
482

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

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

    
494
    return result
495

    
496

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

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

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

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

    
518
    retval[name] = size
519

    
520
  return retval
521

    
522

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

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

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

    
532

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

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

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

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

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

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

    
566

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

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

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

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

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

    
601
    if not process_string in cmdline:
602
      return False
603

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

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

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

    
617
  return True
618

    
619

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

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

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

    
635

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

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

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

    
646

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

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

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

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

    
660

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

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

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

    
676

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

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

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

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

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

    
692

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

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

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

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

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

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

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

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

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

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

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

    
734
  return value
735

    
736

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

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

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

    
763

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

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

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

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

    
794

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

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

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

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

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

    
831

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

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

    
854
          out.write(line)
855

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

    
867

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

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

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

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

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

    
892
  return backup_name
893

    
894

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

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

    
904

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

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

    
911

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

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

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

    
927
  sucess = False
928

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

    
936
  sock.settimeout(timeout)
937

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

    
947
  return success
948

    
949

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

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

    
958

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

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

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

    
979

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

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

    
990

    
991
def WriteFile(file_name, fn=None, data=None,
992
              mode=None, uid=-1, gid=-1,
993
              atime=None, mtime=None):
994
  """(Over)write a file atomically.
995

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

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

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

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

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

    
1019

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

    
1041

    
1042
def all(seq, pred=bool):
1043
  "Returns True if pred(x) is True for every element in the iterable"
1044
  for elem in itertools.ifilterfalse(pred, seq):
1045
    return False
1046
  return True
1047

    
1048

    
1049
def any(seq, pred=bool):
1050
  "Returns True if pred(x) is True for at least one element in the iterable"
1051
  for elem in itertools.ifilter(pred, seq):
1052
    return True
1053
  return False
1054

    
1055

    
1056
def UniqueSequence(seq):
1057
  """Returns a list with unique elements.
1058

1059
  Element order is preserved.
1060
  """
1061
  seen = set()
1062
  return [i for i in seq if i not in seen and not seen.add(i)]
1063

    
1064

    
1065
def IsValidMac(mac):
1066
  """Predicate to check if a MAC address is valid.
1067

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

    
1074

    
1075
def TestDelay(duration):
1076
  """Sleep for a fixed amount of time.
1077

1078
  """
1079
  if duration < 0:
1080
    return False
1081
  time.sleep(duration)
1082
  return True