Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ fee80e90

History | View | Annotate | Download (29.9 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
import logging
43
import signal
44

    
45
from cStringIO import StringIO
46

    
47
from ganeti import errors
48
from ganeti import constants
49

    
50

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

    
54
debug = False
55
no_fork = False
56

    
57

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

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

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

    
76

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

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

    
92
    if self.failed:
93
      logging.debug("Command '%s' failed (%s); output: %s",
94
                    self.cmd, self.fail_reason, self.output)
95

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

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

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

    
104

    
105
def RunCmd(cmd):
106
  """Execute a (shell) command.
107

108
  The command should not read from its standard input, as it will be
109
  closed.
110

111
  Args:
112
    cmd: command to run. (str)
113

114
  Returns: `RunResult` instance
115

116
  """
117
  if no_fork:
118
    raise errors.ProgrammerError("utils.RunCmd() called with fork() disabled")
119

    
120
  if isinstance(cmd, list):
121
    cmd = [str(val) for val in cmd]
122
    strcmd = " ".join(cmd)
123
    shell = False
124
  else:
125
    strcmd = cmd
126
    shell = True
127
  logging.debug("RunCmd '%s'", strcmd)
128
  env = os.environ.copy()
129
  env["LC_ALL"] = "C"
130
  poller = select.poll()
131
  child = subprocess.Popen(cmd, shell=shell,
132
                           stderr=subprocess.PIPE,
133
                           stdout=subprocess.PIPE,
134
                           stdin=subprocess.PIPE,
135
                           close_fds=True, env=env)
136

    
137
  child.stdin.close()
138
  poller.register(child.stdout, select.POLLIN)
139
  poller.register(child.stderr, select.POLLIN)
140
  out = StringIO()
141
  err = StringIO()
142
  fdmap = {
143
    child.stdout.fileno(): (out, child.stdout),
144
    child.stderr.fileno(): (err, child.stderr),
145
    }
146
  for fd in fdmap:
147
    status = fcntl.fcntl(fd, fcntl.F_GETFL)
148
    fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
149

    
150
  while fdmap:
151
    for fd, event in poller.poll():
152
      if event & select.POLLIN or event & select.POLLPRI:
153
        data = fdmap[fd][1].read()
154
        # no data from read signifies EOF (the same as POLLHUP)
155
        if not data:
156
          poller.unregister(fd)
157
          del fdmap[fd]
158
          continue
159
        fdmap[fd][0].write(data)
160
      if (event & select.POLLNVAL or event & select.POLLHUP or
161
          event & select.POLLERR):
162
        poller.unregister(fd)
163
        del fdmap[fd]
164

    
165
  out = out.getvalue()
166
  err = err.getvalue()
167

    
168
  status = child.wait()
169
  if status >= 0:
170
    exitcode = status
171
    signal = None
172
  else:
173
    exitcode = None
174
    signal = -status
175

    
176
  return RunResult(exitcode, signal, out, err, strcmd)
177

    
178

    
179
def RemoveFile(filename):
180
  """Remove a file ignoring some errors.
181

182
  Remove a file, ignoring non-existing ones or directories. Other
183
  errors are passed.
184

185
  """
186
  try:
187
    os.unlink(filename)
188
  except OSError, err:
189
    if err.errno not in (errno.ENOENT, errno.EISDIR):
190
      raise
191

    
192

    
193
def _FingerprintFile(filename):
194
  """Compute the fingerprint of a file.
195

196
  If the file does not exist, a None will be returned
197
  instead.
198

199
  Args:
200
    filename - Filename (str)
201

202
  """
203
  if not (os.path.exists(filename) and os.path.isfile(filename)):
204
    return None
205

    
206
  f = open(filename)
207

    
208
  fp = sha.sha()
209
  while True:
210
    data = f.read(4096)
211
    if not data:
212
      break
213

    
214
    fp.update(data)
215

    
216
  return fp.hexdigest()
217

    
218

    
219
def FingerprintFiles(files):
220
  """Compute fingerprints for a list of files.
221

222
  Args:
223
    files - array of filenames.  ( [str, ...] )
224

225
  Return value:
226
    dictionary of filename: fingerprint for the files that exist
227

228
  """
229
  ret = {}
230

    
231
  for filename in files:
232
    cksum = _FingerprintFile(filename)
233
    if cksum:
234
      ret[filename] = cksum
235

    
236
  return ret
237

    
238

    
239
def CheckDict(target, template, logname=None):
240
  """Ensure a dictionary has a required set of keys.
241

242
  For the given dictionaries `target` and `template`, ensure target
243
  has all the keys from template. Missing keys are added with values
244
  from template.
245

246
  Args:
247
    target   - the dictionary to check
248
    template - template dictionary
249
    logname  - a caller-chosen string to identify the debug log
250
               entry; if None, no logging will be done
251

252
  Returns value:
253
    None
254

255
  """
256
  missing = []
257
  for k in template:
258
    if k not in target:
259
      missing.append(k)
260
      target[k] = template[k]
261

    
262
  if missing and logname:
263
    logging.warning('%s missing keys %s', logname, ', '.join(missing))
264

    
265

    
266
def IsProcessAlive(pid):
267
  """Check if a given pid exists on the system.
268

269
  Returns: true or false, depending on if the pid exists or not
270

271
  Remarks: zombie processes treated as not alive
272

273
  """
274
  try:
275
    f = open("/proc/%d/status" % pid)
276
  except IOError, err:
277
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
278
      return False
279

    
280
  alive = True
281
  try:
282
    data = f.readlines()
283
    if len(data) > 1:
284
      state = data[1].split()
285
      if len(state) > 1 and state[1] == "Z":
286
        alive = False
287
  finally:
288
    f.close()
289

    
290
  return alive
291

    
292

    
293
def IsPidFileAlive(pidfile):
294
  """Check whether the given pidfile points to a live process.
295

296
    @param pidfile: Path to a file containing the pid to be checked
297
    @type  pidfile: string (filename)
298

299
  """
300
  try:
301
    pf = open(pidfile, 'r')
302
  except EnvironmentError, open_err:
303
    if open_err.errno == errno.ENOENT:
304
      return False
305
    else:
306
      raise errors.GenericError("Cannot open file %s. Error: %s" %
307
                                (pidfile, str(open_err)))
308

    
309
  try:
310
    pid = int(pf.read())
311
  except ValueError:
312
    raise errors.GenericError("Invalid pid string in %s" %
313
                              (pidfile,))
314

    
315
  return IsProcessAlive(pid)
316

    
317

    
318
def MatchNameComponent(key, name_list):
319
  """Try to match a name against a list.
320

321
  This function will try to match a name like test1 against a list
322
  like ['test1.example.com', 'test2.example.com', ...]. Against this
323
  list, 'test1' as well as 'test1.example' will match, but not
324
  'test1.ex'. A multiple match will be considered as no match at all
325
  (e.g. 'test1' against ['test1.example.com', 'test1.example.org']).
326

327
  Args:
328
    key: the name to be searched
329
    name_list: the list of strings against which to search the key
330

331
  Returns:
332
    None if there is no match *or* if there are multiple matches
333
    otherwise the element from the list which matches
334

335
  """
336
  mo = re.compile("^%s(\..*)?$" % re.escape(key))
337
  names_filtered = [name for name in name_list if mo.match(name) is not None]
338
  if len(names_filtered) != 1:
339
    return None
340
  return names_filtered[0]
341

    
342

    
343
class HostInfo:
344
  """Class implementing resolver and hostname functionality
345

346
  """
347
  def __init__(self, name=None):
348
    """Initialize the host name object.
349

350
    If the name argument is not passed, it will use this system's
351
    name.
352

353
    """
354
    if name is None:
355
      name = self.SysName()
356

    
357
    self.query = name
358
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
359
    self.ip = self.ipaddrs[0]
360

    
361
  def ShortName(self):
362
    """Returns the hostname without domain.
363

364
    """
365
    return self.name.split('.')[0]
366

    
367
  @staticmethod
368
  def SysName():
369
    """Return the current system's name.
370

371
    This is simply a wrapper over socket.gethostname()
372

373
    """
374
    return socket.gethostname()
375

    
376
  @staticmethod
377
  def LookupHostname(hostname):
378
    """Look up hostname
379

380
    Args:
381
      hostname: hostname to look up
382

383
    Returns:
384
      a tuple (name, aliases, ipaddrs) as returned by socket.gethostbyname_ex
385
      in case of errors in resolving, we raise a ResolverError
386

387
    """
388
    try:
389
      result = socket.gethostbyname_ex(hostname)
390
    except socket.gaierror, err:
391
      # hostname not found in DNS
392
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
393

    
394
    return result
395

    
396

    
397
def ListVolumeGroups():
398
  """List volume groups and their size
399

400
  Returns:
401
     Dictionary with keys volume name and values the size of the volume
402

403
  """
404
  command = "vgs --noheadings --units m --nosuffix -o name,size"
405
  result = RunCmd(command)
406
  retval = {}
407
  if result.failed:
408
    return retval
409

    
410
  for line in result.stdout.splitlines():
411
    try:
412
      name, size = line.split()
413
      size = int(float(size))
414
    except (IndexError, ValueError), err:
415
      logging.error("Invalid output from vgs (%s): %s", err, line)
416
      continue
417

    
418
    retval[name] = size
419

    
420
  return retval
421

    
422

    
423
def BridgeExists(bridge):
424
  """Check whether the given bridge exists in the system
425

426
  Returns:
427
     True if it does, false otherwise.
428

429
  """
430
  return os.path.isdir("/sys/class/net/%s/bridge" % bridge)
431

    
432

    
433
def NiceSort(name_list):
434
  """Sort a list of strings based on digit and non-digit groupings.
435

436
  Given a list of names ['a1', 'a10', 'a11', 'a2'] this function will
437
  sort the list in the logical order ['a1', 'a2', 'a10', 'a11'].
438

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

443
  Return value
444
    - a copy of the list sorted according to our algorithm
445

446
  """
447
  _SORTER_BASE = "(\D+|\d+)"
448
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
449
                                                  _SORTER_BASE, _SORTER_BASE,
450
                                                  _SORTER_BASE, _SORTER_BASE,
451
                                                  _SORTER_BASE, _SORTER_BASE)
