Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ 7b4126b7

History | View | Annotate | Download (30.2 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 TryConvert(fn, val):
573
  """Try to convert a value ignoring errors.
574

575
  This function tries to apply function `fn` to `val`. If no
576
  ValueError or TypeError exceptions are raised, it will return the
577
  result, else it will return the original value. Any other exceptions
578
  are propagated to the caller.
579

580
  """
581
  try:
582
    nv = fn(val)
583
  except (ValueError, TypeError), err:
584
    nv = val
585
  return nv
586

    
587

    
588
def IsValidIP(ip):
589
  """Verifies the syntax of an IP address.
590

591
  This function checks if the ip address passes is valid or not based
592
  on syntax (not ip range, class calculations or anything).
593

594
  """
595
  unit = "(0|[1-9]\d{0,2})"
596
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
597

    
598

    
599
def IsValidShellParam(word):
600
  """Verifies is the given word is safe from the shell's p.o.v.
601

602
  This means that we can pass this to a command via the shell and be
603
  sure that it doesn't alter the command line and is passed as such to
604
  the actual command.
605

606
  Note that we are overly restrictive here, in order to be on the safe
607
  side.
608

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

    
612

    
613
def BuildShellCmd(template, *args):
614
  """Build a safe shell command line from the given arguments.
615

616
  This function will check all arguments in the args list so that they
617
  are valid shell parameters (i.e. they don't contain shell
618
  metacharaters). If everything is ok, it will return the result of
619
  template % args.
620

621
  """
622
  for word in args:
623
    if not IsValidShellParam(word):
624
      raise errors.ProgrammerError("Shell argument '%s' contains"
625
                                   " invalid characters" % word)
626
  return template % args
627

    
628

    
629
def FormatUnit(value):
630
  """Formats an incoming number of MiB with the appropriate unit.
631

632
  Value needs to be passed as a numeric type. Return value is always a string.
633

634
  """
635
  if value < 1024:
636
    return "%dM" % round(value, 0)
637

    
638
  elif value < (1024 * 1024):
639
    return "%0.1fG" % round(float(value) / 1024, 1)
640

    
641
  else:
642
    return "%0.1fT" % round(float(value) / 1024 / 1024, 1)
643

    
644

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

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

651
  """
652
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', input_string)
653
  if not m:
654
    raise errors.UnitParseError("Invalid format")
655

    
656
  value = float(m.groups()[0])
657

    
658
  unit = m.groups()[1]
659
  if unit:
660
    lcunit = unit.lower()
661
  else:
662
    lcunit = 'm'
663

    
664
  if lcunit in ('m', 'mb', 'mib'):
665
    # Value already in MiB
666
    pass
667

    
668
  elif lcunit in ('g', 'gb', 'gib'):
669
    value *= 1024
670

    
671
  elif lcunit in ('t', 'tb', 'tib'):
672
    value *= 1024 * 1024
673

    
674
  else:
675
    raise errors.UnitParseError("Unknown unit: %s" % unit)
676

    
677
  # Make sure we round up
678
  if int(value) < value:
679
    value += 1
680

    
681
  # Round up to the next multiple of 4
682
  value = int(value)
683
  if value % 4:
684
    value += 4 - value % 4
685

    
686
  return value
687

    
688

    
689
def AddAuthorizedKey(file_name, key):
690
  """Adds an SSH public key to an authorized_keys file.
691

692
  Args:
693
    file_name: Path to authorized_keys file
694
    key: String containing key
695
  """
696
  key_fields = key.split()
697

    
698
  f = open(file_name, 'a+')
699
  try:
700
    nl = True
701
    for line in f:
702
      # Ignore whitespace changes
703
      if line.split() == key_fields:
704
        break
705
      nl = line.endswith('\n')
706
    else:
707
      if not nl:
708
        f.write("\n")
709
      f.write(key.rstrip('\r\n'))
710
      f.write("\n")
711
      f.flush()
712
  finally:
713
    f.close()
714

    
715

    
716
def RemoveAuthorizedKey(file_name, key):
717
  """Removes an SSH public key from an authorized_keys file.
718

719
  Args:
720
    file_name: Path to authorized_keys file
721
    key: String containing key
722
  """
723
  key_fields = key.split()
724

    
725
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
726
  try:
727
    out = os.fdopen(fd, 'w')
728
    try:
729
      f = open(file_name, 'r')
730
      try:
731
        for line in f:
732
          # Ignore whitespace changes while comparing lines
733
          if line.split() != key_fields:
734
            out.write(line)
735

    
736
        out.flush()
737
        os.rename(tmpname, file_name)
738
      finally:
739
        f.close()
740
    finally:
741
      out.close()
742
  except:
743
    RemoveFile(tmpname)
744
    raise
745

    
746

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

750
  """
751
  # Ensure aliases are unique
752
  aliases = UniqueSequence([hostname] + aliases)[1:]
753

    
754
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
755
  try:
756
    out = os.fdopen(fd, 'w')
757
    try:
758
      f = open(file_name, 'r')
759
      try:
760
        written = False
761
        for line in f:
762
          fields = line.split()
763
          if fields and not fields[0].startswith('#') and ip == fields[0]:
764
            continue
765
          out.write(line)
766

    
767
        out.write("%s\t%s" % (ip, hostname))
768
        if aliases:
769
          out.write(" %s" % ' '.join(aliases))
770
        out.write('\n')
771

    
772
        out.flush()
773
        os.fsync(out)
774
        os.rename(tmpname, file_name)
775
      finally:
776
        f.close()
777
    finally:
778
      out.close()
779
  except:
780
    RemoveFile(tmpname)
781
    raise
782

    
783

    
784
def AddHostToEtcHosts(hostname):
785
  """Wrapper around SetEtcHostsEntry.
786

787
  """
788
  hi = HostInfo(name=hostname)
789
  SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])
790

    
791

    
792
def RemoveEtcHostsEntry(file_name, hostname):
793
  """Removes a hostname from /etc/hosts.
794

795
  IP addresses without names are removed from the file.
796
  """
797
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
798
  try:
799
    out = os.fdopen(fd, 'w')
800
    try:
801
      f = open(file_name, 'r')
802
      try:
803
        for line in f:
804
          fields = line.split()
805
          if len(fields) > 1 and not fields[0].startswith('#'):
806
            names = fields[1:]
807
            if hostname in names:
808
              while hostname in names:
809
                names.remove(hostname)
810
              if names:
811
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
812
              continue
813

    
814
          out.write(line)
815

    
816
        out.flush()
817
        os.fsync(out)
818
        os.rename(tmpname, file_name)
819
      finally:
820
        f.close()
821
    finally:
822
      out.close()
823
  except:
824
    RemoveFile(tmpname)
825
    raise
826

    
827

    
828
def RemoveHostFromEtcHosts(hostname):
829
  """Wrapper around RemoveEtcHostsEntry.
830

831
  """
832
  hi = HostInfo(name=hostname)
833
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.name)
834
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.ShortName())
835

    
836

    
837
def CreateBackup(file_name):
838
  """Creates a backup of a file.
839

840
  Returns: the path to the newly created backup file.
841

842
  """
843
  if not os.path.isfile(file_name):
844
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
845
                                file_name)
846

    
847
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
848
  dir_name = os.path.dirname(file_name)
849

    
850
  fsrc = open(file_name, 'rb')
851
  try:
852
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
853
    fdst = os.fdopen(fd, 'wb')
854
    try:
855
      shutil.copyfileobj(fsrc, fdst)
856
    finally:
857
      fdst.close()
858
  finally:
859
    fsrc.close()
860

    
861
  return backup_name
862

    
863

    
864
def ShellQuote(value):
865
  """Quotes shell argument according to POSIX.
866

867
  """
868
  if _re_shell_unquoted.match(value):
869
    return value
870
  else:
871
    return "'%s'" % value.replace("'", "'\\''")
872

    
873

    
874
def ShellQuoteArgs(args):
875
  """Quotes all given shell arguments and concatenates using spaces.
876

877
  """
878
  return ' '.join([ShellQuote(i) for i in args])
879

    
880

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

884
  Try to do a TCP connect(2) from an optional source IP to the
885
  specified target IP and the specified target port. If the optional
886
  parameter live_port_needed is set to true, requires the remote end
887
  to accept the connection. The timeout is specified in seconds and
888
  defaults to 10 seconds. If the source optional argument is not
889
  passed, the source address selection is left to the kernel,
890
  otherwise we try to connect using the passed address (failures to
891
  bind other than EADDRNOTAVAIL will be ignored).
892

893
  """
894
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
895

    
896
  sucess = False
897

    
898
  if source is not None:
899
    try:
900
      sock.bind((source, 0))
901
    except socket.error, (errcode, errstring):
902
      if errcode == errno.EADDRNOTAVAIL:
903
        success = False
904

    
905
  sock.settimeout(timeout)
906

    
907
  try:
908
    sock.connect((target, port))
909
    sock.close()
910
    success = True
911
  except socket.timeout:
912
    success = False
913
  except socket.error, (errcode, errstring):
914
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
915

    
916
  return success
917

    
918

    
919
def ListVisibleFiles(path):
920
  """Returns a list of all visible files in a directory.
921

922
  """
923
  files = [i for i in os.listdir(path) if not i.startswith(".")]
924
  files.sort()
925
  return files
926

    
927

    
928
def GetHomeDir(user, default=None):
929
  """Try to get the homedir of the given user.
930

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

935
  """
936
  try:
937
    if isinstance(user, basestring):
938
      result = pwd.getpwnam(user)
939
    elif isinstance(user, (int, long)):
940
      result = pwd.getpwuid(user)
941
    else:
942
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
943
                                   type(user))
