Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ ab3e6da8

History | View | Annotate | Download (64.1 kB)

1
#
2
#
3

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

    
21

    
22
"""Ganeti 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=None):
1230
  """Reads a file.
1231

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

1237
  """
1238
  f = open(file_name, "r")
1239
  try:
1240
    if size is None:
1241
      return f.read()
1242
    else:
1243
      return f.read(size)
1244
  finally:
1245
    f.close()
1246

    
1247

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

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

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

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

1290
  @rtype: None or int
1291
  @return: None if the 'close' parameter evaluates to True,
1292
      otherwise the file descriptor
1293

1294
  @raise errors.ProgrammerError: if any of the arguments are not valid
1295

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

    
1301
  if [fn, data].count(None) != 1:
1302
    raise errors.ProgrammerError("fn or data required")
1303

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

    
1308
  if backup and not dry_run and os.path.isfile(file_name):
1309
    CreateBackup(file_name)
1310

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

    
1344
  return result
1345

    
1346

    
1347
def FirstFree(seq, base=0):
1348
  """Returns the first non-existing integer from seq.
1349

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

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

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

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

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

    
1374

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

    
1381

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

    
1388

    
1389
def UniqueSequence(seq):
1390
  """Returns a list with unique elements.
1391

1392
  Element order is preserved.
1393

1394
  @type seq: sequence
1395
  @param seq: the sequence with the source elements
1396
  @rtype: list
1397
  @return: list of unique elements from seq
1398

1399
  """
1400
  seen = set()
1401
  return [i for i in seq if i not in seen and not seen.add(i)]
1402

    
1403

    
1404
def IsValidMac(mac):
1405
  """Predicate to check if a MAC address is valid.
1406

1407
  Checks whether the supplied MAC address is formally correct, only
1408
  accepts colon separated format.
1409

1410
  @type mac: str
1411
  @param mac: the MAC to be validated
1412
  @rtype: boolean
1413
  @return: True is the MAC seems valid
1414

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

    
1419

    
1420
def TestDelay(duration):
1421
  """Sleep for a fixed amount of time.
1422

1423
  @type duration: float
1424
  @param duration: the sleep duration
1425
  @rtype: boolean
1426
  @return: False for negative value, True otherwise
1427

1428
  """
1429
  if duration < 0:
1430
    return False, "Invalid sleep duration"
1431
  time.sleep(duration)
1432
  return True, None
1433

    
1434

    
1435
def _CloseFDNoErr(fd, retries=5):
1436
  """Close a file descriptor ignoring errors.
1437

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

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

    
1454

    
1455
def CloseFDs(noclose_fds=None):
1456
  """Close file descriptors.
1457

1458
  This closes all file descriptors above 2 (i.e. except
1459
  stdin/out/err).
1460

1461
  @type noclose_fds: list or None
1462
  @param noclose_fds: if given, it denotes a list of file descriptor
1463
      that should not be closed
1464

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

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

    
1486

    
1487
def Daemonize(logfile):
1488
  """Daemonize the current process.
1489

1490
  This detaches the current process from the controlling terminal and
1491
  runs it in the background as a daemon.
1492

1493
  @type logfile: str
1494
  @param logfile: the logfile to which we should redirect stdout/stderr
1495
  @rtype: int
1496
  @return: the value zero
1497

1498
  """
1499
  UMASK = 077
1500
  WORKDIR = "/"
1501

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

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

    
1527

    
1528
def DaemonPidFileName(name):
1529
  """Compute a ganeti pid file absolute path
1530

1531
  @type name: str
1532
  @param name: the daemon name
1533
  @rtype: str
1534
  @return: the full path to the pidfile corresponding to the given
1535
      daemon name
1536

1537
  """
1538
  return os.path.join(constants.RUN_GANETI_DIR, "%s.pid" % name)
1539

    
1540

    
1541
def WritePidFile(name):
1542
  """Write the current process pidfile.