452
  _SORTER_RE = re.compile(_SORTER_FULL)
453
  _SORTER_NODIGIT = re.compile("^\D*$")
454
  def _TryInt(val):
455
    """Attempts to convert a variable to integer."""
456
    if val is None or _SORTER_NODIGIT.match(val):
457
      return val
458
    rval = int(val)
459
    return rval
460

    
461
  to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
462
             for name in name_list]
463
  to_sort.sort()
464
  return [tup[1] for tup in to_sort]
465

    
466

    
467
def TryConvert(fn, val):
468
  """Try to convert a value ignoring errors.
469

470
  This function tries to apply function `fn` to `val`. If no
471
  ValueError or TypeError exceptions are raised, it will return the
472
  result, else it will return the original value. Any other exceptions
473
  are propagated to the caller.
474

475
  """
476
  try:
477
    nv = fn(val)
478
  except (ValueError, TypeError), err:
479
    nv = val
480
  return nv
481

    
482

    
483
def IsValidIP(ip):
484
  """Verifies the syntax of an IP address.
485

486
  This function checks if the ip address passes is valid or not based
487
  on syntax (not ip range, class calculations or anything).
488

489
  """
490
  unit = "(0|[1-9]\d{0,2})"
491
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
492

    
493

    
494
def IsValidShellParam(word):
495
  """Verifies is the given word is safe from the shell's p.o.v.
496

497
  This means that we can pass this to a command via the shell and be
498
  sure that it doesn't alter the command line and is passed as such to
499
  the actual command.
500

501
  Note that we are overly restrictive here, in order to be on the safe
502
  side.
503

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

    
507

    
508
def BuildShellCmd(template, *args):
509
  """Build a safe shell command line from the given arguments.
