Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ 016308cb

History | View | Annotate | Download (64 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
      dirname = os.path.dirname(new)
323
      try:
324
        os.makedirs(dirname, mode=mkdir_mode)
325
      except OSError, err:
326
        # Ignore EEXIST. This is only handled in os.makedirs as included in
327
        # Python 2.5 and above.
328
        if err.errno != errno.EEXIST or not os.path.exists(dirname):
329
          raise
330

    
331
      return os.rename(old, new)
332

    
333
    raise
334

    
335

    
336
def _FingerprintFile(filename):
337
  """Compute the fingerprint of a file.
338

339
  If the file does not exist, a None will be returned
340
  instead.
341

342
  @type filename: str
343
  @param filename: the filename to checksum
344
  @rtype: str
345
  @return: the hex digest of the sha checksum of the contents
346
      of the file
347

348
  """
349
  if not (os.path.exists(filename) and os.path.isfile(filename)):
350
    return None
351

    
352
  f = open(filename)
353

    
354
  fp = sha1()
355
  while True:
356
    data = f.read(4096)
357
    if not data:
358
      break
359

    
360
    fp.update(data)
361

    
362
  return fp.hexdigest()
363

    
364

    
365
def FingerprintFiles(files):
366
  """Compute fingerprints for a list of files.
367

368
  @type files: list
369
  @param files: the list of filename to fingerprint
370
  @rtype: dict
371
  @return: a dictionary filename: fingerprint, holding only
372
      existing files
373

374
  """
375
  ret = {}
376

    
377
  for filename in files:
378
    cksum = _FingerprintFile(filename)
379
    if cksum:
380
      ret[filename] = cksum
381

    
382
  return ret
383

    
384

    
385
def ForceDictType(target, key_types, allowed_values=None):
386
  """Force the values of a dict to have certain types.
387

388
  @type target: dict
389
  @param target: the dict to update
390
  @type key_types: dict
391
  @param key_types: dict mapping target dict keys to types
392
                    in constants.ENFORCEABLE_TYPES
393
  @type allowed_values: list
394
  @keyword allowed_values: list of specially allowed values
395

396
  """
397
  if allowed_values is None:
398
    allowed_values = []
399

    
400
  if not isinstance(target, dict):
401
    msg = "Expected dictionary, got '%s'" % target
402
    raise errors.TypeEnforcementError(msg)
403

    
404
  for key in target:
405
    if key not in key_types:
406
      msg = "Unknown key '%s'" % key
407
      raise errors.TypeEnforcementError(msg)
408

    
409
    if target[key] in allowed_values:
410
      continue
411

    
412
    ktype = key_types[key]
413
    if ktype not in constants.ENFORCEABLE_TYPES:
414
      msg = "'%s' has non-enforceable type %s" % (key, ktype)
415
      raise errors.ProgrammerError(msg)
416

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

    
451

    
452
def IsProcessAlive(pid):
453
  """Check if a given pid exists on the system.
454

455
  @note: zombie status is not handled, so zombie processes
456
      will be returned as alive
457
  @type pid: int
458
  @param pid: the process ID to check
459
  @rtype: boolean
460
  @return: True if the process exists
461

462
  """
463
  if pid <= 0:
464
    return False
465

    
466
  try:
467
    os.stat("/proc/%d/status" % pid)
468
    return True
469
  except EnvironmentError, err:
470
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
471
      return False
472
    raise
473

    
474

    
475
def ReadPidFile(pidfile):
476
  """Read a pid from a file.
477

478
  @type  pidfile: string
479
  @param pidfile: path to the file containing the pid
480
  @rtype: int
481
  @return: The process id, if the file exists and contains a valid PID,
482
           otherwise 0
483

484
  """
485
  try:
486
    raw_data = ReadFile(pidfile)
487
  except EnvironmentError, err:
488
    if err.errno != errno.ENOENT:
489
      logging.exception("Can't read pid file")
490
    return 0
491

    
492
  try:
493
    pid = int(raw_data)
494
  except ValueError, err:
495
    logging.info("Can't parse pid file contents", exc_info=True)
496
    return 0
497

    
498
  return pid
499

    
500

    
501
def MatchNameComponent(key, name_list, case_sensitive=True):
502
  """Try to match a name against a list.
503

504
  This function will try to match a name like test1 against a list
505
  like C{['test1.example.com', 'test2.example.com', ...]}. Against
506
  this list, I{'test1'} as well as I{'test1.example'} will match, but
507
  not I{'test1.ex'}. A multiple match will be considered as no match
508
  at all (e.g. I{'test1'} against C{['test1.example.com',
509
  'test1.example.org']}), except when the key fully matches an entry
510
  (e.g. I{'test1'} against C{['test1', 'test1.example.com']}).
511

512
  @type key: str
513
  @param key: the name to be searched
514
  @type name_list: list
515
  @param name_list: the list of strings against which to search the key
516
  @type case_sensitive: boolean
517
  @param case_sensitive: whether to provide a case-sensitive match
518

519
  @rtype: None or str
520
  @return: None if there is no match I{or} if there are multiple matches,
521
      otherwise the element from the list which matches
522

523
  """
524
  if key in name_list:
525
    return key
526

    
527
  re_flags = 0
528
  if not case_sensitive:
529
    re_flags |= re.IGNORECASE
530
    key = key.upper()
531
  mo = re.compile("^%s(\..*)?$" % re.escape(key), re_flags)
532
  names_filtered = []
533
  string_matches = []
534
  for name in name_list:
535
    if mo.match(name) is not None:
536
      names_filtered.append(name)
537
      if not case_sensitive and key == name.upper():
538
        string_matches.append(name)
539

    
540
  if len(string_matches) == 1:
541
    return string_matches[0]
542
  if len(names_filtered) == 1:
543
    return names_filtered[0]
544
  return None
545

    
546

    
547
class HostInfo:
548
  """Class implementing resolver and hostname functionality
549

550
  """
551
  def __init__(self, name=None):
552
    """Initialize the host name object.
553

554
    If the name argument is not passed, it will use this system's
555
    name.
556

557
    """
558
    if name is None:
559
      name = self.SysName()
560

    
561
    self.query = name
562
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
563
    self.ip = self.ipaddrs[0]
564

    
565
  def ShortName(self):
566
    """Returns the hostname without domain.
567

568
    """
569
    return self.name.split('.')[0]
570

    
571
  @staticmethod
572
  def SysName():
573
    """Return the current system's name.
574

575
    This is simply a wrapper over C{socket.gethostname()}.
576

577
    """
578
    return socket.gethostname()
579

    
580
  @staticmethod
581
  def LookupHostname(hostname):
582
    """Look up hostname
583

584
    @type hostname: str
585
    @param hostname: hostname to look up
586

587
    @rtype: tuple
588
    @return: a tuple (name, aliases, ipaddrs) as returned by
589
        C{socket.gethostbyname_ex}
590
    @raise errors.ResolverError: in case of errors in resolving
591

592
    """
593
    try:
594
      result = socket.gethostbyname_ex(hostname)
595
    except socket.gaierror, err:
596
      # hostname not found in DNS
597
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
598

    
599
    return result
600

    
601

    
602
def GetHostInfo(name=None):
603
  """Lookup host name and raise an OpPrereqError for failures"""
604

    
605
  try:
606
    return HostInfo(name)
607
  except errors.ResolverError, err:
608
    raise errors.OpPrereqError("The given name (%s) does not resolve: %s" %
609
                               (err[0], err[2]), errors.ECODE_RESOLVER)
610

    
611

    
612
def ListVolumeGroups():
613
  """List volume groups and their size
614

615
  @rtype: dict
616
  @return:
617
       Dictionary with keys volume name and values
618
       the size of the volume
619

620
  """
621
  command = "vgs --noheadings --units m --nosuffix -o name,size"
622
  result = RunCmd(command)
623
  retval = {}
624
  if result.failed:
625
    return retval
626

    
627
  for line in result.stdout.splitlines():
628
    try:
629
      name, size = line.split()
630
      size = int(float(size))
631
    except (IndexError, ValueError), err:
632
      logging.error("Invalid output from vgs (%s): %s", err, line)
633
      continue
634

    
635
    retval[name] = size
636

    
637
  return retval
638

    
639

    
640
def BridgeExists(bridge):
641
  """Check whether the given bridge exists in the system
642

643
  @type bridge: str
644
  @param bridge: the bridge name to check
645
  @rtype: boolean
646
  @return: True if it does
647

648
  """
649
  return os.path.isdir("/sys/class/net/%s/bridge" % bridge)
650

    
651

    
652
def NiceSort(name_list):
653
  """Sort a list of strings based on digit and non-digit groupings.
654

655
  Given a list of names C{['a1', 'a10', 'a11', 'a2']} this function
656
  will sort the list in the logical order C{['a1', 'a2', 'a10',
657
  'a11']}.
658

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

663
  @type name_list: list
664
  @param name_list: the names to be sorted
665
  @rtype: list
666
  @return: a copy of the name list sorted with our algorithm
667

668
  """
669
  _SORTER_BASE = "(\D+|\d+)"
670
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
671
                                                  _SORTER_BASE, _SORTER_BASE,
672
                                                  _SORTER_BASE, _SORTER_BASE,
673
                                                  _SORTER_BASE, _SORTER_BASE)