1543

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

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

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

    
1557
  WriteFile(pidfilename, data="%d\n" % pid)
1558

    
1559

    
1560
def RemovePidFile(name):
1561
  """Remove the current process pidfile.
1562

1563
  Any errors are ignored.
1564

1565
  @type name: str
1566
  @param name: the daemon name used to derive the pidfile name
1567

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

    
1576

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

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

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

    
1604
  if pid <= 0:
1605
    # kill with pid=0 == suicide
1606
    raise errors.ProgrammerError("Invalid pid given '%s'" % pid)
1607

    
1608
  if not IsProcessAlive(pid):
1609
    return
1610

    
1611
  _helper(pid, signal_, waitpid)
1612

    
1613
  if timeout <= 0:
1614
    return
1615

    
1616
  def _CheckProcess():
1617
    if not IsProcessAlive(pid):
1618
      return
1619

    
1620
    try:
1621
      (result_pid, _) = os.waitpid(pid, os.WNOHANG)
1622
    except OSError:
1623
      raise RetryAgain()
1624

    
1625
    if result_pid > 0:
1626
      return
1627

    
1628
    raise RetryAgain()
1629

    
1630
  try:
1631
    # Wait up to $timeout seconds
1632
    Retry(_CheckProcess, (0.01, 1.5, 0.1), timeout)
1633
  except RetryTimeout:
1634
    pass
1635

    
1636
  if IsProcessAlive(pid):
1637
    # Kill process if it's still alive
1638
    _helper(pid, signal.SIGKILL, waitpid)
1639

    
1640

    
1641
def FindFile(name, search_path, test=os.path.exists):
1642
  """Look for a filesystem object in a given path.
1643

1644
  This is an abstract method to search for filesystem object (files,
1645
  dirs) under a given search path.
1646

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

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

    
1665

    
1666
def CheckVolumeGroupSize(vglist, vgname, minsize):
1667
  """Checks if the volume group list is valid.
1668

1669
  The function will check if a given volume group is in the list of
1670
  volume groups and has a minimum size.
1671

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

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

    
1690

    
1691
def SplitTime(value):
1692
  """Splits time as floating point number into a tuple.
1693

1694
  @param value: Time in seconds
1695
  @type value: int or float
1696
  @return: Tuple containing (seconds, microseconds)
1697

1698
  """
1699
  (seconds, microseconds) = divmod(int(value * 1000000), 1000000)
1700

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

    
1706
  return (int(seconds), int(microseconds))
1707

    
1708

    
1709
def MergeTime(timetuple):
1710
  """Merges a tuple into time as a floating point number.
1711

1712
  @param timetuple: Time as tuple, (seconds, microseconds)
1713
  @type timetuple: tuple
1714
  @return: Time as a floating point number expressed in seconds
1715

1716
  """
1717
  (seconds, microseconds) = timetuple
1718

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

    
1724
  return float(seconds) + (float(microseconds) * 0.000001)
1725

    
1726

    
1727
def GetDaemonPort(daemon_name):
1728
  """Get the daemon port for this cluster.
1729

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

1734
  @type daemon_name: string
1735
  @param daemon_name: daemon name (in constants.DAEMONS_PORTS)
1736
  @rtype: int
1737

1738
  """
1739
  if daemon_name not in constants.DAEMONS_PORTS:
1740
    raise errors.ProgrammerError("Unknown daemon: %s" % daemon_name)
1741

    
1742
  (proto, default_port) = constants.DAEMONS_PORTS[daemon_name]
1743
  try:
1744
    port = socket.getservbyname(daemon_name, proto)
1745
  except socket.error:
1746
    port = default_port
1747

    
1748
  return port
1749

    
1750

    
1751
def SetupLogging(logfile, debug=False, stderr_logging=False, program="",
1752
                 multithreaded=False):
1753
  """Configures the logging module.
1754

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

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

    
1778
  root_logger = logging.getLogger("")