510

511
  This function will check all arguments in the args list so that they
512
  are valid shell parameters (i.e. they don't contain shell
513
  metacharaters). If everything is ok, it will return the result of
514
  template % args.
515

516
  """
517
  for word in args:
518
    if not IsValidShellParam(word):
519
      raise errors.ProgrammerError("Shell argument '%s' contains"
520
                                   " invalid characters" % word)
521
  return template % args
522

    
523

    
524
def FormatUnit(value):
525
  """Formats an incoming number of MiB with the appropriate unit.
526

527
  Value needs to be passed as a numeric type. Return value is always a string.
528

529
  """
530
  if value < 1024:
531
    return "%dM" % round(value, 0)
532

    
533
  elif value < (1024 * 1024):
534
    return "%0.1fG" % round(float(value) / 1024, 1)
535

    
536
  else:
537
    return "%0.1fT" % round(float(value) / 1024 / 1024, 1)
538

    
539

    
540
def ParseUnit(input_string):
541
  """Tries to extract number and scale from the given string.
542

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

546
  """
547
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', input_string)
548
  if not m:
549
    raise errors.UnitParseError("Invalid format")
550

    
551
  value = float(m.groups()[0])
552

    
553
  unit = m.groups()[1]
554
  if unit:
555
    lcunit = unit.lower()
556
  else:
557
    lcunit = 'm'
558

    
559
  if lcunit in ('m', 'mb', 'mib'):
560
    # Value already in MiB
561
    pass
562

    
563
  elif lcunit in ('g', 'gb', 'gib'):
564
    value *= 1024
565

    
566
  elif lcunit in ('t', 'tb', 'tib'):
567
    value *= 1024 * 1024
568

    
569
  else:
570
    raise errors.UnitParseError("Unknown unit: %s" % unit)
571

    
572
  # Make sure we round up
573
  if int(value) < value:
574
    value += 1
575

    
576
  # Round up to the next multiple of 4
577
  value = int(value)
578
  if value % 4:
579
    value += 4 - value % 4
580

    
581
  return value
582

    
583

    
584
def AddAuthorizedKey(file_name, key):
585
  """Adds an SSH public key to an authorized_keys file.