944
  except KeyError:
945
    return default
946
  return result.pw_dir
947

    
948

    
949
def NewUUID():
950
  """Returns a random UUID.
951

952
  """
953
  f = open("/proc/sys/kernel/random/uuid", "r")
954
  try:
955
    return f.read(128).rstrip("\n")
956
  finally:
957
    f.close()
958

    
959

    
960
def WriteFile(file_name, fn=None, data=None,
961
              mode=None, uid=-1, gid=-1,
962
              atime=None, mtime=None, close=True,
963
              dry_run=False, backup=False,
964
              prewrite=None, postwrite=None):
965
  """(Over)write a file atomically.
966

967
  The file_name and either fn (a function taking one argument, the
968
  file descriptor, and which should write the data to it) or data (the
969
  contents of the file) must be passed. The other arguments are
970
  optional and allow setting the file mode, owner and group, and the
971
  mtime/atime of the file.
972

973
  If the function doesn't raise an exception, it has succeeded and the
974
  target file has the new contents. If the file has raised an
975
  exception, an existing target file should be unmodified and the
976
  temporary file should be removed.
977

978
  Args:
979
    file_name: New filename
980
    fn: Content writing function, called with file descriptor as parameter
981
    data: Content as string
982
    mode: File mode
983
    uid: Owner
984
    gid: Group
985
    atime: Access time
986
    mtime: Modification time
987
    close: Whether to close file after writing it
988
    prewrite: Function object called before writing content
989
    postwrite: Function object called after writing content
990

991
  Returns:
992
    None if "close" parameter evaluates to True, otherwise file descriptor.
993

994
  """