1779
  root_logger.setLevel(logging.NOTSET)
1780

    
1781
  # Remove all previously setup handlers
1782
  for handler in root_logger.handlers:
1783
    handler.close()
1784
    root_logger.removeHandler(handler)
1785

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

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

    
1814

    
1815
def IsNormAbsPath(path):
1816
  """Check whether a path is absolute and also normalized
1817

1818
  This avoids things like /dir/../../other/path to be valid.
1819

1820
  """
1821
  return os.path.normpath(path) == path and os.path.isabs(path)
1822

    
1823

    
1824
def TailFile(fname, lines=20):
1825
  """Return the last lines from a file.
1826

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

1831
  @param fname: the file name
1832
  @type lines: int
1833
  @param lines: the (maximum) number of lines to return
1834

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

    
1846
  rows = raw_data.splitlines()
1847
  return rows[-lines:]
1848

    
1849

    
1850
def SafeEncode(text):
1851
  """Return a 'safe' version of a source string.
1852

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

1862
  @type text: str or unicode
1863
  @param text: input data
1864
  @rtype: str
1865
  @return: a safe version of text
1866

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

    
1886

    
1887
def CommaJoin(names):
1888
  """Nicely join a set of identifiers.
1889

1890
  @param names: set, list or tuple
1891
  @return: a string with the formatted results
1892

1893
  """
1894
  return ", ".join(["'%s'" % val for val in names])
1895

    
1896

    
1897
def BytesToMebibyte(value):
1898
  """Converts bytes to mebibytes.
1899

1900
  @type value: int
1901
  @param value: Value in bytes
1902
  @rtype: int
1903
  @return: Value in mebibytes
1904

1905
  """
1906
  return int(round(value / (1024.0 * 1024.0), 0))
1907

    
1908

    
1909
def CalculateDirectorySize(path):
1910
  """Calculates the size of a directory recursively.
1911

1912
  @type path: string
1913
  @param path: Path to directory
1914
  @rtype: int
1915
  @return: Size in mebibytes
1916

1917
  """
1918
  size = 0
1919

    
1920
  for (curpath, _, files) in os.walk(path):
1921
    for filename in files:
1922
      st = os.lstat(os.path.join(curpath, filename))
1923
      size += st.st_size
1924

    
1925
  return BytesToMebibyte(size)
1926

    
1927

    
1928
def GetFilesystemStats(path):
1929
  """Returns the total and free space on a filesystem.
1930

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

1936
  """
1937
  st = os.statvfs(path)
1938

    
1939
  fsize = BytesToMebibyte(st.f_bavail * st.f_frsize)
1940
  tsize = BytesToMebibyte(st.f_blocks * st.f_frsize)
1941
  return (tsize, fsize)
1942

    
1943

    
1944
def LockedMethod(fn):
1945
  """Synchronized object access decorator.
1946

1947
  This decorator is intended to protect access to an object using the
1948
  object's own lock which is hardcoded to '_lock'.
1949

1950
  """
1951
  def _LockDebug(*args, **kwargs):
1952
    if debug_locks:
1953
      logging.debug(*args, **kwargs)
1954

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

    
1970

    
1971
def LockFile(fd):
1972
  """Locks a file using POSIX locks.
1973

1974
  @type fd: int
1975
  @param fd: the file descriptor we need to lock
1976

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

    
1985

    
1986
def FormatTime(val):
1987
  """Formats a time value.
1988

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

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

    
2000

    
2001
def ReadWatcherPauseFile(filename, now=None, remove_after=3600):
2002
  """Reads the watcher pause file.
2003

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

2012
  """
2013
  if now is None:
2014
    now = time.time()
2015

    
2016
  try:
2017
    value = ReadFile(filename)
2018
  except IOError, err:
2019
    if err.errno != errno.ENOENT:
2020
      raise
2021
    value = None
2022

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

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

    
2038
      elif now > value:
2039
        value = None