674
  _SORTER_RE = re.compile(_SORTER_FULL)
675
  _SORTER_NODIGIT = re.compile("^\D*$")
676
  def _TryInt(val):
677
    """Attempts to convert a variable to integer."""
678
    if val is None or _SORTER_NODIGIT.match(val):
679
      return val
680
    rval = int(val)
681
    return rval
682

    
683
  to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
684
             for name in name_list]
685
  to_sort.sort()
686
  return [tup[1] for tup in to_sort]
687

    
688

    
689
def TryConvert(fn, val):
690
  """Try to convert a value ignoring errors.
691

692
  This function tries to apply function I{fn} to I{val}. If no
693
  C{ValueError} or C{TypeError} exceptions are raised, it will return
694
  the result, else it will return the original value. Any other
695
  exceptions are propagated to the caller.
696

697
  @type fn: callable
698
  @param fn: function to apply to the value
699
  @param val: the value to be converted
700
  @return: The converted value if the conversion was successful,
701
      otherwise the original value.
702

703
  """
704
  try:
705
    nv = fn(val)
706
  except (ValueError, TypeError):
707
    nv = val
708
  return nv
709

    
710

    
711
def IsValidIP(ip):
712
  """Verifies the syntax of an IPv4 address.
713

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

717
  @type ip: str
718
  @param ip: the address to be checked
719
  @rtype: a regular expression match object
720
  @return: a regular expression match object, or None if the
721
      address is not valid
722

723
  """
724
  unit = "(0|[1-9]\d{0,2})"
725
  #TODO: convert and return only boolean
726
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
727

    
728

    
729
def IsValidShellParam(word):
730
  """Verifies is the given word is safe from the shell's p.o.v.
731

732
  This means that we can pass this to a command via the shell and be
733
  sure that it doesn't alter the command line and is passed as such to
734
  the actual command.
735

736
  Note that we are overly restrictive here, in order to be on the safe
737
  side.
738

739
  @type word: str
740
  @param word: the word to check
741
  @rtype: boolean
742
  @return: True if the word is 'safe'
743

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

    
747

    
748
def BuildShellCmd(template, *args):
749
  """Build a safe shell command line from the given arguments.
750

751
  This function will check all arguments in the args list so that they
752
  are valid shell parameters (i.e. they don't contain shell
753
  metacharacters). If everything is ok, it will return the result of
754
  template % args.
755

756
  @type template: str
757
  @param template: the string holding the template for the
758
      string formatting
759
  @rtype: str
760
  @return: the expanded command line
761

762
  """
763
  for word in args:
764
    if not IsValidShellParam(word):
765
      raise errors.ProgrammerError("Shell argument '%s' contains"
766
                                   " invalid characters" % word)
767
  return template % args
768

    
769

    
770
def FormatUnit(value, units):
771
  """Formats an incoming number of MiB with the appropriate unit.
772

773
  @type value: int
774
  @param value: integer representing the value in MiB (1048576)
775
  @type units: char
776
  @param units: the type of formatting we should do:
777
      - 'h' for automatic scaling
778
      - 'm' for MiBs
779
      - 'g' for GiBs
780
      - 't' for TiBs
781
  @rtype: str
782
  @return: the formatted value (with suffix)
783

784
  """
785
  if units not in ('m', 'g', 't', 'h'):
786
    raise errors.ProgrammerError("Invalid unit specified '%s'" % str(units))
787

    
788
  suffix = ''
789

    
790
  if units == 'm' or (units == 'h' and value < 1024):
791
    if units == 'h':
792
      suffix = 'M'
793
    return "%d%s" % (round(value, 0), suffix)
794

    
795
  elif units == 'g' or (units == 'h' and value < (1024 * 1024)):
796
    if units == 'h':
797
      suffix = 'G'
798
    return "%0.1f%s" % (round(float(value) / 1024, 1), suffix)
799

    
800
  else:
801
    if units == 'h':
802
      suffix = 'T'
803
    return "%0.1f%s" % (round(float(value) / 1024 / 1024, 1), suffix)
804

    
805

    
806
def ParseUnit(input_string):
807
  """Tries to extract number and scale from the given string.
808

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

813
  """
814
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', str(input_string))
815
  if not m:
816
    raise errors.UnitParseError("Invalid format")
817

    
818
  value = float(m.groups()[0])
819

    
820
  unit = m.groups()[1]
821
  if unit:
822
    lcunit = unit.lower()
823
  else:
824
    lcunit = 'm'
825

    
826
  if lcunit in ('m', 'mb', 'mib'):
827
    # Value already in MiB
828
    pass
829

    
830
  elif lcunit in ('g', 'gb', 'gib'):
831
    value *= 1024
832

    
833
  elif lcunit in ('t', 'tb', 'tib'):
834
    value *= 1024 * 1024
835

    
836
  else:
837
    raise errors.UnitParseError("Unknown unit: %s" % unit)
838

    
839
  # Make sure we round up
840
  if int(value) < value:
841
    value += 1
842

    
843
  # Round up to the next multiple of 4
844
  value = int(value)
845
  if value % 4:
846
    value += 4 - value % 4
847

    
848
  return value
849

    
850

    
851
def AddAuthorizedKey(file_name, key):
852
  """Adds an SSH public key to an authorized_keys file.
853

854
  @type file_name: str
855
  @param file_name: path to authorized_keys file
856
  @type key: str
857
  @param key: string containing key
858

859
  """
860
  key_fields = key.split()
861

    
862
  f = open(file_name, 'a+')
863
  try:
864
    nl = True
865
    for line in f:
866
      # Ignore whitespace changes
867
      if line.split() == key_fields:
868
        break
869
      nl = line.endswith('\n')
870
    else:
871
      if not nl:
872
        f.write("\n")
873
      f.write(key.rstrip('\r\n'))
874
      f.write("\n")
875
      f.flush()
876
  finally:
877
    f.close()
878

    
879

    
880
def RemoveAuthorizedKey(file_name, key):
881
  """Removes an SSH public key from an authorized_keys file.
882

883
  @type file_name: str
884
  @param file_name: path to authorized_keys file
885
  @type key: str
886
  @param key: string containing key
887

888
  """
889
  key_fields = key.split()
890

    
891
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
892
  try:
893
    out = os.fdopen(fd, 'w')
894
    try:
895
      f = open(file_name, 'r')
896
      try:
897
        for line in f:
898
          # Ignore whitespace changes while comparing lines
899
          if line.split() != key_fields:
900
            out.write(line)
901

    
902
        out.flush()
903
        os.rename(tmpname, file_name)
904
      finally:
905
        f.close()
906
    finally:
907
      out.close()
908
  except:
909
    RemoveFile(tmpname)
910
    raise
911

    
912

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

916
  @type file_name: str
917
  @param file_name: path to the file to modify (usually C{/etc/hosts})
918
  @type ip: str
919
  @param ip: the IP address
920
  @type hostname: str
921
  @param hostname: the hostname to be added
922
  @type aliases: list
923
  @param aliases: the list of aliases to add for the hostname
924

925
  """
