Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ 099c52ad

History | View | Annotate | Download (63.3 kB)

1
#
2
#
3

    
4
# Copyright (C) 2006, 2007 Google Inc.
5
#
6
# This program is free software; you can redistribute it and/or modify
7
# it under the terms of the GNU General Public License as published by
8
# the Free Software Foundation; either version 2 of the License, or
9
# (at your option) any later version.
10
#
11
# This program is distributed in the hope that it will be useful, but
12
# WITHOUT ANY WARRANTY; without even the implied warranty of
13
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14
# General Public License for more details.
15
#
16
# You should have received a copy of the GNU General Public License
17
# along with this program; if not, write to the Free Software
18
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
19
# 02110-1301, USA.
20

    
21

    
22
"""Ganeti utility module.
23

24
This module holds functions that can be used in both daemons (all) and
25
the command line scripts.
26

27
"""
28

    
29

    
30
import os
31
import time
32
import subprocess
33
import re
34
import socket
35
import tempfile
36
import shutil
37
import errno
38
import pwd
39
import itertools
40
import select
41
import fcntl
42
import resource
43
import logging
44
import signal
45

    
46
from cStringIO import StringIO
47

    
48
try:
49
  from hashlib import sha1
50
except ImportError:
51
  import sha
52
  sha1 = sha.new
53

    
54
from ganeti import errors
55
from ganeti import constants
56

    
57

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

    
61
debug_locks = False
62

    
63
#: when set to True, L{RunCmd} is disabled
64
no_fork = False
65

    
66
_RANDOM_UUID_FILE = "/proc/sys/kernel/random/uuid"
67

    
68

    
69
class RunResult(object):
70
  """Holds the result of running external programs.
71

72
  @type exit_code: int
73
  @ivar exit_code: the exit code of the program, or None (if the program
74
      didn't exit())
75
  @type signal: int or None
76
  @ivar signal: the signal that caused the program to finish, or None
77
      (if the program wasn't terminated by a signal)
78
  @type stdout: str
79
  @ivar stdout: the standard output of the program
80
  @type stderr: str
81
  @ivar stderr: the standard error of the program
82
  @type failed: boolean
83
  @ivar failed: True in case the program was
84
      terminated by a signal or exited with a non-zero exit code
85
  @ivar fail_reason: a string detailing the termination reason
86

87
  """
88
  __slots__ = ["exit_code", "signal", "stdout", "stderr",
89
               "failed", "fail_reason", "cmd"]
90

    
91

    
92
  def __init__(self, exit_code, signal_, stdout, stderr, cmd):
93
    self.cmd = cmd
94
    self.exit_code = exit_code
95
    self.signal = signal_
96
    self.stdout = stdout
97
    self.stderr = stderr
98
    self.failed = (signal_ is not None or exit_code != 0)
99

    
100
    if self.signal is not None:
101
      self.fail_reason = "terminated by signal %s" % self.signal
102
    elif self.exit_code is not None:
103
      self.fail_reason = "exited with exit code %s" % self.exit_code
104
    else:
105
      self.fail_reason = "unable to determine termination reason"
106

    
107
    if self.failed:
108
      logging.debug("Command '%s' failed (%s); output: %s",
109
                    self.cmd, self.fail_reason, self.output)
110

    
111
  def _GetOutput(self):
112
    """Returns the combined stdout and stderr for easier usage.
113

114
    """
115
    return self.stdout + self.stderr
116

    
117
  output = property(_GetOutput, None, None, "Return full output")
118

    
119

    
120
def RunCmd(cmd, env=None, output=None, cwd='/'):
121
  """Execute a (shell) command.
122

123
  The command should not read from its standard input, as it will be
124
  closed.
125

126
  @type  cmd: string or list
127
  @param cmd: Command to run
128
  @type env: dict
129
  @param env: Additional environment
130
  @type output: str
131
  @param output: if desired, the output of the command can be
132
      saved in a file instead of the RunResult instance; this
133
      parameter denotes the file name (if not None)
134
  @type cwd: string
135
  @param cwd: if specified, will be used as the working
136
      directory for the command; the default will be /
137
  @rtype: L{RunResult}
138
  @return: RunResult instance
139
  @raise errors.ProgrammerError: if we call this when forks are disabled
140

141
  """
142
  if no_fork:
143
    raise errors.ProgrammerError("utils.RunCmd() called with fork() disabled")
144

    
145
  if isinstance(cmd, list):
146
    cmd = [str(val) for val in cmd]
147
    strcmd = " ".join(cmd)
148
    shell = False
149
  else:
150
    strcmd = cmd
151
    shell = True
152
  logging.debug("RunCmd '%s'", strcmd)
153

    
154
  cmd_env = os.environ.copy()
155
  cmd_env["LC_ALL"] = "C"
156
  if env is not None:
157
    cmd_env.update(env)
158

    
159
  try:
160
    if output is None:
161
      out, err, status = _RunCmdPipe(cmd, cmd_env, shell, cwd)
162
    else:
163
      status = _RunCmdFile(cmd, cmd_env, shell, output, cwd)
164
      out = err = ""
165
  except OSError, err:
166
    if err.errno == errno.ENOENT:
167
      raise errors.OpExecError("Can't execute '%s': not found (%s)" %
168
                               (strcmd, err))
169
    else:
170
      raise
171

    
172
  if status >= 0:
173
    exitcode = status
174
    signal_ = None
175
  else:
176
    exitcode = None
177
    signal_ = -status
178

    
179
  return RunResult(exitcode, signal_, out, err, strcmd)
180

    
181

    
182
def _RunCmdPipe(cmd, env, via_shell, cwd):
183
  """Run a command and return its output.
184

185
  @type  cmd: string or list
186
  @param cmd: Command to run
187
  @type env: dict
188
  @param env: The environment to use
189
  @type via_shell: bool
190
  @param via_shell: if we should run via the shell
191
  @type cwd: string
192
  @param cwd: the working directory for the program
193
  @rtype: tuple
194
  @return: (out, err, status)
195

196
  """
197
  poller = select.poll()
198
  child = subprocess.Popen(cmd, shell=via_shell,
199
                           stderr=subprocess.PIPE,
200
                           stdout=subprocess.PIPE,
201
                           stdin=subprocess.PIPE,
202
                           close_fds=True, env=env,
203
                           cwd=cwd)
204

    
205
  child.stdin.close()
206
  poller.register(child.stdout, select.POLLIN)
207
  poller.register(child.stderr, select.POLLIN)
208
  out = StringIO()
209
  err = StringIO()
210
  fdmap = {
211
    child.stdout.fileno(): (out, child.stdout),
212
    child.stderr.fileno(): (err, child.stderr),
213
    }
214
  for fd in fdmap:
215
    status = fcntl.fcntl(fd, fcntl.F_GETFL)
216
    fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
217

    
218
  while fdmap:
219
    try:
220
      pollresult = poller.poll()
221
    except EnvironmentError, eerr:
222
      if eerr.errno == errno.EINTR:
223
        continue
224
      raise
225
    except select.error, serr:
226
      if serr[0] == errno.EINTR:
227
        continue
228
      raise
229

    
230
    for fd, event in pollresult:
231
      if event & select.POLLIN or event & select.POLLPRI:
232
        data = fdmap[fd][1].read()
233
        # no data from read signifies EOF (the same as POLLHUP)
234
        if not data:
235
          poller.unregister(fd)
236
          del fdmap[fd]
237
          continue
238
        fdmap[fd][0].write(data)
239
      if (event & select.POLLNVAL or event & select.POLLHUP or
240
          event & select.POLLERR):
241
        poller.unregister(fd)
242
        del fdmap[fd]
243

    
244
  out = out.getvalue()
245
  err = err.getvalue()
246

    
247
  status = child.wait()
248
  return out, err, status
249

    
250

    
251
def _RunCmdFile(cmd, env, via_shell, output, cwd):
252
  """Run a command and save its output to a file.
253

254
  @type  cmd: string or list
255
  @param cmd: Command to run
256
  @type env: dict
257
  @param env: The environment to use
258
  @type via_shell: bool
259
  @param via_shell: if we should run via the shell
260
  @type output: str
261
  @param output: the filename in which to save the output
262
  @type cwd: string
263
  @param cwd: the working directory for the program
264
  @rtype: int
265
  @return: the exit status
266

267
  """
268
  fh = open(output, "a")
269
  try:
270
    child = subprocess.Popen(cmd, shell=via_shell,
271
                             stderr=subprocess.STDOUT,
272
                             stdout=fh,
273
                             stdin=subprocess.PIPE,
274
                             close_fds=True, env=env,
275
                             cwd=cwd)
276

    
277
    child.stdin.close()
278
    status = child.wait()
279
  finally:
280
    fh.close()
281
  return status
282

    
283

    
284
def RemoveFile(filename):
285
  """Remove a file ignoring some errors.
286

287
  Remove a file, ignoring non-existing ones or directories. Other
288
  errors are passed.
289

290
  @type filename: str
291
  @param filename: the file to be removed
292

293
  """
294
  try:
295
    os.unlink(filename)
296
  except OSError, err:
297
    if err.errno not in (errno.ENOENT, errno.EISDIR):
298
      raise
299

    
300

    
301
def RenameFile(old, new, mkdir=False, mkdir_mode=0750):
302
  """Renames a file.
303

304
  @type old: string
305
  @param old: Original path
306
  @type new: string
307
  @param new: New path
308
  @type mkdir: bool
309
  @param mkdir: Whether to create target directory if it doesn't exist
310
  @type mkdir_mode: int
311
  @param mkdir_mode: Mode for newly created directories
312

313
  """
314
  try:
315
    return os.rename(old, new)
316
  except OSError, err:
317
    # In at least one use case of this function, the job queue, directory
318
    # creation is very rare. Checking for the directory before renaming is not
319
    # as efficient.
320
    if mkdir and err.errno == errno.ENOENT:
321
      # Create directory and try again
322
      os.makedirs(os.path.dirname(new), mkdir_mode)
323
      return os.rename(old, new)
324
    raise
325

    
326

    
327
def _FingerprintFile(filename):
328
  """Compute the fingerprint of a file.
329

330
  If the file does not exist, a None will be returned
331
  instead.
332

333
  @type filename: str
334
  @param filename: the filename to checksum
335
  @rtype: str
336
  @return: the hex digest of the sha checksum of the contents
337
      of the file
338

339
  """
340
  if not (os.path.exists(filename) and os.path.isfile(filename)):
341
    return None
342

    
343
  f = open(filename)
344

    
345
  fp = sha1()
346
  while True:
347
    data = f.read(4096)
348
    if not data:
349
      break
350

    
351
    fp.update(data)
352

    
353
  return fp.hexdigest()
354

    
355

    
356
def FingerprintFiles(files):
357
  """Compute fingerprints for a list of files.
358

359
  @type files: list
360
  @param files: the list of filename to fingerprint
361
  @rtype: dict
362
  @return: a dictionary filename: fingerprint, holding only
363
      existing files
364

365
  """
366
  ret = {}
367

    
368
  for filename in files:
369
    cksum = _FingerprintFile(filename)
370
    if cksum:
371
      ret[filename] = cksum
372

    
373
  return ret
374

    
375

    
376
def ForceDictType(target, key_types, allowed_values=None):
377
  """Force the values of a dict to have certain types.
378

379
  @type target: dict
380
  @param target: the dict to update
381
  @type key_types: dict
382
  @param key_types: dict mapping target dict keys to types
383
                    in constants.ENFORCEABLE_TYPES
384
  @type allowed_values: list
385
  @keyword allowed_values: list of specially allowed values
386

387
  """
388
  if allowed_values is None:
389
    allowed_values = []
390

    
391
  if not isinstance(target, dict):
392
    msg = "Expected dictionary, got '%s'" % target
393
    raise errors.TypeEnforcementError(msg)
394

    
395
  for key in target:
396
    if key not in key_types:
397
      msg = "Unknown key '%s'" % key
398
      raise errors.TypeEnforcementError(msg)
399

    
400
    if target[key] in allowed_values:
401
      continue
402

    
403
    ktype = key_types[key]
404
    if ktype not in constants.ENFORCEABLE_TYPES:
405
      msg = "'%s' has non-enforceable type %s" % (key, ktype)
406
      raise errors.ProgrammerError(msg)
407

    
408
    if ktype == constants.VTYPE_STRING:
409
      if not isinstance(target[key], basestring):
410
        if isinstance(target[key], bool) and not target[key]:
411
          target[key] = ''
412
        else:
413
          msg = "'%s' (value %s) is not a valid string" % (key, target[key])
414
          raise errors.TypeEnforcementError(msg)
415
    elif ktype == constants.VTYPE_BOOL:
416
      if isinstance(target[key], basestring) and target[key]:
417
        if target[key].lower() == constants.VALUE_FALSE:
418
          target[key] = False
419
        elif target[key].lower() == constants.VALUE_TRUE:
420
          target[key] = True
421
        else:
422
          msg = "'%s' (value %s) is not a valid boolean" % (key, target[key])
423
          raise errors.TypeEnforcementError(msg)
424
      elif target[key]:
425
        target[key] = True
426
      else:
427
        target[key] = False
428
    elif ktype == constants.VTYPE_SIZE:
429
      try:
430
        target[key] = ParseUnit(target[key])
431
      except errors.UnitParseError, err:
432
        msg = "'%s' (value %s) is not a valid size. error: %s" % \
433
              (key, target[key], err)
434
        raise errors.TypeEnforcementError(msg)
435
    elif ktype == constants.VTYPE_INT:
436
      try:
437
        target[key] = int(target[key])
438
      except (ValueError, TypeError):
439
        msg = "'%s' (value %s) is not a valid integer" % (key, target[key])
440
        raise errors.TypeEnforcementError(msg)
441

    
442

    
443
def IsProcessAlive(pid):
444
  """Check if a given pid exists on the system.
445

446
  @note: zombie status is not handled, so zombie processes
447
      will be returned as alive
448
  @type pid: int
449
  @param pid: the process ID to check
450
  @rtype: boolean
451
  @return: True if the process exists
452

453
  """
454
  if pid <= 0:
455
    return False
456

    
457
  try:
458
    os.stat("/proc/%d/status" % pid)
459
    return True
460
  except EnvironmentError, err:
461
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
462
      return False
463
    raise
464

    
465

    
466
def ReadPidFile(pidfile):
467
  """Read a pid from a file.
468

469
  @type  pidfile: string
470
  @param pidfile: path to the file containing the pid
471
  @rtype: int
472
  @return: The process id, if the file exists and contains a valid PID,
473
           otherwise 0
474

475
  """
476
  try:
477
    raw_data = ReadFile(pidfile)
478
  except EnvironmentError, err:
479
    if err.errno != errno.ENOENT:
480
      logging.exception("Can't read pid file")
481
    return 0
482

    
483
  try:
484
    pid = int(raw_data)
485
  except ValueError, err:
486
    logging.info("Can't parse pid file contents", exc_info=True)
487
    return 0
488

    
489
  return pid
490

    
491

    
492
def MatchNameComponent(key, name_list, case_sensitive=True):
493
  """Try to match a name against a list.
494

495
  This function will try to match a name like test1 against a list
496
  like C{['test1.example.com', 'test2.example.com', ...]}. Against
497
  this list, I{'test1'} as well as I{'test1.example'} will match, but
498
  not I{'test1.ex'}. A multiple match will be considered as no match
499
  at all (e.g. I{'test1'} against C{['test1.example.com',
500
  'test1.example.org']}), except when the key fully matches an entry
501
  (e.g. I{'test1'} against C{['test1', 'test1.example.com']}).
502

503
  @type key: str
504
  @param key: the name to be searched
505
  @type name_list: list
506
  @param name_list: the list of strings against which to search the key
507
  @type case_sensitive: boolean
508
  @param case_sensitive: whether to provide a case-sensitive match
509

510
  @rtype: None or str
511
  @return: None if there is no match I{or} if there are multiple matches,
512
      otherwise the element from the list which matches
513

514
  """
515
  if key in name_list:
516
    return key
517

    
518
  re_flags = 0
519
  if not case_sensitive:
520
    re_flags |= re.IGNORECASE
521
    key = key.upper()
522
  mo = re.compile("^%s(\..*)?$" % re.escape(key), re_flags)
523
  names_filtered = []
524
  string_matches = []
525
  for name in name_list:
526
    if mo.match(name) is not None:
527
      names_filtered.append(name)
528
      if not case_sensitive and key == name.upper():
529
        string_matches.append(name)
530

    
531
  if len(string_matches) == 1:
532
    return string_matches[0]
533
  if len(names_filtered) == 1:
534
    return names_filtered[0]
535
  return None
536

    
537

    
538
class HostInfo:
539
  """Class implementing resolver and hostname functionality
540

541
  """
542
  def __init__(self, name=None):
543
    """Initialize the host name object.
544

545
    If the name argument is not passed, it will use this system's
546
    name.
547

548
    """
549
    if name is None:
550
      name = self.SysName()
551

    
552
    self.query = name
553
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
554
    self.ip = self.ipaddrs[0]
555

    
556
  def ShortName(self):
557
    """Returns the hostname without domain.
558

559
    """
560
    return self.name.split('.')[0]
561

    
562
  @staticmethod
563
  def SysName():
564
    """Return the current system's name.
565

566
    This is simply a wrapper over C{socket.gethostname()}.
567

568
    """
569
    return socket.gethostname()
570

    
571
  @staticmethod
572
  def LookupHostname(hostname):
573
    """Look up hostname
574

575
    @type hostname: str
576
    @param hostname: hostname to look up
577

578
    @rtype: tuple
579
    @return: a tuple (name, aliases, ipaddrs) as returned by
580
        C{socket.gethostbyname_ex}
581
    @raise errors.ResolverError: in case of errors in resolving
582

583
    """
584
    try:
585
      result = socket.gethostbyname_ex(hostname)
586
    except socket.gaierror, err:
587
      # hostname not found in DNS
588
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
589

    
590
    return result
591

    
592

    
593
def ListVolumeGroups():
594
  """List volume groups and their size
595

596
  @rtype: dict
597
  @return:
598
       Dictionary with keys volume name and values
599
       the size of the volume
600

601
  """
602
  command = "vgs --noheadings --units m --nosuffix -o name,size"
603
  result = RunCmd(command)
604
  retval = {}
605
  if result.failed:
606
    return retval
607

    
608
  for line in result.stdout.splitlines():
609
    try:
610
      name, size = line.split()
611
      size = int(float(size))
612
    except (IndexError, ValueError), err:
613
      logging.error("Invalid output from vgs (%s): %s", err, line)
614
      continue
615

    
616
    retval[name] = size
617

    
618
  return retval
619

    
620

    
621
def BridgeExists(bridge):
622
  """Check whether the given bridge exists in the system
623

624
  @type bridge: str
625
  @param bridge: the bridge name to check
626
  @rtype: boolean
627
  @return: True if it does
628

629
  """
630
  return os.path.isdir("/sys/class/net/%s/bridge" % bridge)