2040

    
2041
  return value
2042

    
2043

    
2044
class RetryTimeout(Exception):
2045
  """Retry loop timed out.
2046

2047
  """
2048

    
2049

    
2050
class RetryAgain(Exception):
2051
  """Retry again.
2052

2053
  """
2054

    
2055

    
2056
class _RetryDelayCalculator(object):
2057
  """Calculator for increasing delays.
2058

2059
  """
2060
  __slots__ = [
2061
    "_factor",
2062
    "_limit",
2063
    "_next",
2064
    "_start",
2065
    ]
2066

    
2067
  def __init__(self, start, factor, limit):
2068
    """Initializes this class.
2069

2070
    @type start: float
2071
    @param start: Initial delay
2072
    @type factor: float
2073
    @param factor: Factor for delay increase
2074
    @type limit: float or None
2075
    @param limit: Upper limit for delay or None for no limit
2076

2077
    """
2078
    assert start > 0.0
2079
    assert factor >= 1.0
2080
    assert limit is None or limit >= 0.0
2081

    
2082
    self._start = start
2083
    self._factor = factor
2084
    self._limit = limit
2085

    
2086
    self._next = start
2087

    
2088
  def __call__(self):
2089
    """Returns current delay and calculates the next one.
2090

2091
    """
2092
    current = self._next
2093

    
2094
    # Update for next run
2095
    if self._limit is None or self._next < self._limit:
2096
      self._next = max(self._limit, self._next * self._factor)
2097

    
2098
    return current
2099

    
2100

    
2101
#: Special delay to specify whole remaining timeout
2102
RETRY_REMAINING_TIME = object()
2103

    
2104

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

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

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

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

2131
  """
2132
  assert callable(fn)
2133
  assert callable(wait_fn)
2134
  assert callable(_time_fn)
2135

    
2136
  if args is None:
2137
    args = []
2138

    
2139
  end_time = _time_fn() + timeout
2140

    
2141
  if callable(delay):
2142
    # External function to calculate delay
2143
    calc_delay = delay
2144

    
2145
  elif isinstance(delay, (tuple, list)):
2146
    # Increasing delay with optional upper boundary
2147
    (start, factor, limit) = delay
2148
    calc_delay = _RetryDelayCalculator(start, factor, limit)
2149

    
2150
  elif delay is RETRY_REMAINING_TIME:
2151
    # Always use the remaining time
2152
    calc_delay = None
2153

    
2154
  else:
2155
    # Static delay
2156
    calc_delay = lambda: delay
2157

    
2158
  assert calc_delay is None or callable(calc_delay)
2159

    
2160
  while True:
2161
    try:
2162
      return fn(*args)
2163
    except RetryAgain:
2164
      pass
2165

    
2166
    remaining_time = end_time - _time_fn()
2167

    
2168
    if remaining_time < 0.0:
2169
      raise RetryTimeout()
2170

    
2171
    assert remaining_time >= 0.0
2172

    
2173
    if calc_delay is None:
2174
      wait_fn(remaining_time)
2175
    else:
2176
      current_delay = calc_delay()
2177
      if current_delay > 0.0:
2178
        wait_fn(current_delay)
2179

    
2180

    
2181
class FileLock(object):
2182
  """Utility class for file locks.
2183

2184
  """
2185
  def __init__(self, filename):
2186
    """Constructor for FileLock.
2187

2188
    This will open the file denoted by the I{filename} argument.
2189

2190
    @type filename: str
2191
    @param filename: path to the file to be locked
2192

2193
    """
2194
    self.filename = filename
2195
    self.fd = open(self.filename, "w")
2196

    
2197
  def __del__(self):
2198
    self.Close()
2199

    
2200
  def Close(self):
2201
    """Close the file and release the lock.
2202

2203
    """
2204
    if self.fd:
2205
      self.fd.close()
2206
      self.fd = None
2207

    
2208
  def _flock(self, flag, blocking, timeout, errmsg):
2209
    """Wrapper for fcntl.flock.