926
  # FIXME: use WriteFile + fn rather than duplicating its efforts
927
  # Ensure aliases are unique
928
  aliases = UniqueSequence([hostname] + aliases)[1:]
929

    
930
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
931
  try:
932
    out = os.fdopen(fd, 'w')
933
    try:
934
      f = open(file_name, 'r')
935
      try:
936
        for line in f:
937
          fields = line.split()
938
          if fields and not fields[0].startswith('#') and ip == fields[0]:
939
            continue
940
          out.write(line)
941

    
942
        out.write("%s\t%s" % (ip, hostname))
943
        if aliases:
944
          out.write(" %s" % ' '.join(aliases))
945
        out.write('\n')
946

    
947
        out.flush()
948
        os.fsync(out)
949
        os.chmod(tmpname, 0644)
950
        os.rename(tmpname, file_name)
951
      finally:
952
        f.close()
953
    finally:
954
      out.close()
955
  except:
956
    RemoveFile(tmpname)
957
    raise
958

    
959

    
960
def AddHostToEtcHosts(hostname):
961
  """Wrapper around SetEtcHostsEntry.
962

963
  @type hostname: str
964
  @param hostname: a hostname that will be resolved and added to
965
      L{constants.ETC_HOSTS}
966

967
  """
968
  hi = HostInfo(name=hostname)
969
  SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])
970

    
971

    
972
def RemoveEtcHostsEntry(file_name, hostname):
973
  """Removes a hostname from /etc/hosts.
974

975
  IP addresses without names are removed from the file.
976

977
  @type file_name: str
978
  @param file_name: path to the file to modify (usually C{/etc/hosts})
979
  @type hostname: str
980
  @param hostname: the hostname to be removed
981

982
  """
983
  # FIXME: use WriteFile + fn rather than duplicating its efforts
984
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
985
  try:
986
    out = os.fdopen(fd, 'w')
987
    try:
988
      f = open(file_name, 'r')
989
      try:
990
        for line in f:
991
          fields = line.split()
992
          if len(fields) > 1 and not fields[0].startswith('#'):
993
            names = fields[1:]
994
            if hostname in names:
995
              while hostname in names:
996
                names.remove(hostname)
997
              if names:
998
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
999
              continue
1000

    
1001
          out.write(line)
1002

    
1003
        out.flush()
1004
        os.fsync(out)
1005
        os.chmod(tmpname, 0644)
1006
        os.rename(tmpname, file_name)
1007
      finally:
1008
        f.close()
1009
    finally:
1010
      out.close()
1011
  except:
1012
    RemoveFile(tmpname)
1013
    raise
1014

    
1015

    
1016
def RemoveHostFromEtcHosts(hostname):
1017
  """Wrapper around RemoveEtcHostsEntry.
1018

1019
  @type hostname: str
1020
  @param hostname: hostname that will be resolved and its
1021
      full and shot name will be removed from
1022
      L{constants.ETC_HOSTS}
1023

1024
  """
1025
  hi = HostInfo(name=hostname)
1026
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.name)
1027
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.ShortName())
1028

    
1029

    
1030
def CreateBackup(file_name):
1031
  """Creates a backup of a file.
1032

1033
  @type file_name: str
1034
  @param file_name: file to be backed up
1035
  @rtype: str
1036
  @return: the path to the newly created backup
1037
  @raise errors.ProgrammerError: for invalid file names
1038

1039
  """
1040
  if not os.path.isfile(file_name):
1041
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
1042
                                file_name)
1043

    
1044
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
1045
  dir_name = os.path.dirname(file_name)
1046

    
1047
  fsrc = open(file_name, 'rb')
1048
  try:
1049
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
1050
    fdst = os.fdopen(fd, 'wb')
1051
    try:
1052
      shutil.copyfileobj(fsrc, fdst)
1053
    finally:
1054
      fdst.close()
1055
  finally:
1056
    fsrc.close()
1057

    
1058
  return backup_name
1059

    
1060

    
1061
def ShellQuote(value):
1062
  """Quotes shell argument according to POSIX.
1063

1064
  @type value: str
1065
  @param value: the argument to be quoted
1066
  @rtype: str
1067
  @return: the quoted value
1068

1069
  """
1070
  if _re_shell_unquoted.match(value):
1071
    return value
1072
  else:
1073
    return "'%s'" % value.replace("'", "'\\''")
1074

    
1075

    
1076
def ShellQuoteArgs(args):
1077
  """Quotes a list of shell arguments.
1078

1079
  @type args: list
1080
  @param args: list of arguments to be quoted
1081
  @rtype: str
1082
  @return: the quoted arguments concatenated with spaces
1083

1084
  """
1085
  return ' '.join([ShellQuote(i) for i in args])
1086

    
1087

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

1091
  Check if the given IP is reachable by doing attempting a TCP connect
1092
  to it.
1093

1094
  @type target: str
1095
  @param target: the IP or hostname to ping
1096
  @type port: int
1097
  @param port: the port to connect to
1098
  @type timeout: int
1099
  @param timeout: the timeout on the connection attempt
1100
  @type live_port_needed: boolean
1101
  @param live_port_needed: whether a closed port will cause the
1102
      function to return failure, as if there was a timeout
1103
  @type source: str or None
1104
  @param source: if specified, will cause the connect to be made
1105
      from this specific source address; failures to bind other
1106
      than C{EADDRNOTAVAIL} will be ignored
1107

1108
  """
1109
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
1110

    
1111
  success = False
1112

    
1113
  if source is not None:
1114
    try:
1115
      sock.bind((source, 0))
1116
    except socket.error, (errcode, _):
1117
      if errcode == errno.EADDRNOTAVAIL:
1118
        success = False
1119

    
1120
  sock.settimeout(timeout)
1121

    
1122
  try:
1123
    sock.connect((target, port))
1124
    sock.close()
1125
    success = True
1126
  except socket.timeout:
1127
    success = False
1128
  except socket.error, (errcode, _):
1129
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
1130

    
1131
  return success
1132

    
1133

    
1134
def OwnIpAddress(address):
1135
  """Check if the current host has the the given IP address.
1136

1137
  Currently this is done by TCP-pinging the address from the loopback
1138
  address.
1139

1140
  @type address: string
1141
  @param address: the address to check
1142
  @rtype: bool
1143
  @return: True if we own the address
1144

1145
  """
1146
  return TcpPing(address, constants.DEFAULT_NODED_PORT,
1147
                 source=constants.LOCALHOST_IP_ADDRESS)
1148

    
1149

    
1150
def ListVisibleFiles(path):
1151
  """Returns a list of visible files in a directory.
1152

1153
  @type path: str
1154
  @param path: the directory to enumerate
1155
  @rtype: list
1156
  @return: the list of all files not starting with a dot
1157

1158
  """
1159
  files = [i for i in os.listdir(path) if not i.startswith(".")]
1160
  files.sort()
1161
  return files
1162

    
1163

    
1164
def GetHomeDir(user, default=None):
1165
  """Try to get the homedir of the given user.
1166

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

1171
  """
1172
  try:
1173
    if isinstance(user, basestring):
1174
      result = pwd.getpwnam(user)
1175
    elif isinstance(user, (int, long)):
1176
      result = pwd.getpwuid(user)
1177
    else:
1178
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
1179
                                   type(user))
1180
  except KeyError:
1181
    return default
1182
  return result.pw_dir
1183

    
1184

    
1185
def NewUUID():
1186
  """Returns a random UUID.
1187

1188
  @note: This is a Linux-specific method as it uses the /proc
1189
      filesystem.
1190
  @rtype: str
1191

1192
  """
1193
  return ReadFile(_RANDOM_UUID_FILE, size=128).rstrip("\n")
1194

    
1195

    
1196
def GenerateSecret(numbytes=20):
1197
  """Generates a random secret.
1198

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

1202
  @param numbytes: the number of bytes which will be represented by the returned
1203
      string (defaulting to 20, the length of a SHA1 hash)
1204
  @rtype: str