586

587
  Args:
588
    file_name: Path to authorized_keys file
589
    key: String containing key
590
  """
591
  key_fields = key.split()
592

    
593
  f = open(file_name, 'a+')
594
  try:
595
    nl = True
596
    for line in f:
597
      # Ignore whitespace changes
598
      if line.split() == key_fields:
599
        break
600
      nl = line.endswith('\n')
601
    else:
602
      if not nl:
603
        f.write("\n")
604
      f.write(key.rstrip('\r\n'))
605
      f.write("\n")
606
      f.flush()
607
  finally:
608
    f.close()
609

    
610

    
611
def RemoveAuthorizedKey(file_name, key):
612
  """Removes an SSH public key from an authorized_keys file.
613

614
  Args:
615
    file_name: Path to authorized_keys file
616
    key: String containing key
617
  """
618
  key_fields = key.split()
619

    
620
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
621
  try:
622
    out = os.fdopen(fd, 'w')
623
    try:
624
      f = open(file_name, 'r')
625
      try:
626
        for line in f:
627
          # Ignore whitespace changes while comparing lines
628
          if line.split() != key_fields:
629
            out.write(line)
630

    
631
        out.flush()
632
        os.rename(tmpname, file_name)
633
      finally:
634
        f.close()
635
    finally:
636
      out.close()
637
  except:
638
    RemoveFile(tmpname)
639
    raise
640

    
641

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

645
  """
646
  # Ensure aliases are unique
647
  aliases = UniqueSequence([hostname] + aliases)[1:]
648

    
649
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
650
  try:
651
    out = os.fdopen(fd, 'w')
652
    try:
653
      f = open(file_name, 'r')
654
      try:
655
        written = False
656
        for line in f:
657
          fields = line.split()
658
          if fields and not fields[0].startswith('#') and ip == fields[0]:
659
            continue
660
          out.write(line)
661

    
662
        out.write("%s\t%s" % (ip, hostname))
663
        if aliases:
664
          out.write(" %s" % ' '.join(aliases))
665
        out.write('\n')
666

    
667
        out.flush()
668
        os.fsync(out)
669
        os.rename(tmpname, file_name)
670
      finally:
671
        f.close()
672
    finally:
673
      out.close()
674
  except:
675
    RemoveFile(tmpname)
676
    raise
677

    
678

    
679
def AddHostToEtcHosts(hostname):
680
  """Wrapper around SetEtcHostsEntry.
681

682
  """
683
  hi = HostInfo(name=hostname)
684
  SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])
685

    
686

    
687
def RemoveEtcHostsEntry(file_name, hostname):
688
  """Removes a hostname from /etc/hosts.
689

690
  IP addresses without names are removed from the file.
691
  """
692
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
693
  try:
694
    out = os.fdopen(fd, 'w')
695
    try:
696
      f = open(file_name, 'r')
697
      try:
698
        for line in f:
699
          fields = line.split()
700
          if len(fields) > 1 and not fields[0].startswith('#'):
701
            names = fields[1:]
702
            if hostname in names:
703
              while hostname in names:
704
                names.remove(hostname)
705
              if names:
706
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
707
              continue
708

    
709
          out.write(line)
710

    
711
        out.flush()
712
        os.fsync(out)
713
        os.rename(tmpname, file_name)
714
      finally:
715
        f.close()
716
    finally:
717
      out.close()
718
  except:
719
    RemoveFile(tmpname)
720
    raise
721

    
722

    
723
def RemoveHostFromEtcHosts(hostname):
724
  """Wrapper around RemoveEtcHostsEntry.
725

726
  """
727
  hi = HostInfo(name=hostname)
728
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.name)
729
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.ShortName())
730

    
731

    
732
def CreateBackup(file_name):
733
  """Creates a backup of a file.
734

735
  Returns: the path to the newly created backup file.
736

737
  """
738
  if not os.path.isfile(file_name):
739
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
740
                                file_name)
741

    
742
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
743
  dir_name = os.path.dirname(file_name)
744

    
745
  fsrc = open(file_name, 'rb')
746
  try:
747
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
748
    fdst = os.fdopen(fd, 'wb')
749
    try:
750
      shutil.copyfileobj(fsrc, fdst)
751
    finally:
752
      fdst.close()
753
  finally:
754
    fsrc.close()
755

    
756
  return backup_name
757

    
758

    
759
def ShellQuote(value):
760
  """Quotes shell argument according to POSIX.
761

762
  """
763
  if _re_shell_unquoted.match(value):
764
    return value
765
  else:
766
    return "'%s'" % value.replace("'", "'\\''")
767

    
768

    
769
def ShellQuoteArgs(args):
770
  """Quotes all given shell arguments and concatenates using spaces.
771

772
  """
773
  return ' '.join([ShellQuote(i) for i in args])
774

    
775

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

779
  Try to do a TCP connect(2) from an optional source IP to the
780
  specified target IP and the specified target port. If the optional
781
  parameter live_port_needed is set to true, requires the remote end
782
  to accept the connection. The timeout is specified in seconds and
783
  defaults to 10 seconds. If the source optional argument is not
784
  passed, the source address selection is left to the kernel,
785
  otherwise we try to connect using the passed address (failures to
786
  bind other than EADDRNOTAVAIL will be ignored).
787