631

    
632

    
633
def NiceSort(name_list):
634
  """Sort a list of strings based on digit and non-digit groupings.
635

636
  Given a list of names C{['a1', 'a10', 'a11', 'a2']} this function
637
  will sort the list in the logical order C{['a1', 'a2', 'a10',
638
  'a11']}.
639

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

644
  @type name_list: list
645
  @param name_list: the names to be sorted
646
  @rtype: list
647
  @return: a copy of the name list sorted with our algorithm
648

649
  """
650
  _SORTER_BASE = "(\D+|\d+)"
651
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
652
                                                  _SORTER_BASE, _SORTER_BASE,
653
                                                  _SORTER_BASE, _SORTER_BASE,
654
                                                  _SORTER_BASE, _SORTER_BASE)
655
  _SORTER_RE = re.compile(_SORTER_FULL)
656
  _SORTER_NODIGIT = re.compile("^\D*$")
657
  def _TryInt(val):
658
    """Attempts to convert a variable to integer."""
659
    if val is None or _SORTER_NODIGIT.match(val):
660
      return val
661
    rval = int(val)
662
    return rval
663

    
664
  to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
665
             for name in name_list]
666
  to_sort.sort()
667
  return [tup[1] for tup in to_sort]
668

    
669

    
670
def TryConvert(fn, val):
671
  """Try to convert a value ignoring errors.
672

673
  This function tries to apply function I{fn} to I{val}. If no
674
  C{ValueError} or C{TypeError} exceptions are raised, it will return
675
  the result, else it will return the original value. Any other
676
  exceptions are propagated to the caller.
677

678
  @type fn: callable
679
  @param fn: function to apply to the value
680
  @param val: the value to be converted
681
  @return: The converted value if the conversion was successful,
682
      otherwise the original value.
683

684
  """
685
  try:
686
    nv = fn(val)
687
  except (ValueError, TypeError):
688
    nv = val
689
  return nv
690

    
691

    
692
def IsValidIP(ip):
693
  """Verifies the syntax of an IPv4 address.
694

695
  This function checks if the IPv4 address passes is valid or not based
696
  on syntax (not IP range, class calculations, etc.).
697

698
  @type ip: str
699
  @param ip: the address to be checked
700
  @rtype: a regular expression match object
701
  @return: a regular expression match object, or None if the
702
      address is not valid
703

704
  """
705
  unit = "(0|[1-9]\d{0,2})"
706
  #TODO: convert and return only boolean
707
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
708

    
709

    
710
def IsValidShellParam(word):
711
  """Verifies is the given word is safe from the shell's p.o.v.
712

713
  This means that we can pass this to a command via the shell and be
714
  sure that it doesn't alter the command line and is passed as such to
715
  the actual command.
716

717
  Note that we are overly restrictive here, in order to be on the safe
718
  side.
719

720
  @type word: str
721
  @param word: the word to check
722
  @rtype: boolean
723
  @return: True if the word is 'safe'
724

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

    
728

    
729
def BuildShellCmd(template, *args):
730
  """Build a safe shell command line from the given arguments.
731

732
  This function will check all arguments in the args list so that they
733
  are valid shell parameters (i.e. they don't contain shell
734
  metacharacters). If everything is ok, it will return the result of
735
  template % args.
736

737
  @type template: str
738
  @param template: the string holding the template for the
739
      string formatting
740
  @rtype: str
741
  @return: the expanded command line
742

743
  """
744
  for word in args:
745
    if not IsValidShellParam(word):
746
      raise errors.ProgrammerError("Shell argument '%s' contains"
747
                                   " invalid characters" % word)
748
  return template % args
749

    
750

    
751
def FormatUnit(value, units):
752
  """Formats an incoming number of MiB with the appropriate unit.
753

754
  @type value: int
755
  @param value: integer representing the value in MiB (1048576)
756
  @type units: char
757
  @param units: the type of formatting we should do:
758
      - 'h' for automatic scaling
759
      - 'm' for MiBs
760
      - 'g' for GiBs
761
      - 't' for TiBs
762
  @rtype: str
763
  @return: the formatted value (with suffix)
764

765
  """
766
  if units not in ('m', 'g', 't', 'h'):
767
    raise errors.ProgrammerError("Invalid unit specified '%s'" % str(units))
768

    
769
  suffix = ''
770

    
771
  if units == 'm' or (units == 'h' and value < 1024):
772
    if units == 'h':
773
      suffix = 'M'
774
    return "%d%s" % (round(value, 0), suffix)
775

    
776
  elif units == 'g' or (units == 'h' and value < (1024 * 1024)):
777
    if units == 'h':
778
      suffix = 'G'
779
    return "%0.1f%s" % (round(float(value) / 1024, 1), suffix)
780

    
781
  else:
782
    if units == 'h':
783
      suffix = 'T'
784
    return "%0.1f%s" % (round(float(value) / 1024 / 1024, 1), suffix)
785

    
786

    
787
def ParseUnit(input_string):
788
  """Tries to extract number and scale from the given string.
789

790
  Input must be in the format C{NUMBER+ [DOT NUMBER+] SPACE*
791
  [UNIT]}. If no unit is specified, it defaults to MiB. Return value
792
  is always an int in MiB.
793

794
  """
795
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', str(input_string))
796
  if not m:
797
    raise errors.UnitParseError("Invalid format")
798

    
799
  value = float(m.groups()[0])
800

    
801
  unit = m.groups()[1]
802
  if unit:
803
    lcunit = unit.lower()
804
  else:
805
    lcunit = 'm'
806

    
807
  if lcunit in ('m', 'mb', 'mib'):
808
    # Value already in MiB
809
    pass
810

    
811
  elif lcunit in ('g', 'gb', 'gib'):
812
    value *= 1024
813

    
814
  elif lcunit in ('t', 'tb', 'tib'):
815
    value *= 1024 * 1024
816

    
817
  else:
818
    raise errors.UnitParseError("Unknown unit: %s" % unit)
819

    
820
  # Make sure we round up
821
  if int(value) < value:
822
    value += 1
823

    
824
  # Round up to the next multiple of 4
825
  value = int(value)
826
  if value % 4:
827
    value += 4 - value % 4
828

    
829
  return value
830

    
831

    
832
def AddAuthorizedKey(file_name, key):
833
  """Adds an SSH public key to an authorized_keys file.
834

835
  @type file_name: str
836
  @param file_name: path to authorized_keys file
837
  @type key: str
838
  @param key: string containing key
839

840
  """
841
  key_fields = key.split()
842

    
843
  f = open(file_name, 'a+')
844
  try:
845
    nl = True
846
    for line in f:
847
      # Ignore whitespace changes
848
      if line.split() == key_fields:
849
        break
850
      nl = line.endswith('\n')
851
    else:
852
      if not nl:
853
        f.write("\n")
854
      f.write(key.rstrip('\r\n'))
855
      f.write("\n")
856
      f.flush()
857
  finally:
858
    f.close()
859

    
860

    
861
def RemoveAuthorizedKey(file_name, key):
862
  """Removes an SSH public key from an authorized_keys file.
863

864
  @type file_name: str
865
  @param file_name: path to authorized_keys file
866
  @type key: str
867
  @param key: string containing key
868

869
  """
870
  key_fields = key.split()
871

    
872
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
873
  try:
874
    out = os.fdopen(fd, 'w')
875
    try:
876
      f = open(file_name, 'r')
877
      try:
878
        for line in f:
879
          # Ignore whitespace changes while comparing lines
880
          if line.split() != key_fields:
881
            out.write(line)
882

    
883
        out.flush()
884
        os.rename(tmpname, file_name)
885
      finally:
886
        f.close()
887
    finally:
888
      out.close()
889
  except:
890
    RemoveFile(tmpname)
891
    raise
892

    
893

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

897
  @type file_name: str
898
  @param file_name: path to the file to modify (usually C{/etc/hosts})
899
  @type ip: str
900
  @param ip: the IP address
901
  @type hostname: str
902
  @param hostname: the hostname to be added
903
  @type aliases: list
904
  @param aliases: the list of aliases to add for the hostname
905

906
  """
907
  # FIXME: use WriteFile + fn rather than duplicating its efforts
908
  # Ensure aliases are unique
909
  aliases = UniqueSequence([hostname] + aliases)[1:]
910

    
911
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
912
  try:
913
    out = os.fdopen(fd, 'w')
914
    try:
915
      f = open(file_name, 'r')
916
      try:
917
        for line in f:
918
          fields = line.split()
919
          if fields and not fields[0].startswith('#') and ip == fields[0]:
920
            continue
921
          out.write(line)
922

    
923
        out.write("%s\t%s" % (ip, hostname))
924
        if aliases:
925
          out.write(" %s" % ' '.join(aliases))
926
        out.write('\n')
927

    
928
        out.flush()
929
        os.fsync(out)
930
        os.chmod(tmpname, 0644)
931
        os.rename(tmpname, file_name)
932
      finally:
933
        f.close()
934
    finally:
935
      out.close()
936
  except:
937
    RemoveFile(tmpname)
938
    raise
939

    
940

    
941
def AddHostToEtcHosts(hostname):
942
  """Wrapper around SetEtcHostsEntry.
943

944
  @type hostname: str
945
  @param hostname: a hostname that will be resolved and added to
946
      L{constants.ETC_HOSTS}
947

948
  """
949
  hi = HostInfo(name=hostname)
950
  SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])
951

    
952

    
953
def RemoveEtcHostsEntry(file_name, hostname):
954
  """Removes a hostname from /etc/hosts.
955

956
  IP addresses without names are removed from the file.
957

958
  @type file_name: str
959
  @param file_name: path to the file to modify (usually C{/etc/hosts})
960
  @type hostname: str
961
  @param hostname: the hostname to be removed
962

963
  """