1205
  @return: an hex representation of the pseudo-random sequence
1206

1207
  """
1208
  return os.urandom(numbytes).encode('hex')
1209

    
1210

    
1211
def EnsureDirs(dirs):
1212
  """Make required directories, if they don't exist.
1213

1214
  @param dirs: list of tuples (dir_name, dir_mode)
1215
  @type dirs: list of (string, integer)
1216

1217
  """
1218
  for dir_name, dir_mode in dirs:
1219
    try:
1220
      os.mkdir(dir_name, dir_mode)
1221
    except EnvironmentError, err:
1222
      if err.errno != errno.EEXIST:
1223
        raise errors.GenericError("Cannot create needed directory"
1224
                                  " '%s': %s" % (dir_name, err))
1225
    if not os.path.isdir(dir_name):
1226
      raise errors.GenericError("%s is not a directory" % dir_name)
1227

    
1228

    
1229
def ReadFile(file_name, size=-1):
1230
  """Reads a file.
1231

1232
  @type size: int
1233
  @param size: Read at most size bytes (if negative, entire file)
1234
  @rtype: str
1235
  @return: the (possibly partial) content of the file
1236

1237
  """
1238
  f = open(file_name, "r")
1239
  try:
1240
    return f.read(size)
1241
  finally:
1242
    f.close()
1243

    
1244

    
1245
def WriteFile(file_name, fn=None, data=None,
1246
              mode=None, uid=-1, gid=-1,
1247
              atime=None, mtime=None, close=True,
1248
              dry_run=False, backup=False,
1249
              prewrite=None, postwrite=None):
1250
  """(Over)write a file atomically.
1251

1252
  The file_name and either fn (a function taking one argument, the
1253
  file descriptor, and which should write the data to it) or data (the
1254
  contents of the file) must be passed. The other arguments are
1255
  optional and allow setting the file mode, owner and group, and the
1256
  mtime/atime of the file.
1257

1258
  If the function doesn't raise an exception, it has succeeded and the
1259
  target file has the new contents. If the function has raised an
1260
  exception, an existing target file should be unmodified and the
1261
  temporary file should be removed.
1262

1263
  @type file_name: str
1264
  @param file_name: the target filename
1265
  @type fn: callable
1266
  @param fn: content writing function, called with
1267
      file descriptor as parameter
1268
  @type data: str
1269
  @param data: contents of the file
1270
  @type mode: int
1271
  @param mode: file mode
1272
  @type uid: int
1273
  @param uid: the owner of the file
1274
  @type gid: int
1275
  @param gid: the group of the file
1276
  @type atime: int
1277
  @param atime: a custom access time to be set on the file
1278
  @type mtime: int
1279
  @param mtime: a custom modification time to be set on the file
1280
  @type close: boolean
1281
  @param close: whether to close file after writing it
1282
  @type prewrite: callable
1283
  @param prewrite: function to be called before writing content
1284
  @type postwrite: callable
1285
  @param postwrite: function to be called after writing content
1286

1287
  @rtype: None or int
1288
  @return: None if the 'close' parameter evaluates to True,
1289
      otherwise the file descriptor
1290

1291
  @raise errors.ProgrammerError: if any of the arguments are not valid
1292

1293
  """
1294
  if not os.path.isabs(file_name):
1295
    raise errors.ProgrammerError("Path passed to WriteFile is not"
1296
                                 " absolute: '%s'" % file_name)
1297

    
1298
  if [fn, data].count(None) != 1:
1299
    raise errors.ProgrammerError("fn or data required")
1300

    
1301
  if [atime, mtime].count(None) == 1:
1302
    raise errors.ProgrammerError("Both atime and mtime must be either"
1303
                                 " set or None")
1304

    
1305
  if backup and not dry_run and os.path.isfile(file_name):
1306
    CreateBackup(file_name)
1307

    
1308
  dir_name, base_name = os.path.split(file_name)
1309
  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
1310
  do_remove = True
1311
  # here we need to make sure we remove the temp file, if any error
1312
  # leaves it in place
1313
  try:
1314
    if uid != -1 or gid != -1:
1315
      os.chown(new_name, uid, gid)
1316
    if mode:
1317
      os.chmod(new_name, mode)
1318
    if callable(prewrite):
1319
      prewrite(fd)
1320
    if data is not None:
1321
      os.write(fd, data)
1322
    else:
1323
      fn(fd)
1324
    if callable(postwrite):
1325
      postwrite(fd)
1326
    os.fsync(fd)
1327
    if atime is not None and mtime is not None:
1328
      os.utime(new_name, (atime, mtime))
1329
    if not dry_run:
1330
      os.rename(new_name, file_name)
1331
      do_remove = False
1332
  finally:
1333
    if close:
1334
      os.close(fd)
1335
      result = None
1336
    else:
1337
      result = fd
1338
    if do_remove:
1339
      RemoveFile(new_name)
1340

    
1341
  return result
1342

    
1343

    
1344
def FirstFree(seq, base=0):
1345
  """Returns the first non-existing integer from seq.
1346

1347
  The seq argument should be a sorted list of positive integers. The
1348
  first time the index of an element is smaller than the element
1349
  value, the index will be returned.
1350

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

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

1356
  @type seq: sequence
1357
  @param seq: the sequence to be analyzed.
1358
  @type base: int
1359
  @param base: use this value as the base index of the sequence
1360
  @rtype: int
1361
  @return: the first non-used index in the sequence
1362

1363
  """
1364
  for idx, elem in enumerate(seq):
1365
    assert elem >= base, "Passed element is higher than base offset"
1366
    if elem > idx + base:
1367
      # idx is not used
1368
      return idx + base
1369
  return None
1370

    
1371

    
1372
def all(seq, pred=bool):
1373
  "Returns True if pred(x) is True for every element in the iterable"
1374
  for _ in itertools.ifilterfalse(pred, seq):
1375
    return False
1376
  return True
1377

    
1378

    
1379
def any(seq, pred=bool):
1380
  "Returns True if pred(x) is True for at least one element in the iterable"
1381
  for _ in itertools.ifilter(pred, seq):
1382
    return True
1383
  return False
1384

    
1385

    
1386
def UniqueSequence(seq):
1387
  """Returns a list with unique elements.
1388

1389
  Element order is preserved.
1390

1391
  @type seq: sequence
1392
  @param seq: the sequence with the source elements
1393
  @rtype: list
1394
  @return: list of unique elements from seq
1395

1396
  """
1397
  seen = set()
1398
  return [i for i in seq if i not in seen and not seen.add(i)]
1399

    
1400

    
1401
def IsValidMac(mac):
1402
  """Predicate to check if a MAC address is valid.
1403

1404
  Checks whether the supplied MAC address is formally correct, only
1405
  accepts colon separated format.
1406

1407
  @type mac: str
1408
  @param mac: the MAC to be validated
1409
  @rtype: boolean
1410
  @return: True is the MAC seems valid
1411

1412
  """
1413
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
1414
  return mac_check.match(mac) is not None
1415

    
1416

    
1417
def TestDelay(duration):
1418
  """Sleep for a fixed amount of time.
1419

1420
  @type duration: float
1421
  @param duration: the sleep duration
1422
  @rtype: boolean
1423
  @return: False for negative value, True otherwise
1424

1425
  """
1426
  if duration < 0:
1427
    return False, "Invalid sleep duration"
1428
  time.sleep(duration)
1429
  return True, None
1430

    
1431

    
1432
def _CloseFDNoErr(fd, retries=5):
1433
  """Close a file descriptor ignoring errors.
1434

1435
  @type fd: int
1436
  @param fd: the file descriptor
1437
  @type retries: int
1438
  @param retries: how many retries to make, in case we get any
1439
      other error than EBADF
1440

1441
  """
1442
  try:
1443
    os.close(fd)
1444
  except OSError, err:
1445
    if err.errno != errno.EBADF:
1446
      if retries > 0:
1447
        _CloseFDNoErr(fd, retries - 1)
1448
    # else either it's closed already or we're out of retries, so we
1449
    # ignore this and go on
1450

    
1451

    
1452
def CloseFDs(noclose_fds=None):
1453
  """Close file descriptors.