995
  if not os.path.isabs(file_name):
996
    raise errors.ProgrammerError("Path passed to WriteFile is not"
997
                                 " absolute: '%s'" % file_name)
998

    
999
  if [fn, data].count(None) != 1:
1000
    raise errors.ProgrammerError("fn or data required")
1001

    
1002
  if [atime, mtime].count(None) == 1:
1003
    raise errors.ProgrammerError("Both atime and mtime must be either"
1004
                                 " set or None")
1005

    
1006
  if backup and not dry_run and os.path.isfile(file_name):
1007
    CreateBackup(file_name)
1008

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

    
1039
  return result
1040

    
1041

    
1042
def FirstFree(seq, base=0):
1043
  """Returns the first non-existing integer from seq.
1044

1045
  The seq argument should be a sorted list of positive integers. The
1046
  first time the index of an element is smaller than the element
1047
  value, the index will be returned.
1048

1049
  The base argument is used to start at a different offset,
1050
  i.e. [3, 4, 6] with offset=3 will return 5.
1051

1052
  Example: [0, 1, 3] will return 2.
1053

1054
  """
1055
  for idx, elem in enumerate(seq):
1056
    assert elem >= base, "Passed element is higher than base offset"
1057
    if elem > idx + base:
1058
      # idx is not used
1059
      return idx + base
1060
  return None
1061

    
1062

    
1063
def all(seq, pred=bool):
1064
  "Returns True if pred(x) is True for every element in the iterable"
1065
  for elem in itertools.ifilterfalse(pred, seq):
1066
    return False
1067
  return True