964
  # FIXME: use WriteFile + fn rather than duplicating its efforts
965
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
966
  try:
967
    out = os.fdopen(fd, 'w')
968
    try:
969
      f = open(file_name, 'r')
970
      try:
971
        for line in f:
972
          fields = line.split()
973
          if len(fields) > 1 and not fields[0].startswith('#'):
974
            names = fields[1:]
975
            if hostname in names:
976
              while hostname in names:
977
                names.remove(hostname)
978
              if names:
979
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
980
              continue
981

    
982
          out.write(line)
983

    
984
        out.flush()
985
        os.fsync(out)
986
        os.chmod(tmpname, 0644)
987
        os.rename(tmpname, file_name)
988
      finally:
989
        f.close()
990
    finally:
991
      out.close()
992
  except:
993
    RemoveFile(tmpname)
994
    raise
995

    
996

    
997
def RemoveHostFromEtcHosts(hostname):
998
  """Wrapper around RemoveEtcHostsEntry.
999

1000
  @type hostname: str
1001
  @param hostname: hostname that will be resolved and its
1002
      full and shot name will be removed from
1003
      L{constants.ETC_HOSTS}
1004

1005
  """
1006
  hi = HostInfo(name=hostname)
1007
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.name)
1008
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.ShortName())
1009

    
1010

    
1011
def CreateBackup(file_name):
1012
  """Creates a backup of a file.
1013

1014
  @type file_name: str
1015
  @param file_name: file to be backed up
1016
  @rtype: str
1017
  @return: the path to the newly created backup
1018
  @raise errors.ProgrammerError: for invalid file names
1019

1020
  """
1021
  if not os.path.isfile(file_name):
1022
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
1023
                                file_name)
1024

    
1025
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
1026
  dir_name = os.path.dirname(file_name)
1027

    
1028
  fsrc = open(file_name, 'rb')
1029
  try:
1030
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
1031
    fdst = os.fdopen(fd, 'wb')
1032
    try:
1033
      shutil.copyfileobj(fsrc, fdst)
1034
    finally:
1035
      fdst.close()
1036
  finally:
1037
    fsrc.close()
1038

    
1039
  return backup_name
1040

    
1041

    
1042
def ShellQuote(value):
1043
  """Quotes shell argument according to POSIX.
1044

1045
  @type value: str
1046
  @param value: the argument to be quoted
1047
  @rtype: str
1048
  @return: the quoted value
1049

1050
  """
1051
  if _re_shell_unquoted.match(value):
1052
    return value
1053
  else:
1054
    return "'%s'" % value.replace("'", "'\\''")
1055

    
1056

    
1057
def ShellQuoteArgs(args):
1058
  """Quotes a list of shell arguments.
1059

1060
  @type args: list
1061
  @param args: list of arguments to be quoted
1062
  @rtype: str
1063
  @return: the quoted arguments concatenated with spaces
1064

1065
  """
1066
  return ' '.join([ShellQuote(i) for i in args])
1067

    
1068

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

1072
  Check if the given IP is reachable by doing attempting a TCP connect
1073
  to it.
1074

1075
  @type target: str
1076
  @param target: the IP or hostname to ping
1077
  @type port: int
1078
  @param port: the port to connect to
1079
  @type timeout: int
1080
  @param timeout: the timeout on the connection attempt
1081
  @type live_port_needed: boolean
1082
  @param live_port_needed: whether a closed port will cause the
1083
      function to return failure, as if there was a timeout
1084
  @type source: str or None
1085
  @param source: if specified, will cause the connect to be made
1086
      from this specific source address; failures to bind other
1087
      than C{EADDRNOTAVAIL} will be ignored
1088

1089
  """
1090
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
1091

    
1092
  success = False
1093

    
1094
  if source is not None:
1095
    try:
1096
      sock.bind((source, 0))
1097
    except socket.error, (errcode, _):
1098
      if errcode == errno.EADDRNOTAVAIL:
1099
        success = False
1100

    
1101
  sock.settimeout(timeout)
1102

    
1103
  try:
1104
    sock.connect((target, port))
1105
    sock.close()
1106
    success = True
1107
  except socket.timeout:
1108
    success = False
1109
  except socket.error, (errcode, _):
1110
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
1111

    
1112
  return success
1113

    
1114

    
1115
def OwnIpAddress(address):
1116
  """Check if the current host has the the given IP address.
1117

1118
  Currently this is done by TCP-pinging the address from the loopback
1119
  address.
1120

1121
  @type address: string
1122
  @param address: the address to check
1123
  @rtype: bool
1124
  @return: True if we own the address
1125

1126
  """
1127
  return TcpPing(address, constants.DEFAULT_NODED_PORT,
1128
                 source=constants.LOCALHOST_IP_ADDRESS)
1129

    
1130

    
1131
def ListVisibleFiles(path):
1132
  """Returns a list of visible files in a directory.
1133

1134
  @type path: str
1135
  @param path: the directory to enumerate
1136
  @rtype: list
1137
  @return: the list of all files not starting with a dot
1138

1139
  """
1140
  files = [i for i in os.listdir(path) if not i.startswith(".")]
1141
  files.sort()
1142
  return files
1143

    
1144

    
1145
def GetHomeDir(user, default=None):
1146
  """Try to get the homedir of the given user.
1147

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

1152
  """
1153
  try:
1154
    if isinstance(user, basestring):
1155
      result = pwd.getpwnam(user)
1156
    elif isinstance(user, (int, long)):
1157
      result = pwd.getpwuid(user)
1158
    else:
1159
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
1160
                                   type(user))
1161
  except KeyError:
1162
    return default
1163
  return result.pw_dir
1164

    
1165

    
1166
def NewUUID():
1167
  """Returns a random UUID.
1168

1169
  @note: This is a Linux-specific method as it uses the /proc
1170
      filesystem.
1171
  @rtype: str
1172

1173
  """
1174
  return ReadFile(_RANDOM_UUID_FILE, size=128).rstrip("\n")
1175

    
1176

    
1177
def GenerateSecret(numbytes=20):
1178
  """Generates a random secret.
1179

1180
  This will generate a pseudo-random secret returning an hex string
1181
  (so that it can be used where an ASCII string is needed).
1182

1183
  @param numbytes: the number of bytes which will be represented by the returned
1184
      string (defaulting to 20, the length of a SHA1 hash)
1185
  @rtype: str
1186
  @return: an hex representation of the pseudo-random sequence
1187

1188
  """
1189
  return os.urandom(numbytes).encode('hex')
1190

    
1191

    
1192
def EnsureDirs(dirs):
1193
  """Make required directories, if they don't exist.
1194

1195
  @param dirs: list of tuples (dir_name, dir_mode)
1196
  @type dirs: list of (string, integer)
1197

1198
  """
1199
  for dir_name, dir_mode in dirs:
1200
    try:
1201
      os.mkdir(dir_name, dir_mode)
1202
    except EnvironmentError, err:
1203
      if err.errno != errno.EEXIST:
1204
        raise errors.GenericError("Cannot create needed directory"
1205
                                  " '%s': %s" % (dir_name, err))
1206
    if not os.path.isdir(dir_name):
1207
      raise errors.GenericError("%s is not a directory" % dir_name)
1208

    
1209

    
1210
def ReadFile(file_name, size=None):
1211
  """Reads a file.
1212

1213
  @type size: None or int
1214
  @param size: Read at most size bytes
1215
  @rtype: str
1216
  @return: the (possibly partial) content of the file
1217

1218
  """
1219
  f = open(file_name, "r")
1220
  try:
1221
    if size is None:
1222
      return f.read()
1223
    else:
1224
      return f.read(size)
1225
  finally:
1226
    f.close()
1227

    
1228

    
1229
def WriteFile(file_name, fn=None, data=None,
1230
              mode=None, uid=-1, gid=-1,
1231
              atime=None, mtime=None, close=True,
1232
              dry_run=False, backup=False,
1233
              prewrite=None, postwrite=None):
1234
  """(Over)write a file atomically.
1235

1236
  The file_name and either fn (a function taking one argument, the
1237
  file descriptor, and which should write the data to it) or data (the
1238
  contents of the file) must be passed. The other arguments are
1239
  optional and allow setting the file mode, owner and group, and the
1240
  mtime/atime of the file.
1241

1242
  If the function doesn't raise an exception, it has succeeded and the
1243
  target file has the new contents. If the function has raised an
1244
  exception, an existing target file should be unmodified and the
1245
  temporary file should be removed.
1246

1247
  @type file_name: str
1248
  @param file_name: the target filename
1249
  @type fn: callable
1250
  @param fn: content writing function, called with
1251
      file descriptor as parameter
1252
  @type data: str
1253
  @param data: contents of the file
1254
  @type mode: int
1255
  @param mode: file mode
1256
  @type uid: int
1257
  @param uid: the owner of the file
1258
  @type gid: int
1259
  @param gid: the group of the file
1260
  @type atime: int
1261
  @param atime: a custom access time to be set on the file
1262
  @type mtime: int
1263
  @param mtime: a custom modification time to be set on the file
1264
  @type close: boolean
1265
  @param close: whether to close file after writing it
1266
  @type prewrite: callable
1267
  @param prewrite: function to be called before writing content
1268
  @type postwrite: callable
1269
  @param postwrite: function to be called after writing content
1270

1271
  @rtype: None or int
1272
  @return: None if the 'close' parameter evaluates to True,
1273
      otherwise the file descriptor
1274

1275
  @raise errors.ProgrammerError: if any of the arguments are not valid