1454

1455
  This closes all file descriptors above 2 (i.e. except
1456
  stdin/out/err).
1457

1458
  @type noclose_fds: list or None
1459
  @param noclose_fds: if given, it denotes a list of file descriptor
1460
      that should not be closed
1461

1462
  """
1463
  # Default maximum for the number of available file descriptors.
1464
  if 'SC_OPEN_MAX' in os.sysconf_names:
1465
    try:
1466
      MAXFD = os.sysconf('SC_OPEN_MAX')
1467
      if MAXFD < 0:
1468
        MAXFD = 1024
1469
    except OSError:
1470
      MAXFD = 1024
1471
  else:
1472
    MAXFD = 1024
1473
  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1474
  if (maxfd == resource.RLIM_INFINITY):
1475
    maxfd = MAXFD
1476

    
1477
  # Iterate through and close all file descriptors (except the standard ones)
1478
  for fd in range(3, maxfd):
1479
    if noclose_fds and fd in noclose_fds:
1480
      continue
1481
    _CloseFDNoErr(fd)
1482

    
1483

    
1484
def Daemonize(logfile):
1485
  """Daemonize the current process.
1486

1487
  This detaches the current process from the controlling terminal and
1488
  runs it in the background as a daemon.
1489

1490
  @type logfile: str
1491
  @param logfile: the logfile to which we should redirect stdout/stderr
1492
  @rtype: int
1493
  @return: the value zero
1494

1495
  """
1496
  UMASK = 077
1497
  WORKDIR = "/"
1498

    
1499
  # this might fail
1500
  pid = os.fork()
1501
  if (pid == 0):  # The first child.
1502
    os.setsid()
1503
    # this might fail
1504
    pid = os.fork() # Fork a second child.
1505
    if (pid == 0):  # The second child.
1506
      os.chdir(WORKDIR)
1507
      os.umask(UMASK)
1508
    else:
1509
      # exit() or _exit()?  See below.
1510
      os._exit(0) # Exit parent (the first child) of the second child.
1511
  else:
1512
    os._exit(0) # Exit parent of the first child.
1513

    
1514
  for fd in range(3):
1515
    _CloseFDNoErr(fd)
1516
  i = os.open("/dev/null", os.O_RDONLY) # stdin
1517
  assert i == 0, "Can't close/reopen stdin"
1518
  i = os.open(logfile, os.O_WRONLY|os.O_CREAT|os.O_APPEND, 0600) # stdout
1519
  assert i == 1, "Can't close/reopen stdout"
1520
  # Duplicate standard output to standard error.
1521
  os.dup2(1, 2)
1522
  return 0
1523

    
1524

    
1525
def DaemonPidFileName(name):
1526
  """Compute a ganeti pid file absolute path
1527

1528
  @type name: str
1529
  @param name: the daemon name
1530
  @rtype: str
1531
  @return: the full path to the pidfile corresponding to the given
1532
      daemon name
1533

1534
  """
1535
  return os.path.join(constants.RUN_GANETI_DIR, "%s.pid" % name)
1536

    
1537

    
1538
def WritePidFile(name):
1539
  """Write the current process pidfile.
1540

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

1543
  @type name: str
1544
  @param name: the daemon name to use
1545
  @raise errors.GenericError: if the pid file already exists and
1546
      points to a live process
1547

1548
  """
1549
  pid = os.getpid()
1550
  pidfilename = DaemonPidFileName(name)
1551
  if IsProcessAlive(ReadPidFile(pidfilename)):
1552
    raise errors.GenericError("%s contains a live process" % pidfilename)
1553

    
1554
  WriteFile(pidfilename, data="%d\n" % pid)
1555

    
1556

    
1557
def RemovePidFile(name):
1558
  """Remove the current process pidfile.
1559

1560
  Any errors are ignored.
1561

1562
  @type name: str
1563
  @param name: the daemon name used to derive the pidfile name
1564

1565
  """
1566
  pidfilename = DaemonPidFileName(name)
1567
  # TODO: we could check here that the file contains our pid
1568
  try:
1569
    RemoveFile(pidfilename)
1570
  except:
1571
    pass
1572

    
1573

    
1574
def KillProcess(pid, signal_=signal.SIGTERM, timeout=30,
1575
                waitpid=False):
1576
  """Kill a process given by its pid.
1577

1578
  @type pid: int
1579
  @param pid: The PID to terminate.
1580
  @type signal_: int
1581
  @param signal_: The signal to send, by default SIGTERM
1582
  @type timeout: int
1583
  @param timeout: The timeout after which, if the process is still alive,
1584
                  a SIGKILL will be sent. If not positive, no such checking
1585
                  will be done
1586
  @type waitpid: boolean
1587
  @param waitpid: If true, we should waitpid on this process after
1588
      sending signals, since it's our own child and otherwise it
1589
      would remain as zombie
1590

1591
  """
1592
  def _helper(pid, signal_, wait):
1593
    """Simple helper to encapsulate the kill/waitpid sequence"""
1594
    os.kill(pid, signal_)
1595
    if wait:
1596
      try:
1597
        os.waitpid(pid, os.WNOHANG)
1598
      except OSError:
1599
        pass
1600

    
1601
  if pid <= 0:
1602
    # kill with pid=0 == suicide
1603
    raise errors.ProgrammerError("Invalid pid given '%s'" % pid)
1604

    
1605
  if not IsProcessAlive(pid):
1606
    return
1607

    
1608
  _helper(pid, signal_, waitpid)
1609

    
1610
  if timeout <= 0:
1611
    return
1612

    
1613
  def _CheckProcess():
1614
    if not IsProcessAlive(pid):
1615
      return
1616

    
1617
    try:
1618
      (result_pid, _) = os.waitpid(pid, os.WNOHANG)
1619
    except OSError:
1620
      raise RetryAgain()
1621

    
1622
    if result_pid > 0:
1623
      return
1624

    
1625
    raise RetryAgain()
1626

    
1627
  try:
1628
    # Wait up to $timeout seconds
1629
    Retry(_CheckProcess, (0.01, 1.5, 0.1), timeout)
1630
  except RetryTimeout:
1631
    pass
1632

    
1633
  if IsProcessAlive(pid):
1634
    # Kill process if it's still alive
1635
    _helper(pid, signal.SIGKILL, waitpid)
1636

    
1637

    
1638
def FindFile(name, search_path, test=os.path.exists):
1639
  """Look for a filesystem object in a given path.
1640

1641
  This is an abstract method to search for filesystem object (files,
1642
  dirs) under a given search path.
1643

1644
  @type name: str
1645
  @param name: the name to look for
1646
  @type search_path: str
1647
  @param search_path: location to start at
1648
  @type test: callable
1649
  @param test: a function taking one argument that should return True
1650
      if the a given object is valid; the default value is
1651
      os.path.exists, causing only existing files to be returned
1652
  @rtype: str or None
1653
  @return: full path to the object if found, None otherwise
1654

1655
  """
1656
  for dir_name in search_path:
1657
    item_name = os.path.sep.join([dir_name, name])
1658
    if test(item_name):
1659
      return item_name
1660
  return None
1661

    
1662

    
1663
def CheckVolumeGroupSize(vglist, vgname, minsize):
1664
  """Checks if the volume group list is valid.
1665

1666
  The function will check if a given volume group is in the list of
1667
  volume groups and has a minimum size.
1668

1669
  @type vglist: dict
1670
  @param vglist: dictionary of volume group names and their size
1671
  @type vgname: str
1672
  @param vgname: the volume group we should check
1673
  @type minsize: int
1674
  @param minsize: the minimum size we accept
1675
  @rtype: None or str
1676
  @return: None for success, otherwise the error message
1677

1678
  """
1679
  vgsize = vglist.get(vgname, None)
1680
  if vgsize is None:
1681
    return "volume group '%s' missing" % vgname
1682
  elif vgsize < minsize:
1683
    return ("volume group '%s' too small (%s MiB required, %d MiB found)" %
1684
            (vgname, minsize, vgsize))
1685
  return None
1686

    
1687

    
1688
def SplitTime(value):
1689
  """Splits time as floating point number into a tuple.