788
  """
789
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
790

    
791
  sucess = False
792

    
793
  if source is not None:
794
    try:
795
      sock.bind((source, 0))
796
    except socket.error, (errcode, errstring):
797
      if errcode == errno.EADDRNOTAVAIL:
798
        success = False
799

    
800
  sock.settimeout(timeout)
801

    
802
  try:
803
    sock.connect((target, port))
804
    sock.close()
805
    success = True
806
  except socket.timeout:
807
    success = False
808
  except socket.error, (errcode, errstring):
809
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
810

    
811
  return success
812

    
813

    
814
def ListVisibleFiles(path):
815
  """Returns a list of all visible files in a directory.
816

817
  """
818
  files = [i for i in os.listdir(path) if not i.startswith(".")]
819
  files.sort()
820
  return files
821

    
822

    
823
def GetHomeDir(user, default=None):
824
  """Try to get the homedir of the given user.
825

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

830
  """
831
  try:
832
    if isinstance(user, basestring):
833
      result = pwd.getpwnam(user)
834
    elif isinstance(user, (int, long)):
835
      result = pwd.getpwuid(user)
836
    else:
837
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
838
                                   type(user))
839
  except KeyError:
840
    return default
841
  return result.pw_dir
842

    
843

    
844
def NewUUID():
845
  """Returns a random UUID.
846

847
  """
848
  f = open("/proc/sys/kernel/random/uuid", "r")
849
  try:
850
    return f.read(128).rstrip("\n")
851
  finally:
852
    f.close()
853

    
854

    
855
def WriteFile(file_name, fn=None, data=None,
856
              mode=None, uid=-1, gid=-1,
857
              atime=None, mtime=None, close=True,
858
              dry_run=False, backup=False,
859
              prewrite=None, postwrite=None):
860
  """(Over)write a file atomically.
861

862
  The file_name and either fn (a function taking one argument, the
863
  file descriptor, and which should write the data to it) or data (the
864
  contents of the file) must be passed. The other arguments are
865
  optional and allow setting the file mode, owner and group, and the
866
  mtime/atime of the file.
867

868
  If the function doesn't raise an exception, it has succeeded and the
869
  target file has the new contents. If the file has raised an
870
  exception, an existing target file should be unmodified and the
871
  temporary file should be removed.
872

873
  Args:
874
    file_name: New filename
875
    fn: Content writing function, called with file descriptor as parameter
876
    data: Content as string
877
    mode: File mode
878
    uid: Owner
879
    gid: Group
880
    atime: Access time
881
    mtime: Modification time
882
    close: Whether to close file after writing it
883
    prewrite: Function object called before writing content
884
    postwrite: Function object called after writing content
885

886
  Returns:
887
    None if "close" parameter evaluates to True, otherwise file descriptor.
888

889
  """
890
  if not os.path.isabs(file_name):
891
    raise errors.ProgrammerError("Path passed to WriteFile is not"
892
                                 " absolute: '%s'" % file_name)
893

    
894
  if [fn, data].count(None) != 1:
895
    raise errors.ProgrammerError("fn or data required")
896

    
897
  if [atime, mtime].count(None) == 1:
898
    raise errors.ProgrammerError("Both atime and mtime must be either"
899
                                 " set or None")
900

    
901
  if backup and not dry_run and os.path.isfile(file_name):
902
    CreateBackup(file_name)
903

    
904
  dir_name, base_name = os.path.split(file_name)
905
  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
906
  # here we need to make sure we remove the temp file, if any error
907
  # leaves it in place
908
  try:
909
    if uid != -1 or gid != -1:
910
      os.chown(new_name, uid, gid)
911
    if mode:
912
      os.chmod(new_name, mode)
913
    if callable(prewrite):
914
      prewrite(fd)
915
    if data is not None:
916
      os.write(fd, data)
917
    else:
918
      fn(fd)
919
    if callable(postwrite):
920
      postwrite(fd)
921
    os.fsync(fd)
922
    if atime is not None and mtime is not None:
923
      os.utime(new_name, (atime, mtime))
924
    if not dry_run:
925
      os.rename(new_name, file_name)
926
  finally:
927
    if close:
928
      os.close(fd)
929
      result = None
930
    else:
931
      result = fd
932
    RemoveFile(new_name)
933

    
934
  return result
935

    
936

    
937
def FirstFree(seq, base=0):
938
  """Returns the first non-existing integer from seq.