1276

1277
  """
1278
  if not os.path.isabs(file_name):
1279
    raise errors.ProgrammerError("Path passed to WriteFile is not"
1280
                                 " absolute: '%s'" % file_name)
1281

    
1282
  if [fn, data].count(None) != 1:
1283
    raise errors.ProgrammerError("fn or data required")
1284

    
1285
  if [atime, mtime].count(None) == 1:
1286
    raise errors.ProgrammerError("Both atime and mtime must be either"
1287
                                 " set or None")
1288

    
1289
  if backup and not dry_run and os.path.isfile(file_name):
1290
    CreateBackup(file_name)
1291

    
1292
  dir_name, base_name = os.path.split(file_name)
1293
  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
1294
  do_remove = True
1295
  # here we need to make sure we remove the temp file, if any error
1296
  # leaves it in place
1297
  try:
1298
    if uid != -1 or gid != -1:
1299
      os.chown(new_name, uid, gid)
1300
    if mode:
1301
      os.chmod(new_name, mode)
1302
    if callable(prewrite):
1303
      prewrite(fd)
1304
    if data is not None:
1305
      os.write(fd, data)
1306
    else:
1307
      fn(fd)
1308
    if callable(postwrite):
1309
      postwrite(fd)
1310
    os.fsync(fd)
1311
    if atime is not None and mtime is not None:
1312
      os.utime(new_name, (atime, mtime))
1313
    if not dry_run:
1314
      os.rename(new_name, file_name)
1315
      do_remove = False
1316
  finally:
1317
    if close:
1318
      os.close(fd)
1319
      result = None
1320
    else:
1321
      result = fd
1322
    if do_remove:
1323
      RemoveFile(new_name)
1324

    
1325
  return result
1326

    
1327

    
1328
def FirstFree(seq, base=0):
1329
  """Returns the first non-existing integer from seq.
1330

1331
  The seq argument should be a sorted list of positive integers. The
1332
  first time the index of an element is smaller than the element
1333
  value, the index will be returned.
1334

1335
  The base argument is used to start at a different offset,
1336
  i.e. C{[3, 4, 6]} with I{offset=3} will return 5.
1337

1338
  Example: C{[0, 1, 3]} will return I{2}.
1339

1340
  @type seq: sequence
1341
  @param seq: the sequence to be analyzed.
1342
  @type base: int
1343
  @param base: use this value as the base index of the sequence
1344
  @rtype: int
1345
  @return: the first non-used index in the sequence
1346

1347
  """
1348
  for idx, elem in enumerate(seq):
1349
    assert elem >= base, "Passed element is higher than base offset"
1350
    if elem > idx + base:
1351
      # idx is not used
1352
      return idx + base
1353
  return None
1354

    
1355

    
1356
def all(seq, pred=bool):
1357
  "Returns True if pred(x) is True for every element in the iterable"
1358
  for _ in itertools.ifilterfalse(pred, seq):
1359
    return False
1360
  return True
1361

    
1362

    
1363
def any(seq, pred=bool):
1364
  "Returns True if pred(x) is True for at least one element in the iterable"
1365
  for _ in itertools.ifilter(pred, seq):
1366
    return True
1367
  return False
1368

    
1369

    
1370
def UniqueSequence(seq):
1371
  """Returns a list with unique elements.
1372

1373
  Element order is preserved.
1374

1375
  @type seq: sequence
1376
  @param seq: the sequence with the source elements
1377
  @rtype: list
1378
  @return: list of unique elements from seq
1379

1380
  """
1381
  seen = set()
1382
  return [i for i in seq if i not in seen and not seen.add(i)]
1383

    
1384

    
1385
def IsValidMac(mac):
1386
  """Predicate to check if a MAC address is valid.
1387

1388
  Checks whether the supplied MAC address is formally correct, only
1389
  accepts colon separated format.
1390

1391
  @type mac: str
1392
  @param mac: the MAC to be validated
1393
  @rtype: boolean
1394
  @return: True is the MAC seems valid
1395

1396
  """
1397
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
1398
  return mac_check.match(mac) is not None
1399

    
1400

    
1401
def TestDelay(duration):
1402
  """Sleep for a fixed amount of time.
1403

1404
  @type duration: float
1405
  @param duration: the sleep duration
1406
  @rtype: boolean
1407
  @return: False for negative value, True otherwise
1408

1409
  """
1410
  if duration < 0:
1411
    return False, "Invalid sleep duration"
1412
  time.sleep(duration)
1413
  return True, None
1414

    
1415

    
1416
def _CloseFDNoErr(fd, retries=5):
1417
  """Close a file descriptor ignoring errors.
1418

1419
  @type fd: int
1420
  @param fd: the file descriptor
1421
  @type retries: int
1422
  @param retries: how many retries to make, in case we get any
1423
      other error than EBADF
1424

1425
  """
1426
  try:
1427
    os.close(fd)
1428
  except OSError, err:
1429
    if err.errno != errno.EBADF:
1430
      if retries > 0:
1431
        _CloseFDNoErr(fd, retries - 1)
1432
    # else either it's closed already or we're out of retries, so we
1433
    # ignore this and go on
1434

    
1435

    
1436
def CloseFDs(noclose_fds=None):
1437
  """Close file descriptors.
1438

1439
  This closes all file descriptors above 2 (i.e. except
1440
  stdin/out/err).
1441

1442
  @type noclose_fds: list or None
1443
  @param noclose_fds: if given, it denotes a list of file descriptor
1444
      that should not be closed
1445

1446
  """
1447
  # Default maximum for the number of available file descriptors.
1448
  if 'SC_OPEN_MAX' in os.sysconf_names:
1449
    try:
1450
      MAXFD = os.sysconf('SC_OPEN_MAX')
1451
      if MAXFD < 0:
1452
        MAXFD = 1024
1453
    except OSError:
1454
      MAXFD = 1024
1455
  else:
1456
    MAXFD = 1024
1457
  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1458
  if (maxfd == resource.RLIM_INFINITY):
1459
    maxfd = MAXFD
1460

    
1461
  # Iterate through and close all file descriptors (except the standard ones)
1462
  for fd in range(3, maxfd):
1463
    if noclose_fds and fd in noclose_fds:
1464
      continue
1465
    _CloseFDNoErr(fd)
1466

    
1467

    
1468
def Daemonize(logfile):
1469
  """Daemonize the current process.
1470

1471
  This detaches the current process from the controlling terminal and
1472
  runs it in the background as a daemon.
1473

1474
  @type logfile: str
1475
  @param logfile: the logfile to which we should redirect stdout/stderr
1476
  @rtype: int
1477
  @return: the value zero
1478

1479
  """
1480
  UMASK = 077
1481
  WORKDIR = "/"
1482

    
1483
  # this might fail
1484
  pid = os.fork()
1485
  if (pid == 0):  # The first child.
1486
    os.setsid()
1487
    # this might fail
1488
    pid = os.fork() # Fork a second child.
1489
    if (pid == 0):  # The second child.
1490
      os.chdir(WORKDIR)
1491
      os.umask(UMASK)
1492
    else:
1493
      # exit() or _exit()?  See below.
1494
      os._exit(0) # Exit parent (the first child) of the second child.
1495
  else:
1496
    os._exit(0) # Exit parent of the first child.
1497

    
1498
  for fd in range(3):
1499
    _CloseFDNoErr(fd)
1500
  i = os.open("/dev/null", os.O_RDONLY) # stdin
1501
  assert i == 0, "Can't close/reopen stdin"
1502
  i = os.open(logfile, os.O_WRONLY|os.O_CREAT|os.O_APPEND, 0600) # stdout
1503
  assert i == 1, "Can't close/reopen stdout"
1504
  # Duplicate standard output to standard error.
1505
  os.dup2(1, 2)
1506
  return 0
1507

    
1508

    
1509
def DaemonPidFileName(name):
1510
  """Compute a ganeti pid file absolute path
1511

1512
  @type name: str
1513
  @param name: the daemon name
1514
  @rtype: str
1515
  @return: the full path to the pidfile corresponding to the given
1516
      daemon name
1517

1518
  """
1519
  return os.path.join(constants.RUN_GANETI_DIR, "%s.pid" % name)
1520

    
1521

    
1522
def WritePidFile(name):
1523
  """Write the current process pidfile.
1524

1525
  The file will be written to L{constants.RUN_GANETI_DIR}I{/name.pid}
1526

1527
  @type name: str
1528
  @param name: the daemon name to use
1529
  @raise errors.GenericError: if the pid file already exists and
1530
      points to a live process
1531

1532
  """
1533
  pid = os.getpid()
1534
  pidfilename = DaemonPidFileName(name)
1535
  if IsProcessAlive(ReadPidFile(pidfilename)):
1536
    raise errors.GenericError("%s contains a live process" % pidfilename)
1537

    
1538
  WriteFile(pidfilename, data="%d\n" % pid)
1539

    
1540

    
1541
def RemovePidFile(name):
1542
  """Remove the current process pidfile.
1543

1544
  Any errors are ignored.
1545

1546
  @type name: str
1547
  @param name: the daemon name used to derive the pidfile name
1548

1549
  """
1550
  pidfilename = DaemonPidFileName(name)
1551
  # TODO: we could check here that the file contains our pid
1552
  try:
1553
    RemoveFile(pidfilename)
1554
  except:
1555
    pass
1556

    
1557

    
1558
def KillProcess(pid, signal_=signal.SIGTERM, timeout=30,
1559
                waitpid=False):
1560
  """Kill a process given by its pid.
1561

1562
  @type pid: int
1563
  @param pid: The PID to terminate.
1564
  @type signal_: int