1690

1691
  @param value: Time in seconds
1692
  @type value: int or float
1693
  @return: Tuple containing (seconds, microseconds)
1694

1695
  """
1696
  (seconds, microseconds) = divmod(int(value * 1000000), 1000000)
1697

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

    
1703
  return (int(seconds), int(microseconds))
1704

    
1705

    
1706
def MergeTime(timetuple):
1707
  """Merges a tuple into time as a floating point number.
1708

1709
  @param timetuple: Time as tuple, (seconds, microseconds)
1710
  @type timetuple: tuple
1711
  @return: Time as a floating point number expressed in seconds
1712

1713
  """
1714
  (seconds, microseconds) = timetuple
1715

    
1716
  assert 0 <= seconds, \
1717
    "Seconds must be larger than or equal to 0, but are %s" % seconds
1718
  assert 0 <= microseconds <= 999999, \
1719
    "Microseconds must be 0-999999, but are %s" % microseconds
1720

    
1721
  return float(seconds) + (float(microseconds) * 0.000001)
1722

    
1723

    
1724
def GetDaemonPort(daemon_name):
1725
  """Get the daemon port for this cluster.
1726

1727
  Note that this routine does not read a ganeti-specific file, but
1728
  instead uses C{socket.getservbyname} to allow pre-customization of
1729
  this parameter outside of Ganeti.
1730

1731
  @type daemon_name: string
1732
  @param daemon_name: daemon name (in constants.DAEMONS_PORTS)
1733
  @rtype: int
1734

1735
  """
1736
  if daemon_name not in constants.DAEMONS_PORTS:
1737
    raise errors.ProgrammerError("Unknown daemon: %s" % daemon_name)
1738

    
1739
  (proto, default_port) = constants.DAEMONS_PORTS[daemon_name]
1740
  try:
1741
    port = socket.getservbyname(daemon_name, proto)
1742
  except socket.error:
1743
    port = default_port
1744

    
1745
  return port
1746

    
1747

    
1748
def SetupLogging(logfile, debug=False, stderr_logging=False, program="",
1749
                 multithreaded=False):
1750
  """Configures the logging module.
1751

1752
  @type logfile: str
1753
  @param logfile: the filename to which we should log
1754
  @type debug: boolean
1755
  @param debug: whether to enable debug messages too or
1756
      only those at C{INFO} and above level
1757
  @type stderr_logging: boolean
1758
  @param stderr_logging: whether we should also log to the standard error
1759
  @type program: str
1760
  @param program: the name under which we should log messages
1761
  @type multithreaded: boolean
1762
  @param multithreaded: if True, will add the thread name to the log file
1763
  @raise EnvironmentError: if we can't open the log file and
1764
      stderr logging is disabled
1765

1766
  """
1767
  fmt = "%(asctime)s: " + program + " pid=%(process)d"
1768
  if multithreaded:
1769
    fmt += "/%(threadName)s"
1770
  if debug:
1771
    fmt += " %(module)s:%(lineno)s"
1772
  fmt += " %(levelname)s %(message)s"
1773
  formatter = logging.Formatter(fmt)
1774

    
1775
  root_logger = logging.getLogger("")
1776
  root_logger.setLevel(logging.NOTSET)
1777

    
1778
  # Remove all previously setup handlers
1779
  for handler in root_logger.handlers:
1780
    handler.close()
1781
    root_logger.removeHandler(handler)
1782

    
1783
  if stderr_logging:
1784
    stderr_handler = logging.StreamHandler()
1785
    stderr_handler.setFormatter(formatter)
1786
    if debug:
1787
      stderr_handler.setLevel(logging.NOTSET)
1788
    else:
1789
      stderr_handler.setLevel(logging.CRITICAL)
1790
    root_logger.addHandler(stderr_handler)
1791

    
1792
  # this can fail, if the logging directories are not setup or we have
1793
  # a permisssion problem; in this case, it's best to log but ignore
1794
  # the error if stderr_logging is True, and if false we re-raise the
1795
  # exception since otherwise we could run but without any logs at all
1796
  try:
1797
    logfile_handler = logging.FileHandler(logfile)
1798
    logfile_handler.setFormatter(formatter)
1799
    if debug:
1800
      logfile_handler.setLevel(logging.DEBUG)
1801
    else:
1802
      logfile_handler.setLevel(logging.INFO)
1803
    root_logger.addHandler(logfile_handler)
1804
  except EnvironmentError:
1805
    if stderr_logging:
1806
      logging.exception("Failed to enable logging to file '%s'", logfile)
1807
    else:
1808
      # we need to re-raise the exception
1809
      raise
1810

    
1811

    
1812
def IsNormAbsPath(path):
1813
  """Check whether a path is absolute and also normalized
1814

1815
  This avoids things like /dir/../../other/path to be valid.
1816

1817
  """
1818
  return os.path.normpath(path) == path and os.path.isabs(path)
1819

    
1820

    
1821
def TailFile(fname, lines=20):
1822
  """Return the last lines from a file.
1823

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

1828
  @param fname: the file name
1829
  @type lines: int
1830
  @param lines: the (maximum) number of lines to return
1831

1832
  """
1833
  fd = open(fname, "r")
1834
  try:
1835
    fd.seek(0, 2)
1836
    pos = fd.tell()
1837
    pos = max(0, pos-4096)
1838
    fd.seek(pos, 0)
1839
    raw_data = fd.read()
1840
  finally:
1841
    fd.close()
1842

    
1843
  rows = raw_data.splitlines()
1844
  return rows[-lines:]
1845

    
1846

    
1847
def SafeEncode(text):
1848
  """Return a 'safe' version of a source string.
1849

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

1859
  @type text: str or unicode
1860
  @param text: input data
1861
  @rtype: str
1862
  @return: a safe version of text
1863

1864
  """
1865
  if isinstance(text, unicode):
1866
    # only if unicode; if str already, we handle it below
1867
    text = text.encode('ascii', 'backslashreplace')
1868
  resu = ""
1869
  for char in text:
1870
    c = ord(char)
1871
    if char  == '\t':
1872
      resu += r'\t'
1873
    elif char == '\n':
1874
      resu += r'\n'
1875
    elif char == '\r':
1876
      resu += r'\'r'
1877
    elif c < 32 or c >= 127: # non-printable
1878
      resu += "\\x%02x" % (c & 0xff)
1879
    else:
1880
      resu += char
1881
  return resu
1882

    
1883

    
1884
def CommaJoin(names):
1885
  """Nicely join a set of identifiers.
1886

1887
  @param names: set, list or tuple
1888
  @return: a string with the formatted results
1889

1890
  """
1891
  return ", ".join([str(val) for val in names])
1892

    
1893

    
1894
def BytesToMebibyte(value):
1895
  """Converts bytes to mebibytes.
1896

1897
  @type value: int
1898
  @param value: Value in bytes
1899
  @rtype: int
1900
  @return: Value in mebibytes
1901

1902
  """
1903
  return int(round(value / (1024.0 * 1024.0), 0))
1904

    
1905

    
1906
def CalculateDirectorySize(path):
1907
  """Calculates the size of a directory recursively.
1908

1909
  @type path: string
1910
  @param path: Path to directory
1911
  @rtype: int
1912
  @return: Size in mebibytes
1913

1914
  """
1915
  size = 0
1916

    
1917
  for (curpath, _, files) in os.walk(path):
1918
    for filename in files:
1919
      st = os.lstat(os.path.join(curpath, filename))
1920
      size += st.st_size
1921

    
1922
  return BytesToMebibyte(size)
1923

    
1924

    
1925
def GetFilesystemStats(path):
1926
  """Returns the total and free space on a filesystem.
1927

1928
  @type path: string
1929
  @param path: Path on filesystem to be examined
1930
  @rtype: int
1931
  @return: tuple of (Total space, Free space) in mebibytes
1932

1933
  """
1934
  st = os.statvfs(path)
1935

    
1936
  fsize = BytesToMebibyte(st.f_bavail * st.f_frsize)
1937
  tsize = BytesToMebibyte(st.f_blocks * st.f_frsize)
1938
  return (tsize, fsize)
1939

    
1940

    
1941
def LockedMethod(fn):
1942
  """Synchronized object access decorator.