1068

    
1069

    
1070
def any(seq, pred=bool):
1071
  "Returns True if pred(x) is True for at least one element in the iterable"
1072
  for elem in itertools.ifilter(pred, seq):
1073
    return True
1074
  return False
1075

    
1076

    
1077
def UniqueSequence(seq):
1078
  """Returns a list with unique elements.
1079

1080
  Element order is preserved.
1081
  """
1082
  seen = set()
1083
  return [i for i in seq if i not in seen and not seen.add(i)]
1084

    
1085

    
1086
def IsValidMac(mac):
1087
  """Predicate to check if a MAC address is valid.
1088

1089
  Checks wether the supplied MAC address is formally correct, only
1090
  accepts colon separated format.
1091
  """
1092
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
1093
  return mac_check.match(mac) is not None
1094

    
1095

    
1096
def TestDelay(duration):
1097
  """Sleep for a fixed amount of time.
1098

1099
  """
1100
  if duration < 0:
1101
    return False
1102
  time.sleep(duration)
1103
  return True
1104

    
1105

    
1106
def Daemonize(logfile, noclose_fds=None):
1107
  """Daemonize the current process.
1108

1109
  This detaches the current process from the controlling terminal and
1110
  runs it in the background as a daemon.
1111

1112
  """
1113
  UMASK = 077
1114
  WORKDIR = "/"
1115
  # Default maximum for the number of available file descriptors.
1116
  if 'SC_OPEN_MAX' in os.sysconf_names:
1117
    try:
1118
      MAXFD = os.sysconf('SC_OPEN_MAX')
1119
      if MAXFD < 0:
1120
        MAXFD = 1024
1121
    except OSError:
1122
      MAXFD = 1024
1123
  else:
1124
    MAXFD = 1024
1125

    
1126
  # this might fail
1127
  pid = os.fork()
1128
  if (pid == 0):  # The first child.
1129
    os.setsid()
1130
    # this might fail
1131
    pid = os.fork() # Fork a second child.
1132
    if (pid == 0):  # The second child.
1133
      os.chdir(WORKDIR)
1134
      os.umask(UMASK)
1135
    else:
1136
      # exit() or _exit()?  See below.
1137
      os._exit(0) # Exit parent (the first child) of the second child.
1138
  else:
1139
    os._exit(0) # Exit parent of the first child.
1140
  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1141
  if (maxfd == resource.RLIM_INFINITY):
1142
    maxfd = MAXFD
1143

    
1144
  # Iterate through and close all file descriptors.
1145
  for fd in range(0, maxfd):
1146
    if noclose_fds and fd in noclose_fds:
1147
      continue
1148
    try:
1149
      os.close(fd)
1150
    except OSError: # ERROR, fd wasn't open to begin with (ignored)
1151
      pass
1152
  os.open(logfile, os.O_RDWR|os.O_CREAT|os.O_APPEND, 0600)
1153
  # Duplicate standard input to standard output and standard error.
1154
  os.dup2(0, 1)     # standard output (1)
1155
  os.dup2(0, 2)     # standard error (2)
1156
  return 0
1157

    
1158

    
1159
def FindFile(name, search_path, test=os.path.exists):
1160
  """Look for a filesystem object in a given path.
1161

1162
  This is an abstract method to search for filesystem object (files,
1163
  dirs) under a given search path.
1164

1165
  Args:
1166
    - name: the name to look for
1167
    - search_path: list of directory names
1168
    - test: the test which the full path must satisfy
1169
      (defaults to os.path.exists)
1170

1171
  Returns:
1172
    - full path to the item if found
1173
    - None otherwise
1174

1175
  """
1176
  for dir_name in search_path:
1177
    item_name = os.path.sep.join([dir_name, name])
1178
    if test(item_name):
1179
      return item_name
1180
  return None
1181

    
1182

    
1183
def CheckVolumeGroupSize(vglist, vgname, minsize):
1184
  """Checks if the volume group list is valid.
1185

1186
  A non-None return value means there's an error, and the return value
1187
  is the error message.
1188

1189
  """
1190
  vgsize = vglist.get(vgname, None)
1191
  if vgsize is None:
1192
    return "volume group '%s' missing" % vgname
1193
  elif vgsize < minsize:
1194
    return ("volume group '%s' too small (%s MiB required, %d MiB found)" %
1195
            (vgname, minsize, vgsize))
1196
  return None