1565
  @param signal_: The signal to send, by default SIGTERM
1566
  @type timeout: int
1567
  @param timeout: The timeout after which, if the process is still alive,
1568
                  a SIGKILL will be sent. If not positive, no such checking
1569
                  will be done
1570
  @type waitpid: boolean
1571
  @param waitpid: If true, we should waitpid on this process after
1572
      sending signals, since it's our own child and otherwise it
1573
      would remain as zombie
1574

1575
  """
1576
  def _helper(pid, signal_, wait):
1577
    """Simple helper to encapsulate the kill/waitpid sequence"""
1578
    os.kill(pid, signal_)
1579
    if wait:
1580
      try:
1581
        os.waitpid(pid, os.WNOHANG)
1582
      except OSError:
1583
        pass
1584

    
1585
  if pid <= 0:
1586
    # kill with pid=0 == suicide
1587
    raise errors.ProgrammerError("Invalid pid given '%s'" % pid)
1588

    
1589
  if not IsProcessAlive(pid):
1590
    return
1591

    
1592
  _helper(pid, signal_, waitpid)
1593

    
1594
  if timeout <= 0:
1595
    return
1596

    
1597
  def _CheckProcess():
1598
    if not IsProcessAlive(pid):
1599
      return
1600

    
1601
    try:
1602
      (result_pid, _) = os.waitpid(pid, os.WNOHANG)
1603
    except OSError:
1604
      raise RetryAgain()
1605

    
1606
    if result_pid > 0:
1607
      return
1608

    
1609
    raise RetryAgain()
1610

    
1611
  try:
1612
    # Wait up to $timeout seconds
1613
    Retry(_CheckProcess, (0.01, 1.5, 0.1), timeout)
1614
  except RetryTimeout:
1615
    pass
1616

    
1617
  if IsProcessAlive(pid):
1618
    # Kill process if it's still alive
1619
    _helper(pid, signal.SIGKILL, waitpid)
1620

    
1621

    
1622
def FindFile(name, search_path, test=os.path.exists):
1623
  """Look for a filesystem object in a given path.
1624

1625
  This is an abstract method to search for filesystem object (files,
1626
  dirs) under a given search path.
1627

1628
  @type name: str
1629
  @param name: the name to look for
1630
  @type search_path: str
1631
  @param search_path: location to start at
1632
  @type test: callable
1633
  @param test: a function taking one argument that should return True
1634
      if the a given object is valid; the default value is
1635
      os.path.exists, causing only existing files to be returned
1636
  @rtype: str or None
1637
  @return: full path to the object if found, None otherwise
1638

1639
  """
1640
  for dir_name in search_path:
1641
    item_name = os.path.sep.join([dir_name, name])
1642
    if test(item_name):
1643
      return item_name
1644
  return None
1645

    
1646

    
1647
def CheckVolumeGroupSize(vglist, vgname, minsize):
1648
  """Checks if the volume group list is valid.
1649

1650
  The function will check if a given volume group is in the list of
1651
  volume groups and has a minimum size.
1652

1653
  @type vglist: dict
1654
  @param vglist: dictionary of volume group names and their size
1655
  @type vgname: str
1656
  @param vgname: the volume group we should check
1657
  @type minsize: int
1658
  @param minsize: the minimum size we accept
1659
  @rtype: None or str
1660
  @return: None for success, otherwise the error message
1661

1662
  """
1663
  vgsize = vglist.get(vgname, None)
1664
  if vgsize is None:
1665
    return "volume group '%s' missing" % vgname
1666
  elif vgsize < minsize:
1667
    return ("volume group '%s' too small (%s MiB required, %d MiB found)" %
1668
            (vgname, minsize, vgsize))
1669
  return None
1670

    
1671

    
1672
def SplitTime(value):
1673
  """Splits time as floating point number into a tuple.
1674

1675
  @param value: Time in seconds
1676
  @type value: int or float
1677
  @return: Tuple containing (seconds, microseconds)
1678

1679
  """
1680
  (seconds, microseconds) = divmod(int(value * 1000000), 1000000)
1681

    
1682
  assert 0 <= seconds, \
1683
    "Seconds must be larger than or equal to 0, but are %s" % seconds
1684
  assert 0 <= microseconds <= 999999, \
1685
    "Microseconds must be 0-999999, but are %s" % microseconds
1686

    
1687
  return (int(seconds), int(microseconds))
1688

    
1689

    
1690
def MergeTime(timetuple):
1691
  """Merges a tuple into time as a floating point number.
1692

1693
  @param timetuple: Time as tuple, (seconds, microseconds)
1694
  @type timetuple: tuple
1695
  @return: Time as a floating point number expressed in seconds
1696

1697
  """
1698
  (seconds, microseconds) = timetuple
1699

    
1700
  assert 0 <= seconds, \
1701
    "Seconds must be larger than or equal to 0, but are %s" % seconds
1702
  assert 0 <= microseconds <= 999999, \
1703
    "Microseconds must be 0-999999, but are %s" % microseconds
1704

    
1705
  return float(seconds) + (float(microseconds) * 0.000001)
1706

    
1707

    
1708
def GetDaemonPort(daemon_name):
1709
  """Get the daemon port for this cluster.
1710

1711
  Note that this routine does not read a ganeti-specific file, but
1712
  instead uses C{socket.getservbyname} to allow pre-customization of
1713
  this parameter outside of Ganeti.
1714

1715
  @type daemon_name: string
1716
  @param daemon_name: daemon name (in constants.DAEMONS_PORTS)
1717
  @rtype: int
1718

1719
  """
1720
  if daemon_name not in constants.DAEMONS_PORTS:
1721
    raise errors.ProgrammerError("Unknown daemon: %s" % daemon_name)
1722

    
1723
  (proto, default_port) = constants.DAEMONS_PORTS[daemon_name]
1724
  try:
1725
    port = socket.getservbyname(daemon_name, proto)
1726
  except socket.error:
1727
    port = default_port
1728

    
1729
  return port
1730

    
1731

    
1732
def SetupLogging(logfile, debug=False, stderr_logging=False, program="",
1733
                 multithreaded=False):
1734
  """Configures the logging module.
1735

1736
  @type logfile: str
1737
  @param logfile: the filename to which we should log
1738
  @type debug: boolean
1739
  @param debug: whether to enable debug messages too or
1740
      only those at C{INFO} and above level
1741
  @type stderr_logging: boolean
1742
  @param stderr_logging: whether we should also log to the standard error
1743
  @type program: str
1744
  @param program: the name under which we should log messages
1745
  @type multithreaded: boolean
1746
  @param multithreaded: if True, will add the thread name to the log file
1747
  @raise EnvironmentError: if we can't open the log file and
1748
      stderr logging is disabled
1749

1750
  """
1751
  fmt = "%(asctime)s: " + program + " pid=%(process)d"
1752
  if multithreaded:
1753
    fmt += "/%(threadName)s"
1754
  if debug:
1755
    fmt += " %(module)s:%(lineno)s"
1756
  fmt += " %(levelname)s %(message)s"
1757
  formatter = logging.Formatter(fmt)
1758

    
1759
  root_logger = logging.getLogger("")
1760
  root_logger.setLevel(logging.NOTSET)
1761

    
1762
  # Remove all previously setup handlers
1763
  for handler in root_logger.handlers:
1764
    handler.close()
1765
    root_logger.removeHandler(handler)
1766

    
1767
  if stderr_logging:
1768
    stderr_handler = logging.StreamHandler()
1769
    stderr_handler.setFormatter(formatter)
1770
    if debug:
1771
      stderr_handler.setLevel(logging.NOTSET)
1772
    else:
1773
      stderr_handler.setLevel(logging.CRITICAL)
1774
    root_logger.addHandler(stderr_handler)
1775

    
1776
  # this can fail, if the logging directories are not setup or we have
1777
  # a permisssion problem; in this case, it's best to log but ignore
1778
  # the error if stderr_logging is True, and if false we re-raise the
1779
  # exception since otherwise we could run but without any logs at all
1780
  try:
1781
    logfile_handler = logging.FileHandler(logfile)
1782
    logfile_handler.setFormatter(formatter)
1783
    if debug:
1784
      logfile_handler.setLevel(logging.DEBUG)
1785
    else:
1786
      logfile_handler.setLevel(logging.INFO)
1787
    root_logger.addHandler(logfile_handler)
1788
  except EnvironmentError:
1789
    if stderr_logging:
1790
      logging.exception("Failed to enable logging to file '%s'", logfile)
1791
    else:
1792
      # we need to re-raise the exception
1793
      raise
1794

    
1795

    
1796
def IsNormAbsPath(path):
1797
  """Check whether a path is absolute and also normalized
1798

1799
  This avoids things like /dir/../../other/path to be valid.
1800

1801
  """
1802
  return os.path.normpath(path) == path and os.path.isabs(path)
1803

    
1804

    
1805
def TailFile(fname, lines=20):
1806
  """Return the last lines from a file.
1807

1808
  @note: this function will only read and parse the last 4KB of
1809
      the file; if the lines are very long, it could be that less
1810
      than the requested number of lines are returned
1811

1812
  @param fname: the file name
1813
  @type lines: int
1814
  @param lines: the (maximum) number of lines to return
1815

1816
  """
1817
  fd = open(fname, "r")
1818
  try:
1819
    fd.seek(0, 2)
1820
    pos = fd.tell()
1821
    pos = max(0, pos-4096)
1822
    fd.seek(pos, 0)
1823
    raw_data = fd.read()
1824
  finally:
1825
    fd.close()
1826

    
1827
  rows = raw_data.splitlines()
1828
  return rows[-lines:]
1829

    
1830

    
1831
def SafeEncode(text):
1832
  """Return a 'safe' version of a source string.