1943

1944
  This decorator is intended to protect access to an object using the
1945
  object's own lock which is hardcoded to '_lock'.
1946

1947
  """
1948
  def _LockDebug(*args, **kwargs):
1949
    if debug_locks:
1950
      logging.debug(*args, **kwargs)
1951

    
1952
  def wrapper(self, *args, **kwargs):
1953
    assert hasattr(self, '_lock')
1954
    lock = self._lock
1955
    _LockDebug("Waiting for %s", lock)
1956
    lock.acquire()
1957
    try:
1958
      _LockDebug("Acquired %s", lock)
1959
      result = fn(self, *args, **kwargs)
1960
    finally:
1961
      _LockDebug("Releasing %s", lock)
1962
      lock.release()
1963
      _LockDebug("Released %s", lock)
1964
    return result
1965
  return wrapper
1966

    
1967

    
1968
def LockFile(fd):
1969
  """Locks a file using POSIX locks.
1970

1971
  @type fd: int
1972
  @param fd: the file descriptor we need to lock
1973

1974
  """
1975
  try:
1976
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
1977
  except IOError, err:
1978
    if err.errno == errno.EAGAIN:
1979
      raise errors.LockError("File already locked")
1980
    raise
1981

    
1982

    
1983
def FormatTime(val):
1984
  """Formats a time value.
1985

1986
  @type val: float or None
1987
  @param val: the timestamp as returned by time.time()
1988
  @return: a string value or N/A if we don't have a valid timestamp
1989

1990
  """
1991
  if val is None or not isinstance(val, (int, float)):
1992
    return "N/A"
1993
  # these two codes works on Linux, but they are not guaranteed on all
1994
  # platforms
1995
  return time.strftime("%F %T", time.localtime(val))
1996

    
1997

    
1998
def ReadWatcherPauseFile(filename, now=None, remove_after=3600):
1999
  """Reads the watcher pause file.
2000

2001
  @type filename: string
2002
  @param filename: Path to watcher pause file
2003
  @type now: None, float or int
2004
  @param now: Current time as Unix timestamp
2005
  @type remove_after: int
2006
  @param remove_after: Remove watcher pause file after specified amount of
2007
    seconds past the pause end time
2008

2009
  """
2010
  if now is None:
2011
    now = time.time()
2012

    
2013
  try:
2014
    value = ReadFile(filename)
2015
  except IOError, err:
2016
    if err.errno != errno.ENOENT:
2017
      raise
2018
    value = None
2019

    
2020
  if value is not None:
2021
    try:
2022
      value = int(value)
2023
    except ValueError:
2024
      logging.warning(("Watcher pause file (%s) contains invalid value,"
2025
                       " removing it"), filename)
2026
      RemoveFile(filename)
2027
      value = None
2028

    
2029
    if value is not None:
2030
      # Remove file if it's outdated
2031
      if now > (value + remove_after):
2032
        RemoveFile(filename)
2033
        value = None
2034

    
2035
      elif now > value:
2036
        value = None
2037

    
2038
  return value
2039

    
2040

    
2041
class RetryTimeout(Exception):
2042
  """Retry loop timed out.
2043

2044
  """
2045

    
2046

    
2047
class RetryAgain(Exception):
2048
  """Retry again.
2049

2050
  """
2051

    
2052

    
2053
class _RetryDelayCalculator(object):
2054
  """Calculator for increasing delays.
2055

2056
  """
2057
  __slots__ = [
2058
    "_factor",
2059
    "_limit",
2060
    "_next",
2061
    "_start",
2062
    ]
2063

    
2064
  def __init__(self, start, factor, limit):
2065
    """Initializes this class.
2066

2067
    @type start: float
2068
    @param start: Initial delay
2069
    @type factor: float
2070
    @param factor: Factor for delay increase
2071
    @type limit: float or None
2072
    @param limit: Upper limit for delay or None for no limit
2073

2074
    """
2075
    assert start > 0.0
2076
    assert factor >= 1.0
2077
    assert limit is None or limit >= 0.0
2078

    
2079
    self._start = start
2080
    self._factor = factor
2081
    self._limit = limit
2082

    
2083
    self._next = start
2084

    
2085
  def __call__(self):
2086
    """Returns current delay and calculates the next one.
2087

2088
    """
2089
    current = self._next
2090

    
2091
    # Update for next run
2092
    if self._limit is None or self._next < self._limit:
2093
      self._next = max(self._limit, self._next * self._factor)
2094

    
2095
    return current
2096

    
2097

    
2098
#: Special delay to specify whole remaining timeout
2099
RETRY_REMAINING_TIME = object()
2100

    
2101

    
2102
def Retry(fn, delay, timeout, args=None, wait_fn=time.sleep,
2103
          _time_fn=time.time):
2104
  """Call a function repeatedly until it succeeds.
2105

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

2110
  C{delay} can be one of the following:
2111
    - callable returning the delay length as a float
2112
    - Tuple of (start, factor, limit)
2113
    - L{RETRY_REMAINING_TIME} to sleep until the timeout expires (this is
2114
      useful when overriding L{wait_fn} to wait for an external event)
2115
    - A static delay as a number (int or float)
2116

2117
  @type fn: callable
2118
  @param fn: Function to be called
2119
  @param delay: Either a callable (returning the delay), a tuple of (start,
2120
                factor, limit) (see L{_RetryDelayCalculator}),
2121
                L{RETRY_REMAINING_TIME} or a number (int or float)
2122
  @type timeout: float
2123
  @param timeout: Total timeout
2124
  @type wait_fn: callable
2125
  @param wait_fn: Waiting function
2126
  @return: Return value of function
2127

2128
  """
2129
  assert callable(fn)
2130
  assert callable(wait_fn)
2131
  assert callable(_time_fn)
2132

    
2133
  if args is None:
2134
    args = []
2135

    
2136
  end_time = _time_fn() + timeout
2137

    
2138
  if callable(delay):
2139
    # External function to calculate delay
2140
    calc_delay = delay
2141

    
2142
  elif isinstance(delay, (tuple, list)):
2143
    # Increasing delay with optional upper boundary
2144
    (start, factor, limit) = delay
2145
    calc_delay = _RetryDelayCalculator(start, factor, limit)
2146

    
2147
  elif delay is RETRY_REMAINING_TIME:
2148
    # Always use the remaining time
2149
    calc_delay = None
2150

    
2151
  else:
2152
    # Static delay
2153
    calc_delay = lambda: delay
2154

    
2155
  assert calc_delay is None or callable(calc_delay)
2156

    
2157
  while True:
2158
    try:
2159
      return fn(*args)
2160
    except RetryAgain:
2161
      pass
2162

    
2163
    remaining_time = end_time - _time_fn()
2164

    
2165
    if remaining_time < 0.0:
2166
      raise RetryTimeout()
2167

    
2168
    assert remaining_time >= 0.0
2169

    
2170
    if calc_delay is None:
2171
      wait_fn(remaining_time)
2172
    else:
2173
      current_delay = calc_delay()
2174
      if current_delay > 0.0:
2175
        wait_fn(current_delay)
2176

    
2177

    
2178
class FileLock(object):
2179
  """Utility class for file locks.
2180

2181
  """
2182
  def __init__(self, filename):
2183
    """Constructor for FileLock.
2184

2185
    This will open the file denoted by the I{filename} argument.
2186

2187
    @type filename: str
2188
    @param filename: path to the file to be locked
2189

2190
    """
2191
    self.filename = filename
2192
    self.fd = open(self.filename, "w")
2193

    
2194
  def __del__(self):
2195
    self.Close()
2196

    
2197
  def Close(self):
2198
    """Close the file and release the lock.
2199

2200
    """
2201
    if self.fd:
2202
      self.fd.close()
2203
      self.fd = None
2204

    
2205
  def _flock(self, flag, blocking, timeout, errmsg):
2206
    """Wrapper for fcntl.flock.
2207

2208
    @type flag: int
2209
    @param flag: operation flag
2210
    @type blocking: bool