2210

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

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

    
2226
    if timeout is not None:
2227
      flag |= fcntl.LOCK_NB
2228
      timeout_end = time.time() + timeout
2229

    
2230
    # Blocking doesn't have effect with timeout
2231
    elif not blocking:
2232
      flag |= fcntl.LOCK_NB
2233
      timeout_end = None
2234

    
2235
    # TODO: Convert to utils.Retry
2236

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

    
2253
  def Exclusive(self, blocking=False, timeout=None):
2254
    """Locks the file in exclusive mode.
2255

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

2263
    """
2264
    self._flock(fcntl.LOCK_EX, blocking, timeout,
2265
                "Failed to lock %s in exclusive mode" % self.filename)
2266

    
2267
  def Shared(self, blocking=False, timeout=None):
2268
    """Locks the file in shared mode.
2269

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

2277
    """
2278
    self._flock(fcntl.LOCK_SH, blocking, timeout,
2279
                "Failed to lock %s in shared mode" % self.filename)
2280

    
2281
  def Unlock(self, blocking=True, timeout=None):
2282
    """Unlocks the file.
2283

2284
    According to C{flock(2)}, unlocking can also be a nonblocking
2285
    operation::
2286

2287
      To make a non-blocking request, include LOCK_NB with any of the above
2288
      operations.
2289

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

2297
    """
2298
    self._flock(fcntl.LOCK_UN, blocking, timeout,
2299
                "Failed to unlock %s" % self.filename)
2300

    
2301

    
2302
def SignalHandled(signums):
2303
  """Signal Handled decoration.
2304

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

2310
  The decorator can be safely stacked with iself, to handle multiple signals
2311
  with different handlers.
2312

2313
  @type signums: list
2314
  @param signums: signals to intercept
2315

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

    
2338

    
2339
class SignalHandler(object):
2340
  """Generic signal handler class.
2341

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

2347
  @type signum: list
2348
  @ivar signum: the signals we handle
2349
  @type called: boolean
2350
  @ivar called: tracks whether any of the signals have been raised
2351

2352
  """
2353
  def __init__(self, signum):
2354
    """Constructs a new SignalHandler instance.
2355

2356
    @type signum: int or list of ints
2357
    @param signum: Single signal number or set of signal numbers
2358

2359
    """
2360
    self.signum = set(signum)
2361
    self.called = False
2362

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

    
2381
  def __del__(self):
2382
    self.Reset()
2383

    
2384
  def Reset(self):
2385
    """Restore previous handler.
2386

2387
    This will reset all the signals to their previous handlers.
2388

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

    
2395
  def Clear(self):
2396
    """Unsets the L{called} flag.
2397

2398
    This function can be used in case a signal may arrive several times.
2399

2400
    """
2401
    self.called = False
2402

    
2403
  def _HandleSignal(self, signum, frame):
2404
    """Actual signal handling function.
2405

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

    
2411

    
2412
class FieldSet(object):
2413
  """A simple field set.
2414

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

2420
  Internally, all fields are held as regular expression objects.
2421

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

    
2426
  def Extend(self, other_set):
2427
    """Extend the field set with the items from another one"""
2428
    self.items.extend(other_set.items)
2429

    
2430
  def Matches(self, field):
2431
    """Checks if a field matches the current set
2432

2433
    @type field: str
2434
    @param field: the string to match
2435
    @return: either None or a regular expression match object
2436

2437
    """
2438
    for m in itertools.ifilter(None, (val.match(field) for val in self.items)):
2439
      return m
2440
    return None
2441

    
2442
  def NonMatching(self, items):
2443
    """Returns the list of fields not matching the current set
2444

2445
    @type items: list
2446
    @param items: the list of fields to check
2447
    @rtype: list
2448
    @return: list of non-matching fields
2449

2450
    """
2451
    return [val for val in items if not self.Matches(val)]