1833

1834
  This function mangles the input string and returns a version that
1835
  should be safe to display/encode as ASCII. To this end, we first
1836
  convert it to ASCII using the 'backslashreplace' encoding which
1837
  should get rid of any non-ASCII chars, and then we process it
1838
  through a loop copied from the string repr sources in the python; we
1839
  don't use string_escape anymore since that escape single quotes and
1840
  backslashes too, and that is too much; and that escaping is not
1841
  stable, i.e. string_escape(string_escape(x)) != string_escape(x).
1842

1843
  @type text: str or unicode
1844
  @param text: input data
1845
  @rtype: str
1846
  @return: a safe version of text
1847

1848
  """
1849
  if isinstance(text, unicode):
1850
    # only if unicode; if str already, we handle it below
1851
    text = text.encode('ascii', 'backslashreplace')
1852
  resu = ""
1853
  for char in text:
1854
    c = ord(char)
1855
    if char  == '\t':
1856
      resu += r'\t'
1857
    elif char == '\n':
1858
      resu += r'\n'
1859
    elif char == '\r':
1860
      resu += r'\'r'
1861
    elif c < 32 or c >= 127: # non-printable
1862
      resu += "\\x%02x" % (c & 0xff)
1863
    else:
1864
      resu += char
1865
  return resu
1866

    
1867

    
1868
def BytesToMebibyte(value):
1869
  """Converts bytes to mebibytes.
1870

1871
  @type value: int
1872
  @param value: Value in bytes
1873
  @rtype: int
1874
  @return: Value in mebibytes
1875

1876
  """
1877
  return int(round(value / (1024.0 * 1024.0), 0))
1878

    
1879

    
1880
def CalculateDirectorySize(path):
1881
  """Calculates the size of a directory recursively.
1882

1883
  @type path: string
1884
  @param path: Path to directory
1885
  @rtype: int
1886
  @return: Size in mebibytes
1887

1888
  """
1889
  size = 0
1890

    
1891
  for (curpath, _, files) in os.walk(path):
1892
    for filename in files:
1893
      st = os.lstat(os.path.join(curpath, filename))
1894
      size += st.st_size
1895

    
1896
  return BytesToMebibyte(size)
1897

    
1898

    
1899
def GetFilesystemStats(path):
1900
  """Returns the total and free space on a filesystem.
1901

1902
  @type path: string
1903
  @param path: Path on filesystem to be examined
1904
  @rtype: int
1905
  @return: tuple of (Total space, Free space) in mebibytes
1906

1907
  """
1908
  st = os.statvfs(path)
1909

    
1910
  fsize = BytesToMebibyte(st.f_bavail * st.f_frsize)
1911
  tsize = BytesToMebibyte(st.f_blocks * st.f_frsize)
1912
  return (tsize, fsize)
1913

    
1914

    
1915
def LockedMethod(fn):
1916
  """Synchronized object access decorator.
1917

1918
  This decorator is intended to protect access to an object using the
1919
  object's own lock which is hardcoded to '_lock'.
1920

1921
  """
1922
  def _LockDebug(*args, **kwargs):
1923
    if debug_locks:
1924
      logging.debug(*args, **kwargs)
1925

    
1926
  def wrapper(self, *args, **kwargs):
1927
    assert hasattr(self, '_lock')
1928
    lock = self._lock
1929
    _LockDebug("Waiting for %s", lock)
1930
    lock.acquire()
1931
    try:
1932
      _LockDebug("Acquired %s", lock)
1933
      result = fn(self, *args, **kwargs)
1934
    finally:
1935
      _LockDebug("Releasing %s", lock)
1936
      lock.release()
1937
      _LockDebug("Released %s", lock)
1938
    return result
1939
  return wrapper
1940

    
1941

    
1942
def LockFile(fd):
1943
  """Locks a file using POSIX locks.
1944

1945
  @type fd: int
1946
  @param fd: the file descriptor we need to lock
1947

1948
  """
1949
  try:
1950
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
1951
  except IOError, err:
1952
    if err.errno == errno.EAGAIN:
1953
      raise errors.LockError("File already locked")
1954
    raise
1955

    
1956

    
1957
def FormatTime(val):
1958
  """Formats a time value.
1959

1960
  @type val: float or None
1961
  @param val: the timestamp as returned by time.time()
1962
  @return: a string value or N/A if we don't have a valid timestamp
1963

1964
  """
1965
  if val is None or not isinstance(val, (int, float)):
1966
    return "N/A"
1967
  # these two codes works on Linux, but they are not guaranteed on all
1968
  # platforms
1969
  return time.strftime("%F %T", time.localtime(val))
1970

    
1971

    
1972
def ReadWatcherPauseFile(filename, now=None, remove_after=3600):
1973
  """Reads the watcher pause file.
1974

1975
  @type filename: string
1976
  @param filename: Path to watcher pause file
1977
  @type now: None, float or int
1978
  @param now: Current time as Unix timestamp
1979
  @type remove_after: int
1980
  @param remove_after: Remove watcher pause file after specified amount of
1981
    seconds past the pause end time
1982

1983
  """
1984
  if now is None:
1985
    now = time.time()
1986

    
1987
  try:
1988
    value = ReadFile(filename)
1989
  except IOError, err:
1990
    if err.errno != errno.ENOENT:
1991
      raise
1992
    value = None
1993

    
1994
  if value is not None:
1995
    try:
1996
      value = int(value)
1997
    except ValueError:
1998
      logging.warning(("Watcher pause file (%s) contains invalid value,"
1999
                       " removing it"), filename)
2000
      RemoveFile(filename)
2001
      value = None
2002

    
2003
    if value is not None:
2004
      # Remove file if it's outdated
2005
      if now > (value + remove_after):
2006
        RemoveFile(filename)
2007
        value = None
2008

    
2009
      elif now > value:
2010
        value = None
2011

    
2012
  return value
2013

    
2014

    
2015
class RetryTimeout(Exception):
2016
  """Retry loop timed out.
2017

2018
  """
2019

    
2020

    
2021
class RetryAgain(Exception):
2022
  """Retry again.
2023

2024
  """
2025

    
2026

    
2027
class _RetryDelayCalculator(object):
2028
  """Calculator for increasing delays.
2029

2030
  """
2031
  __slots__ = [
2032
    "_factor",
2033
    "_limit",
2034
    "_next",
2035
    "_start",
2036
    ]
2037

    
2038
  def __init__(self, start, factor, limit):
2039
    """Initializes this class.
2040

2041
    @type start: float
2042
    @param start: Initial delay
2043
    @type factor: float
2044
    @param factor: Factor for delay increase
2045
    @type limit: float or None
2046
    @param limit: Upper limit for delay or None for no limit
2047

2048
    """
2049
    assert start > 0.0
2050
    assert factor >= 1.0
2051
    assert limit is None or limit >= 0.0
2052

    
2053
    self._start = start
2054
    self._factor = factor
2055
    self._limit = limit
2056

    
2057
    self._next = start
2058

    
2059
  def __call__(self):
2060
    """Returns current delay and calculates the next one.
2061

2062
    """
2063
    current = self._next
2064

    
2065
    # Update for next run
2066
    if self._limit is None or self._next < self._limit:
2067
      self._next = max(self._limit, self._next * self._factor)
2068

    
2069
    return current
2070

    
2071

    
2072
#: Special delay to specify whole remaining timeout
2073
RETRY_REMAINING_TIME = object()
2074

    
2075

    
2076
def Retry(fn, delay, timeout, args=None, wait_fn=time.sleep,
2077
          _time_fn=time.time):
2078
  """Call a function repeatedly until it succeeds.
2079

2080
  The function C{fn} is called repeatedly until it doesn't throw L{RetryAgain}
2081
  anymore. Between calls a delay, specified by C{delay}, is inserted. After a
2082
  total of C{timeout} seconds, this function throws L{RetryTimeout}.
2083

2084
  C{delay} can be one of the following:
2085
    - callable returning the delay length as a float
2086
    - Tuple of (start, factor, limit)
2087
    - L{RETRY_REMAINING_TIME} to sleep until the timeout expires (this is
2088
      useful when overriding L{wait_fn} to wait for an external event)
2089
    - A static delay as a number (int or float)
2090

2091
  @type fn: callable
2092
  @param fn: Function to be called
2093
  @param delay: Either a callable (returning the delay), a tuple of (start,
2094
                factor, limit) (see L{_RetryDelayCalculator}),
2095
                L{RETRY_REMAINING_TIME} or a number (int or float)
2096
  @type timeout: float
2097
  @param timeout: Total timeout
2098
  @type wait_fn: callable
2099
  @param wait_fn: Waiting function
2100
  @return: Return value of function
2101

2102
  """
2103
  assert callable(fn)
2104
  assert callable(wait_fn)
2105
  assert callable(_time_fn)
2106

    
2107
  if args is None:
2108
    args = []
2109

    
2110
  end_time = _time_fn() + timeout
2111

    
2112
  if callable(delay):
2113
    # External function to calculate delay
2114
    calc_delay = delay
2115

    
2116
  elif isinstance(delay, (tuple, list)):
2117
    # Increasing delay with optional upper boundary
2118
    (start, factor, limit) = delay
2119
    calc_delay = _RetryDelayCalculator(start, factor, limit)
2120

    
2121
  elif delay is RETRY_REMAINING_TIME:
2122
    # Always use the remaining time
2123
    calc_delay = None
2124

    
2125
  else:
2126
    # Static delay
2127
    calc_delay = lambda: delay
2128

    
2129
  assert calc_delay is None or callable(calc_delay)
2130

    
2131
  while True:
2132
    try:
2133
      return fn(*args)
2134
    except RetryAgain:
2135
      pass
2136

    
2137
    remaining_time = end_time - _time_fn()
2138

    
2139
    if remaining_time < 0.0:
2140
      raise RetryTimeout()
2141

    
2142
    assert remaining_time >= 0.0
2143

    
2144
    if calc_delay is None:
2145
      wait_fn(remaining_time)
2146
    else:
2147
      current_delay = calc_delay()
2148
      if current_delay > 0.0:
2149
        wait_fn(current_delay)
2150

    
2151

    
2152
class FileLock(object):
2153
  """Utility class for file locks.