2211
    @param blocking: whether the operation should be done in blocking mode.
2212
    @type timeout: None or float
2213
    @param timeout: for how long the operation should be retried (implies
2214
                    non-blocking mode).
2215
    @type errmsg: string
2216
    @param errmsg: error message in case operation fails.
2217

2218
    """
2219
    assert self.fd, "Lock was closed"
2220
    assert timeout is None or timeout >= 0, \
2221
      "If specified, timeout must be positive"
2222

    
2223
    if timeout is not None:
2224
      flag |= fcntl.LOCK_NB
2225
      timeout_end = time.time() + timeout
2226

    
2227
    # Blocking doesn't have effect with timeout
2228
    elif not blocking:
2229
      flag |= fcntl.LOCK_NB
2230
      timeout_end = None
2231

    
2232
    # TODO: Convert to utils.Retry
2233

    
2234
    retry = True
2235
    while retry:
2236
      try:
2237
        fcntl.flock(self.fd, flag)
2238
        retry = False
2239
      except IOError, err:
2240
        if err.errno in (errno.EAGAIN, ):
2241
          if timeout_end is not None and time.time() < timeout_end:
2242
            # Wait before trying again
2243
            time.sleep(max(0.1, min(1.0, timeout)))
2244
          else:
2245
            raise errors.LockError(errmsg)
2246
        else:
2247
          logging.exception("fcntl.flock failed")
2248
          raise
2249

    
2250
  def Exclusive(self, blocking=False, timeout=None):
2251
    """Locks the file in exclusive mode.
2252

2253
    @type blocking: boolean
2254
    @param blocking: whether to block and wait until we
2255
        can lock the file or return immediately
2256
    @type timeout: int or None
2257
    @param timeout: if not None, the duration to wait for the lock
2258
        (in blocking mode)
2259

2260
    """
2261
    self._flock(fcntl.LOCK_EX, blocking, timeout,
2262
                "Failed to lock %s in exclusive mode" % self.filename)
2263

    
2264
  def Shared(self, blocking=False, timeout=None):
2265
    """Locks the file in shared mode.
2266

2267
    @type blocking: boolean
2268
    @param blocking: whether to block and wait until we
2269
        can lock the file or return immediately
2270
    @type timeout: int or None
2271
    @param timeout: if not None, the duration to wait for the lock
2272
        (in blocking mode)
2273

2274
    """
2275
    self._flock(fcntl.LOCK_SH, blocking, timeout,
2276
                "Failed to lock %s in shared mode" % self.filename)
2277

    
2278
  def Unlock(self, blocking=True, timeout=None):
2279
    """Unlocks the file.
2280

2281
    According to C{flock(2)}, unlocking can also be a nonblocking
2282
    operation::
2283

2284
      To make a non-blocking request, include LOCK_NB with any of the above
2285
      operations.
2286

2287
    @type blocking: boolean
2288
    @param blocking: whether to block and wait until we
2289
        can lock the file or return immediately
2290
    @type timeout: int or None
2291
    @param timeout: if not None, the duration to wait for the lock
2292
        (in blocking mode)
2293

2294
    """
2295
    self._flock(fcntl.LOCK_UN, blocking, timeout,
2296
                "Failed to unlock %s" % self.filename)
2297

    
2298

    
2299
def SignalHandled(signums):
2300
  """Signal Handled decoration.
2301

2302
  This special decorator installs a signal handler and then calls the target
2303
  function. The function must accept a 'signal_handlers' keyword argument,
2304
  which will contain a dict indexed by signal number, with SignalHandler
2305
  objects as values.
2306

2307
  The decorator can be safely stacked with iself, to handle multiple signals
2308
  with different handlers.
2309

2310
  @type signums: list
2311
  @param signums: signals to intercept
2312

2313
  """
2314
  def wrap(fn):
2315
    def sig_function(*args, **kwargs):
2316
      assert 'signal_handlers' not in kwargs or \
2317
             kwargs['signal_handlers'] is None or \
2318
             isinstance(kwargs['signal_handlers'], dict), \
2319
             "Wrong signal_handlers parameter in original function call"
2320
      if 'signal_handlers' in kwargs and kwargs['signal_handlers'] is not None:
2321
        signal_handlers = kwargs['signal_handlers']
2322
      else:
2323
        signal_handlers = {}
2324
        kwargs['signal_handlers'] = signal_handlers
2325
      sighandler = SignalHandler(signums)
2326
      try:
2327
        for sig in signums:
2328
          signal_handlers[sig] = sighandler
2329
        return fn(*args, **kwargs)
2330
      finally:
2331
        sighandler.Reset()
2332
    return sig_function
2333
  return wrap
2334

    
2335

    
2336
class SignalHandler(object):
2337
  """Generic signal handler class.
2338

2339
  It automatically restores the original handler when deconstructed or
2340
  when L{Reset} is called. You can either pass your own handler
2341
  function in or query the L{called} attribute to detect whether the
2342
  signal was sent.
2343

2344
  @type signum: list
2345
  @ivar signum: the signals we handle
2346
  @type called: boolean
2347
  @ivar called: tracks whether any of the signals have been raised
2348

2349
  """
2350
  def __init__(self, signum):
2351
    """Constructs a new SignalHandler instance.
2352

2353
    @type signum: int or list of ints
2354
    @param signum: Single signal number or set of signal numbers
2355

2356
    """
2357
    self.signum = set(signum)
2358
    self.called = False
2359

    
2360
    self._previous = {}
2361
    try:
2362
      for signum in self.signum:
2363
        # Setup handler
2364
        prev_handler = signal.signal(signum, self._HandleSignal)
2365
        try:
2366
          self._previous[signum] = prev_handler
2367
        except:
2368
          # Restore previous handler
2369
          signal.signal(signum, prev_handler)
2370
          raise
2371
    except:
2372
      # Reset all handlers
2373
      self.Reset()
2374
      # Here we have a race condition: a handler may have already been called,
2375
      # but there's not much we can do about it at this point.
2376
      raise
2377

    
2378
  def __del__(self):
2379
    self.Reset()
2380

    
2381
  def Reset(self):
2382
    """Restore previous handler.
2383

2384
    This will reset all the signals to their previous handlers.
2385

2386
    """
2387
    for signum, prev_handler in self._previous.items():
2388
      signal.signal(signum, prev_handler)
2389
      # If successful, remove from dict
2390
      del self._previous[signum]
2391

    
2392
  def Clear(self):
2393
    """Unsets the L{called} flag.
2394

2395
    This function can be used in case a signal may arrive several times.
2396

2397
    """
2398
    self.called = False
2399

    
2400
  def _HandleSignal(self, signum, frame):
2401
    """Actual signal handling function.
2402

2403
    """
2404
    # This is not nice and not absolutely atomic, but it appears to be the only
2405
    # solution in Python -- there are no atomic types.
2406
    self.called = True
2407

    
2408

    
2409
class FieldSet(object):
2410
  """A simple field set.
2411

2412
  Among the features are:
2413
    - checking if a string is among a list of static string or regex objects
2414
    - checking if a whole list of string matches
2415
    - returning the matching groups from a regex match
2416

2417
  Internally, all fields are held as regular expression objects.
2418

2419
  """
2420
  def __init__(self, *items):
2421
    self.items = [re.compile("^%s$" % value) for value in items]
2422

    
2423
  def Extend(self, other_set):
2424
    """Extend the field set with the items from another one"""
2425
    self.items.extend(other_set.items)
2426

    
2427
  def Matches(self, field):
2428
    """Checks if a field matches the current set
2429

2430
    @type field: str
2431
    @param field: the string to match
2432
    @return: either None or a regular expression match object
2433

2434
    """
2435
    for m in itertools.ifilter(None, (val.match(field) for val in self.items)):
2436
      return m
2437
    return None
2438

    
2439
  def NonMatching(self, items):
2440
    """Returns the list of fields not matching the current set
2441

2442
    @type items: list
2443
    @param items: the list of fields to check
2444
    @rtype: list
2445
    @return: list of non-matching fields
2446

2447
    """
2448
    return [val for val in items if not self.Matches(val)]