939

940
  The seq argument should be a sorted list of positive integers. The
941
  first time the index of an element is smaller than the element
942
  value, the index will be returned.
943

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

947
  Example: [0, 1, 3] will return 2.
948

949
  """
950
  for idx, elem in enumerate(seq):
951
    assert elem >= base, "Passed element is higher than base offset"
952
    if elem > idx + base:
953
      # idx is not used
954
      return idx + base
955
  return None
956

    
957

    
958
def all(seq, pred=bool):
959
  "Returns True if pred(x) is True for every element in the iterable"
960
  for elem in itertools.ifilterfalse(pred, seq):
961
    return False
962
  return True
963

    
964

    
965
def any(seq, pred=bool):
966
  "Returns True if pred(x) is True for at least one element in the iterable"
967
  for elem in itertools.ifilter(pred, seq):
968
    return True
969
  return False
970

    
971

    
972
def UniqueSequence(seq):
973
  """Returns a list with unique elements.
974

975
  Element order is preserved.
976
  """
977
  seen = set()
978
  return [i for i in seq if i not in seen and not seen.add(i)]
979

    
980

    
981
def IsValidMac(mac):
982
  """Predicate to check if a MAC address is valid.
983

984
  Checks wether the supplied MAC address is formally correct, only
985
  accepts colon separated format.
986
  """
987
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
988
  return mac_check.match(mac) is not None
989

    
990

    
991
def TestDelay(duration):
992
  """Sleep for a fixed amount of time.
993

994
  """
995
  if duration < 0:
996
    return False
997
  time.sleep(duration)
998
  return True
999

    
1000

    
1001
def Daemonize(logfile, noclose_fds=None):
1002
  """Daemonize the current process.
1003

1004
  This detaches the current process from the controlling terminal and
1005
  runs it in the background as a daemon.
1006

1007
  """
1008
  UMASK = 077
1009
  WORKDIR = "/"
1010
  # Default maximum for the number of available file descriptors.
1011
  if 'SC_OPEN_MAX' in os.sysconf_names:
1012
    try:
1013
      MAXFD = os.sysconf('SC_OPEN_MAX')
1014
      if MAXFD < 0:
1015
        MAXFD = 1024
1016
    except OSError:
1017
      MAXFD = 1024
1018
  else:
1019
    MAXFD = 1024
1020

    
1021
  # this might fail
1022
  pid = os.fork()
1023
  if (pid == 0):  # The first child.
1024
    os.setsid()
1025
    # this might fail
1026
    pid = os.fork() # Fork a second child.
1027
    if (pid == 0):  # The second child.
1028
      os.chdir(WORKDIR)
1029
      os.umask(UMASK)
1030
    else:
1031
      # exit() or _exit()?  See below.
1032
      os._exit(0) # Exit parent (the first child) of the second child.
1033
  else:
1034
    os._exit(0) # Exit parent of the first child.
1035
  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1036
  if (maxfd == resource.RLIM_INFINITY):
1037
    maxfd = MAXFD
1038

    
1039
  # Iterate through and close all file descriptors.
1040
  for fd in range(0, maxfd):
1041
    if noclose_fds and fd in noclose_fds:
1042
      continue
1043
    try:
1044
      os.close(fd)
1045
    except OSError: # ERROR, fd wasn't open to begin with (ignored)
1046
      pass
1047
  os.open(logfile, os.O_RDWR|os.O_CREAT|os.O_APPEND, 0600)
1048
  # Duplicate standard input to standard output and standard error.
1049
  os.dup2(0, 1)     # standard output (1)
1050
  os.dup2(0, 2)     # standard error (2)
1051
  return 0
1052

    
1053

    
1054
def FindFile(name, search_path, test=os.path.exists):
1055
  """Look for a filesystem object in a given path.
1056

1057
  This is an abstract method to search for filesystem object (files,
1058
  dirs) under a given search path.
1059

1060
  Args:
1061
    - name: the name to look for
1062
    - search_path: list of directory names
1063
    - test: the test which the full path must satisfy
1064
      (defaults to os.path.exists)
1065

1066
  Returns:
1067
    - full path to the item if found
1068
    - None otherwise
1069

1070
  """
1071
  for dir_name in search_path:
1072
    item_name = os.path.sep.join([dir_name, name])
1073
    if test(item_name):
1074
      return item_name
1075
  return None
1076

    
1077

    
1078
def CheckVolumeGroupSize(vglist, vgname, minsize):
1079
  """Checks if the volume group list is valid.
1080

1081
  A non-None return value means there's an error, and the return value
1082
  is the error message.
1083

1084
  """
1085
  vgsize = vglist.get(vgname, None)
1086
  if vgsize is None:
1087
    return "volume group '%s' missing" % vgname
1088
  elif vgsize < minsize:
1089
    return ("volume group '%s' too small (%s MiB required, %d MiB found)" %
1090
            (vgname, minsize, vgsize))
1091
  return None
1092

    
1093

    
1094
def LockedMethod(fn):
1095
  """Synchronized object access decorator.
1096

1097
  This decorator is intended to protect access to an object using the
1098
  object's own lock which is hardcoded to '_lock'.
1099

1100
  """
1101
  def wrapper(self, *args, **kwargs):
1102
    assert hasattr(self, '_lock')
1103
    lock = self._lock
1104
    lock.acquire()
1105
    try:
1106
      result = fn(self, *args, **kwargs)
1107
    finally:
1108
      lock.release()
1109
    return result
1110
  return wrapper
1111

    
1112

    
1113
def LockFile(fd):
1114
  """Locks a file using POSIX locks.
1115

1116
  """
1117
  try:
1118
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
1119
  except IOError, err:
1120
    if err.errno == errno.EAGAIN:
1121
      raise errors.LockError("File already locked")
1122
    raise
1123

    
1124

    
1125
class SignalHandler(object):
1126
  """Generic signal handler class.
1127

1128
  It automatically restores the original handler when deconstructed or when
1129
  Reset() is called. You can either pass your own handler function in or query
1130
  the "called" attribute to detect whether the signal was sent.
1131

1132
  """
1133
  def __init__(self, signum):
1134
    """Constructs a new SignalHandler instance.
1135

1136
    @param signum: Single signal number or set of signal numbers
1137

1138
    """
1139
    if isinstance(signum, (int, long)):
1140
      self.signum = set([signum])
1141
    else:
1142
      self.signum = set(signum)
1143

    
1144
    self.called = False
1145

    
1146
    self._previous = {}
1147
    try:
1148
      for signum in self.signum:
1149
        # Setup handler
1150
        prev_handler = signal.signal(signum, self._HandleSignal)
1151
        try:
1152
          self._previous[signum] = prev_handler
1153
        except:
1154
          # Restore previous handler
1155
          signal.signal(signum, prev_handler)
1156
          raise
1157
    except:
1158
      # Reset all handlers
1159
      self.Reset()
1160
      # Here we have a race condition: a handler may have already been called,
1161
      # but there's not much we can do about it at this point.
1162
      raise
1163

    
1164
  def __del__(self):
1165
    self.Reset()
1166

    
1167
  def Reset(self):
1168
    """Restore previous handler.
1169

1170
    """
1171
    for signum, prev_handler in self._previous.items():
1172
      signal.signal(signum, prev_handler)
1173
      # If successful, remove from dict
1174
      del self._previous[signum]
1175

    
1176
  def Clear(self):
1177
    """Unsets "called" flag.
1178

1179
    This function can be used in case a signal may arrive several times.
1180

1181
    """
1182
    self.called = False
1183

    
1184
  def _HandleSignal(self, signum, frame):
1185
    """Actual signal handling function.
1186

1187
    """
1188
    # This is not nice and not absolutely atomic, but it appears to be the only
1189
    # solution in Python -- there are no atomic types.
1190
    self.called = True