2154

2155
  """
2156
  def __init__(self, filename):
2157
    """Constructor for FileLock.
2158

2159
    This will open the file denoted by the I{filename} argument.
2160

2161
    @type filename: str
2162
    @param filename: path to the file to be locked
2163

2164
    """
2165
    self.filename = filename
2166
    self.fd = open(self.filename, "w")
2167

    
2168
  def __del__(self):
2169
    self.Close()
2170

    
2171
  def Close(self):
2172
    """Close the file and release the lock.
2173

2174
    """
2175
    if self.fd:
2176
      self.fd.close()
2177
      self.fd = None
2178

    
2179
  def _flock(self, flag, blocking, timeout, errmsg):
2180
    """Wrapper for fcntl.flock.
2181

2182
    @type flag: int
2183
    @param flag: operation flag
2184
    @type blocking: bool
2185
    @param blocking: whether the operation should be done in blocking mode.
2186
    @type timeout: None or float
2187
    @param timeout: for how long the operation should be retried (implies
2188
                    non-blocking mode).
2189
    @type errmsg: string
2190
    @param errmsg: error message in case operation fails.
2191

2192
    """
2193
    assert self.fd, "Lock was closed"
2194
    assert timeout is None or timeout >= 0, \
2195
      "If specified, timeout must be positive"
2196

    
2197
    if timeout is not None:
2198
      flag |= fcntl.LOCK_NB
2199
      timeout_end = time.time() + timeout
2200

    
2201
    # Blocking doesn't have effect with timeout
2202
    elif not blocking:
2203
      flag |= fcntl.LOCK_NB
2204
      timeout_end = None
2205

    
2206
    # TODO: Convert to utils.Retry
2207

    
2208
    retry = True
2209
    while retry:
2210
      try:
2211
        fcntl.flock(self.fd, flag)
2212
        retry = False
2213
      except IOError, err:
2214
        if err.errno in (errno.EAGAIN, ):
2215
          if timeout_end is not None and time.time() < timeout_end:
2216
            # Wait before trying again
2217
            time.sleep(max(0.1, min(1.0, timeout)))
2218
          else:
2219
            raise errors.LockError(errmsg)
2220
        else:
2221
          logging.exception("fcntl.flock failed")
2222
          raise
2223

    
2224
  def Exclusive(self, blocking=False, timeout=None):
2225
    """Locks the file in exclusive mode.
2226

2227
    @type blocking: boolean
2228
    @param blocking: whether to block and wait until we
2229
        can lock the file or return immediately
2230
    @type timeout: int or None
2231
    @param timeout: if not None, the duration to wait for the lock
2232
        (in blocking mode)
2233

2234
    """
2235
    self._flock(fcntl.LOCK_EX, blocking, timeout,
2236
                "Failed to lock %s in exclusive mode" % self.filename)
2237

    
2238
  def Shared(self, blocking=False, timeout=None):
2239
    """Locks the file in shared mode.
2240

2241
    @type blocking: boolean
2242
    @param blocking: whether to block and wait until we
2243
        can lock the file or return immediately
2244
    @type timeout: int or None
2245
    @param timeout: if not None, the duration to wait for the lock
2246
        (in blocking mode)
2247

2248
    """
2249
    self._flock(fcntl.LOCK_SH, blocking, timeout,
2250
                "Failed to lock %s in shared mode" % self.filename)
2251

    
2252
  def Unlock(self, blocking=True, timeout=None):
2253
    """Unlocks the file.
2254

2255
    According to C{flock(2)}, unlocking can also be a nonblocking
2256
    operation::
2257

2258
      To make a non-blocking request, include LOCK_NB with any of the above
2259
      operations.
2260

2261
    @type blocking: boolean
2262
    @param blocking: whether to block and wait until we
2263
        can lock the file or return immediately
2264
    @type timeout: int or None
2265
    @param timeout: if not None, the duration to wait for the lock
2266
        (in blocking mode)
2267

2268
    """
2269
    self._flock(fcntl.LOCK_UN, blocking, timeout,
2270
                "Failed to unlock %s" % self.filename)
2271

    
2272

    
2273
def SignalHandled(signums):
2274
  """Signal Handled decoration.
2275

2276
  This special decorator installs a signal handler and then calls the target
2277
  function. The function must accept a 'signal_handlers' keyword argument,
2278
  which will contain a dict indexed by signal number, with SignalHandler
2279
  objects as values.
2280

2281
  The decorator can be safely stacked with iself, to handle multiple signals
2282
  with different handlers.
2283

2284
  @type signums: list
2285
  @param signums: signals to intercept
2286

2287
  """
2288
  def wrap(fn):
2289
    def sig_function(*args, **kwargs):
2290
      assert 'signal_handlers' not in kwargs or \
2291
             kwargs['signal_handlers'] is None or \
2292
             isinstance(kwargs['signal_handlers'], dict), \
2293
             "Wrong signal_handlers parameter in original function call"
2294
      if 'signal_handlers' in kwargs and kwargs['signal_handlers'] is not None:
2295
        signal_handlers = kwargs['signal_handlers']
2296
      else:
2297
        signal_handlers = {}
2298
        kwargs['signal_handlers'] = signal_handlers
2299
      sighandler = SignalHandler(signums)
2300
      try:
2301
        for sig in signums:
2302
          signal_handlers[sig] = sighandler
2303
        return fn(*args, **kwargs)
2304
      finally:
2305
        sighandler.Reset()
2306
    return sig_function
2307
  return wrap
2308

    
2309

    
2310
class SignalHandler(object):
2311
  """Generic signal handler class.
2312

2313
  It automatically restores the original handler when deconstructed or
2314
  when L{Reset} is called. You can either pass your own handler
2315
  function in or query the L{called} attribute to detect whether the
2316
  signal was sent.
2317

2318
  @type signum: list
2319
  @ivar signum: the signals we handle
2320
  @type called: boolean
2321
  @ivar called: tracks whether any of the signals have been raised
2322

2323
  """
2324
  def __init__(self, signum):
2325
    """Constructs a new SignalHandler instance.
2326

2327
    @type signum: int or list of ints
2328
    @param signum: Single signal number or set of signal numbers
2329

2330
    """
2331
    self.signum = set(signum)
2332
    self.called = False
2333

    
2334
    self._previous = {}
2335
    try:
2336
      for signum in self.signum:
2337
        # Setup handler
2338
        prev_handler = signal.signal(signum, self._HandleSignal)
2339
        try:
2340
          self._previous[signum] = prev_handler
2341
        except:
2342
          # Restore previous handler
2343
          signal.signal(signum, prev_handler)
2344
          raise
2345
    except:
2346
      # Reset all handlers
2347
      self.Reset()
2348
      # Here we have a race condition: a handler may have already been called,
2349
      # but there's not much we can do about it at this point.
2350
      raise
2351

    
2352
  def __del__(self):
2353
    self.Reset()
2354

    
2355
  def Reset(self):
2356
    """Restore previous handler.
2357

2358
    This will reset all the signals to their previous handlers.
2359

2360
    """
2361
    for signum, prev_handler in self._previous.items():
2362
      signal.signal(signum, prev_handler)
2363
      # If successful, remove from dict
2364
      del self._previous[signum]
2365

    
2366
  def Clear(self):
2367
    """Unsets the L{called} flag.
2368

2369
    This function can be used in case a signal may arrive several times.
2370

2371
    """
2372
    self.called = False
2373

    
2374
  def _HandleSignal(self, signum, frame):
2375
    """Actual signal handling function.
2376

2377
    """
2378
    # This is not nice and not absolutely atomic, but it appears to be the only
2379
    # solution in Python -- there are no atomic types.
2380
    self.called = True
2381

    
2382

    
2383
class FieldSet(object):
2384
  """A simple field set.
2385

2386
  Among the features are:
2387
    - checking if a string is among a list of static string or regex objects
2388
    - checking if a whole list of string matches
2389
    - returning the matching groups from a regex match
2390

2391
  Internally, all fields are held as regular expression objects.
2392

2393
  """
2394
  def __init__(self, *items):
2395
    self.items = [re.compile("^%s$" % value) for value in items]
2396

    
2397
  def Extend(self, other_set):
2398
    """Extend the field set with the items from another one"""
2399
    self.items.extend(other_set.items)
2400

    
2401
  def Matches(self, field):
2402
    """Checks if a field matches the current set
2403

2404
    @type field: str
2405
    @param field: the string to match
2406
    @return: either False or a regular expression match object
2407

2408
    """
2409
    for m in itertools.ifilter(None, (val.match(field) for val in self.items)):
2410
      return m
2411
    return False
2412

    
2413
  def NonMatching(self, items):
2414
    """Returns the list of fields not matching the current set
2415

2416
    @type items: list
2417
    @param items: the list of fields to check
2418
    @rtype: list
2419
    @return: list of non-matching fields
2420

2421
    """
2422
    return [val for val in items if not self.Matches(val)]