Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ 82187135

History | View | Annotate | Download (64.9 kB)

1
#
2
#
3

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

    
21

    
22
"""Ganeti 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): # pylint: disable-msg=W0622
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): # pylint: disable-msg=W0622
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 NormalizeAndValidateMac(mac):
1402
  """Normalizes and check if a MAC address is valid.
1403

1404
  Checks whether the supplied MAC address is formally correct, only
1405
  accepts colon separated format. Normalize it to all lower.
1406

1407
  @type mac: str
1408
  @param mac: the MAC to be validated
1409
  @rtype: str
1410
  @return: returns the normalized and validated MAC.
1411

1412
  @raise errors.OpPrereqError: If the MAC isn't valid
1413

1414
  """
1415
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$", re.I)
1416
  if not mac_check.match(mac):
1417
    raise errors.OpPrereqError("Invalid MAC address specified: %s" %
1418
                               mac, errors.ECODE_INVAL)
1419

    
1420
  return mac.lower()
1421

    
1422

    
1423
def TestDelay(duration):
1424
  """Sleep for a fixed amount of time.
1425

1426
  @type duration: float
1427
  @param duration: the sleep duration
1428
  @rtype: boolean
1429
  @return: False for negative value, True otherwise
1430

1431
  """
1432
  if duration < 0:
1433
    return False, "Invalid sleep duration"
1434
  time.sleep(duration)
1435
  return True, None
1436

    
1437

    
1438
def _CloseFDNoErr(fd, retries=5):
1439
  """Close a file descriptor ignoring errors.
1440

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

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

    
1457

    
1458
def CloseFDs(noclose_fds=None):
1459
  """Close file descriptors.
1460

1461
  This closes all file descriptors above 2 (i.e. except
1462
  stdin/out/err).
1463

1464
  @type noclose_fds: list or None
1465
  @param noclose_fds: if given, it denotes a list of file descriptor
1466
      that should not be closed
1467

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

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

    
1489

    
1490
def Daemonize(logfile):
1491
  """Daemonize the current process.
1492

1493
  This detaches the current process from the controlling terminal and
1494
  runs it in the background as a daemon.
1495

1496
  @type logfile: str
1497
  @param logfile: the logfile to which we should redirect stdout/stderr
1498
  @rtype: int
1499
  @return: the value zero
1500

1501
  """
1502
  # pylint: disable-msg=W0212
1503
  # yes, we really want os._exit
1504
  UMASK = 077
1505
  WORKDIR = "/"
1506

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

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

    
1532

    
1533
def DaemonPidFileName(name):
1534
  """Compute a ganeti pid file absolute path
1535

1536
  @type name: str
1537
  @param name: the daemon name
1538
  @rtype: str
1539
  @return: the full path to the pidfile corresponding to the given
1540
      daemon name
1541

1542
  """
1543
  return os.path.join(constants.RUN_GANETI_DIR, "%s.pid" % name)
1544

    
1545

    
1546
def WritePidFile(name):
1547
  """Write the current process pidfile.
1548

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

1551
  @type name: str
1552
  @param name: the daemon name to use
1553
  @raise errors.GenericError: if the pid file already exists and
1554
      points to a live process
1555

1556
  """
1557
  pid = os.getpid()
1558
  pidfilename = DaemonPidFileName(name)
1559
  if IsProcessAlive(ReadPidFile(pidfilename)):
1560
    raise errors.GenericError("%s contains a live process" % pidfilename)
1561

    
1562
  WriteFile(pidfilename, data="%d\n" % pid)
1563

    
1564

    
1565
def RemovePidFile(name):
1566
  """Remove the current process pidfile.
1567

1568
  Any errors are ignored.
1569

1570
  @type name: str
1571
  @param name: the daemon name used to derive the pidfile name
1572

1573
  """
1574
  pidfilename = DaemonPidFileName(name)
1575
  # TODO: we could check here that the file contains our pid
1576
  try:
1577
    RemoveFile(pidfilename)
1578
  except: # pylint: disable-msg=W0702
1579
    pass
1580

    
1581

    
1582
def KillProcess(pid, signal_=signal.SIGTERM, timeout=30,
1583
                waitpid=False):
1584
  """Kill a process given by its pid.
1585

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

1599
  """
1600
  def _helper(pid, signal_, wait):
1601
    """Simple helper to encapsulate the kill/waitpid sequence"""
1602
    os.kill(pid, signal_)
1603
    if wait:
1604
      try:
1605
        os.waitpid(pid, os.WNOHANG)
1606
      except OSError:
1607
        pass
1608

    
1609
  if pid <= 0:
1610
    # kill with pid=0 == suicide
1611
    raise errors.ProgrammerError("Invalid pid given '%s'" % pid)
1612

    
1613
  if not IsProcessAlive(pid):
1614
    return
1615

    
1616
  _helper(pid, signal_, waitpid)
1617

    
1618
  if timeout <= 0:
1619
    return
1620

    
1621
  def _CheckProcess():
1622
    if not IsProcessAlive(pid):
1623
      return
1624

    
1625
    try:
1626
      (result_pid, _) = os.waitpid(pid, os.WNOHANG)
1627
    except OSError:
1628
      raise RetryAgain()
1629

    
1630
    if result_pid > 0:
1631
      return
1632

    
1633
    raise RetryAgain()
1634

    
1635
  try:
1636
    # Wait up to $timeout seconds
1637
    Retry(_CheckProcess, (0.01, 1.5, 0.1), timeout)
1638
  except RetryTimeout:
1639
    pass
1640

    
1641
  if IsProcessAlive(pid):
1642
    # Kill process if it's still alive
1643
    _helper(pid, signal.SIGKILL, waitpid)
1644

    
1645

    
1646
def FindFile(name, search_path, test=os.path.exists):
1647
  """Look for a filesystem object in a given path.
1648

1649
  This is an abstract method to search for filesystem object (files,
1650
  dirs) under a given search path.
1651

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

1663
  """
1664
  # validate the filename mask
1665
  if constants.EXT_PLUGIN_MASK.match(name) is None:
1666
    logging.critical("Invalid value passed for external script name: '%s'",
1667
                     name)
1668
    return None
1669

    
1670
  for dir_name in search_path:
1671
    item_name = os.path.sep.join([dir_name, name])
1672
    # check the user test and that we're indeed resolving to the given
1673
    # basename
1674
    if test(item_name) and os.path.basename(item_name) == name:
1675
      return item_name
1676
  return None
1677

    
1678

    
1679
def CheckVolumeGroupSize(vglist, vgname, minsize):
1680
  """Checks if the volume group list is valid.
1681

1682
  The function will check if a given volume group is in the list of
1683
  volume groups and has a minimum size.
1684

1685
  @type vglist: dict
1686
  @param vglist: dictionary of volume group names and their size
1687
  @type vgname: str
1688
  @param vgname: the volume group we should check
1689
  @type minsize: int
1690
  @param minsize: the minimum size we accept
1691
  @rtype: None or str
1692
  @return: None for success, otherwise the error message
1693

1694
  """
1695
  vgsize = vglist.get(vgname, None)
1696
  if vgsize is None:
1697
    return "volume group '%s' missing" % vgname
1698
  elif vgsize < minsize:
1699
    return ("volume group '%s' too small (%s MiB required, %d MiB found)" %
1700
            (vgname, minsize, vgsize))
1701
  return None
1702

    
1703

    
1704
def SplitTime(value):
1705
  """Splits time as floating point number into a tuple.
1706

1707
  @param value: Time in seconds
1708
  @type value: int or float
1709
  @return: Tuple containing (seconds, microseconds)
1710

1711
  """
1712
  (seconds, microseconds) = divmod(int(value * 1000000), 1000000)
1713

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

    
1719
  return (int(seconds), int(microseconds))
1720

    
1721

    
1722
def MergeTime(timetuple):
1723
  """Merges a tuple into time as a floating point number.
1724

1725
  @param timetuple: Time as tuple, (seconds, microseconds)
1726
  @type timetuple: tuple
1727
  @return: Time as a floating point number expressed in seconds
1728

1729
  """
1730
  (seconds, microseconds) = timetuple
1731

    
1732
  assert 0 <= seconds, \
1733
    "Seconds must be larger than or equal to 0, but are %s" % seconds
1734
  assert 0 <= microseconds <= 999999, \
1735
    "Microseconds must be 0-999999, but are %s" % microseconds
1736

    
1737
  return float(seconds) + (float(microseconds) * 0.000001)
1738

    
1739

    
1740
def GetDaemonPort(daemon_name):
1741
  """Get the daemon port for this cluster.
1742

1743
  Note that this routine does not read a ganeti-specific file, but
1744
  instead uses C{socket.getservbyname} to allow pre-customization of
1745
  this parameter outside of Ganeti.
1746

1747
  @type daemon_name: string
1748
  @param daemon_name: daemon name (in constants.DAEMONS_PORTS)
1749
  @rtype: int
1750

1751
  """
1752
  if daemon_name not in constants.DAEMONS_PORTS:
1753
    raise errors.ProgrammerError("Unknown daemon: %s" % daemon_name)
1754

    
1755
  (proto, default_port) = constants.DAEMONS_PORTS[daemon_name]
1756
  try:
1757
    port = socket.getservbyname(daemon_name, proto)
1758
  except socket.error:
1759
    port = default_port
1760

    
1761
  return port
1762

    
1763

    
1764
def SetupLogging(logfile, debug=False, stderr_logging=False, program="",
1765
                 multithreaded=False):
1766
  """Configures the logging module.
1767

1768
  @type logfile: str
1769
  @param logfile: the filename to which we should log
1770
  @type debug: boolean
1771
  @param debug: whether to enable debug messages too or
1772
      only those at C{INFO} and above level
1773
  @type stderr_logging: boolean
1774
  @param stderr_logging: whether we should also log to the standard error
1775
  @type program: str
1776
  @param program: the name under which we should log messages
1777
  @type multithreaded: boolean
1778
  @param multithreaded: if True, will add the thread name to the log file
1779
  @raise EnvironmentError: if we can't open the log file and
1780
      stderr logging is disabled
1781

1782
  """
1783
  fmt = "%(asctime)s: " + program + " pid=%(process)d"
1784
  if multithreaded:
1785
    fmt += "/%(threadName)s"
1786
  if debug:
1787
    fmt += " %(module)s:%(lineno)s"
1788
  fmt += " %(levelname)s %(message)s"
1789
  formatter = logging.Formatter(fmt)
1790

    
1791
  root_logger = logging.getLogger("")
1792
  root_logger.setLevel(logging.NOTSET)
1793

    
1794
  # Remove all previously setup handlers
1795
  for handler in root_logger.handlers:
1796
    handler.close()
1797
    root_logger.removeHandler(handler)
1798

    
1799
  if stderr_logging:
1800
    stderr_handler = logging.StreamHandler()
1801
    stderr_handler.setFormatter(formatter)
1802
    if debug:
1803
      stderr_handler.setLevel(logging.NOTSET)
1804
    else:
1805
      stderr_handler.setLevel(logging.CRITICAL)
1806
    root_logger.addHandler(stderr_handler)
1807

    
1808
  # this can fail, if the logging directories are not setup or we have
1809
  # a permisssion problem; in this case, it's best to log but ignore
1810
  # the error if stderr_logging is True, and if false we re-raise the
1811
  # exception since otherwise we could run but without any logs at all
1812
  try:
1813
    logfile_handler = logging.FileHandler(logfile)
1814
    logfile_handler.setFormatter(formatter)
1815
    if debug:
1816
      logfile_handler.setLevel(logging.DEBUG)
1817
    else:
1818
      logfile_handler.setLevel(logging.INFO)
1819
    root_logger.addHandler(logfile_handler)
1820
  except EnvironmentError:
1821
    if stderr_logging:
1822
      logging.exception("Failed to enable logging to file '%s'", logfile)
1823
    else:
1824
      # we need to re-raise the exception
1825
      raise
1826

    
1827

    
1828
def IsNormAbsPath(path):
1829
  """Check whether a path is absolute and also normalized
1830

1831
  This avoids things like /dir/../../other/path to be valid.
1832

1833
  """
1834
  return os.path.normpath(path) == path and os.path.isabs(path)
1835

    
1836

    
1837
def TailFile(fname, lines=20):
1838
  """Return the last lines from a file.
1839

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

1844
  @param fname: the file name
1845
  @type lines: int
1846
  @param lines: the (maximum) number of lines to return
1847

1848
  """
1849
  fd = open(fname, "r")
1850
  try:
1851
    fd.seek(0, 2)
1852
    pos = fd.tell()
1853
    pos = max(0, pos-4096)
1854
    fd.seek(pos, 0)
1855
    raw_data = fd.read()
1856
  finally:
1857
    fd.close()
1858

    
1859
  rows = raw_data.splitlines()
1860
  return rows[-lines:]
1861

    
1862

    
1863
def SafeEncode(text):
1864
  """Return a 'safe' version of a source string.
1865

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

1875
  @type text: str or unicode
1876
  @param text: input data
1877
  @rtype: str
1878
  @return: a safe version of text
1879

1880
  """
1881
  if isinstance(text, unicode):
1882
    # only if unicode; if str already, we handle it below
1883
    text = text.encode('ascii', 'backslashreplace')
1884
  resu = ""
1885
  for char in text:
1886
    c = ord(char)
1887
    if char  == '\t':
1888
      resu += r'\t'
1889
    elif char == '\n':
1890
      resu += r'\n'
1891
    elif char == '\r':
1892
      resu += r'\'r'
1893
    elif c < 32 or c >= 127: # non-printable
1894
      resu += "\\x%02x" % (c & 0xff)
1895
    else:
1896
      resu += char
1897
  return resu
1898

    
1899

    
1900
def CommaJoin(names):
1901
  """Nicely join a set of identifiers.
1902

1903
  @param names: set, list or tuple
1904
  @return: a string with the formatted results
1905

1906
  """
1907
  return ", ".join([str(val) for val in names])
1908

    
1909

    
1910
def BytesToMebibyte(value):
1911
  """Converts bytes to mebibytes.
1912

1913
  @type value: int
1914
  @param value: Value in bytes
1915
  @rtype: int
1916
  @return: Value in mebibytes
1917

1918
  """
1919
  return int(round(value / (1024.0 * 1024.0), 0))
1920

    
1921

    
1922
def CalculateDirectorySize(path):
1923
  """Calculates the size of a directory recursively.
1924

1925
  @type path: string
1926
  @param path: Path to directory
1927
  @rtype: int
1928
  @return: Size in mebibytes
1929

1930
  """
1931
  size = 0
1932

    
1933
  for (curpath, _, files) in os.walk(path):
1934
    for filename in files:
1935
      st = os.lstat(os.path.join(curpath, filename))
1936
      size += st.st_size
1937

    
1938
  return BytesToMebibyte(size)
1939

    
1940

    
1941
def GetFilesystemStats(path):
1942
  """Returns the total and free space on a filesystem.
1943

1944
  @type path: string
1945
  @param path: Path on filesystem to be examined
1946
  @rtype: int
1947
  @return: tuple of (Total space, Free space) in mebibytes
1948

1949
  """
1950
  st = os.statvfs(path)
1951

    
1952
  fsize = BytesToMebibyte(st.f_bavail * st.f_frsize)
1953
  tsize = BytesToMebibyte(st.f_blocks * st.f_frsize)
1954
  return (tsize, fsize)
1955

    
1956

    
1957
def LockedMethod(fn):
1958
  """Synchronized object access decorator.
1959

1960
  This decorator is intended to protect access to an object using the
1961
  object's own lock which is hardcoded to '_lock'.
1962

1963
  """
1964
  def _LockDebug(*args, **kwargs):
1965
    if debug_locks:
1966
      logging.debug(*args, **kwargs)
1967

    
1968
  def wrapper(self, *args, **kwargs):
1969
    # pylint: disable-msg=W0212
1970
    assert hasattr(self, '_lock')
1971
    lock = self._lock
1972
    _LockDebug("Waiting for %s", lock)
1973
    lock.acquire()
1974
    try:
1975
      _LockDebug("Acquired %s", lock)
1976
      result = fn(self, *args, **kwargs)
1977
    finally:
1978
      _LockDebug("Releasing %s", lock)
1979
      lock.release()
1980
      _LockDebug("Released %s", lock)
1981
    return result
1982
  return wrapper
1983

    
1984

    
1985
def LockFile(fd):
1986
  """Locks a file using POSIX locks.
1987

1988
  @type fd: int
1989
  @param fd: the file descriptor we need to lock
1990

1991
  """
1992
  try:
1993
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
1994
  except IOError, err:
1995
    if err.errno == errno.EAGAIN:
1996
      raise errors.LockError("File already locked")
1997
    raise
1998

    
1999

    
2000
def FormatTime(val):
2001
  """Formats a time value.
2002

2003
  @type val: float or None
2004
  @param val: the timestamp as returned by time.time()
2005
  @return: a string value or N/A if we don't have a valid timestamp
2006

2007
  """
2008
  if val is None or not isinstance(val, (int, float)):
2009
    return "N/A"
2010
  # these two codes works on Linux, but they are not guaranteed on all
2011
  # platforms
2012
  return time.strftime("%F %T", time.localtime(val))
2013

    
2014

    
2015
def ReadWatcherPauseFile(filename, now=None, remove_after=3600):
2016
  """Reads the watcher pause file.
2017

2018
  @type filename: string
2019
  @param filename: Path to watcher pause file
2020
  @type now: None, float or int
2021
  @param now: Current time as Unix timestamp
2022
  @type remove_after: int
2023
  @param remove_after: Remove watcher pause file after specified amount of
2024
    seconds past the pause end time
2025

2026
  """
2027
  if now is None:
2028
    now = time.time()
2029

    
2030
  try:
2031
    value = ReadFile(filename)
2032
  except IOError, err:
2033
    if err.errno != errno.ENOENT:
2034
      raise
2035
    value = None
2036

    
2037
  if value is not None:
2038
    try:
2039
      value = int(value)
2040
    except ValueError:
2041
      logging.warning(("Watcher pause file (%s) contains invalid value,"
2042
                       " removing it"), filename)
2043
      RemoveFile(filename)
2044
      value = None
2045

    
2046
    if value is not None:
2047
      # Remove file if it's outdated
2048
      if now > (value + remove_after):
2049
        RemoveFile(filename)
2050
        value = None
2051

    
2052
      elif now > value:
2053
        value = None
2054

    
2055
  return value
2056

    
2057

    
2058
class RetryTimeout(Exception):
2059
  """Retry loop timed out.
2060

2061
  """
2062

    
2063

    
2064
class RetryAgain(Exception):
2065
  """Retry again.
2066

2067
  """
2068

    
2069

    
2070
class _RetryDelayCalculator(object):
2071
  """Calculator for increasing delays.
2072

2073
  """
2074
  __slots__ = [
2075
    "_factor",
2076
    "_limit",
2077
    "_next",
2078
    "_start",
2079
    ]
2080

    
2081
  def __init__(self, start, factor, limit):
2082
    """Initializes this class.
2083

2084
    @type start: float
2085
    @param start: Initial delay
2086
    @type factor: float
2087
    @param factor: Factor for delay increase
2088
    @type limit: float or None
2089
    @param limit: Upper limit for delay or None for no limit
2090

2091
    """
2092
    assert start > 0.0
2093
    assert factor >= 1.0
2094
    assert limit is None or limit >= 0.0
2095

    
2096
    self._start = start
2097
    self._factor = factor
2098
    self._limit = limit
2099

    
2100
    self._next = start
2101

    
2102
  def __call__(self):
2103
    """Returns current delay and calculates the next one.
2104

2105
    """
2106
    current = self._next
2107

    
2108
    # Update for next run
2109
    if self._limit is None or self._next < self._limit:
2110
      self._next = max(self._limit, self._next * self._factor)
2111

    
2112
    return current
2113

    
2114

    
2115
#: Special delay to specify whole remaining timeout
2116
RETRY_REMAINING_TIME = object()
2117

    
2118

    
2119
def Retry(fn, delay, timeout, args=None, wait_fn=time.sleep,
2120
          _time_fn=time.time):
2121
  """Call a function repeatedly until it succeeds.
2122

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

2127
  C{delay} can be one of the following:
2128
    - callable returning the delay length as a float
2129
    - Tuple of (start, factor, limit)
2130
    - L{RETRY_REMAINING_TIME} to sleep until the timeout expires (this is
2131
      useful when overriding L{wait_fn} to wait for an external event)
2132
    - A static delay as a number (int or float)
2133

2134
  @type fn: callable
2135
  @param fn: Function to be called
2136
  @param delay: Either a callable (returning the delay), a tuple of (start,
2137
                factor, limit) (see L{_RetryDelayCalculator}),
2138
                L{RETRY_REMAINING_TIME} or a number (int or float)
2139
  @type timeout: float
2140
  @param timeout: Total timeout
2141
  @type wait_fn: callable
2142
  @param wait_fn: Waiting function
2143
  @return: Return value of function
2144

2145
  """
2146
  assert callable(fn)
2147
  assert callable(wait_fn)
2148
  assert callable(_time_fn)
2149

    
2150
  if args is None:
2151
    args = []
2152

    
2153
  end_time = _time_fn() + timeout
2154

    
2155
  if callable(delay):
2156
    # External function to calculate delay
2157
    calc_delay = delay
2158

    
2159
  elif isinstance(delay, (tuple, list)):
2160
    # Increasing delay with optional upper boundary
2161
    (start, factor, limit) = delay
2162
    calc_delay = _RetryDelayCalculator(start, factor, limit)
2163

    
2164
  elif delay is RETRY_REMAINING_TIME:
2165
    # Always use the remaining time
2166
    calc_delay = None
2167

    
2168
  else:
2169
    # Static delay
2170
    calc_delay = lambda: delay
2171

    
2172
  assert calc_delay is None or callable(calc_delay)
2173

    
2174
  while True:
2175
    try:
2176
      # pylint: disable-msg=W0142
2177
      return fn(*args)
2178
    except RetryAgain:
2179
      pass
2180

    
2181
    remaining_time = end_time - _time_fn()
2182

    
2183
    if remaining_time < 0.0:
2184
      raise RetryTimeout()
2185

    
2186
    assert remaining_time >= 0.0
2187

    
2188
    if calc_delay is None:
2189
      wait_fn(remaining_time)
2190
    else:
2191
      current_delay = calc_delay()
2192
      if current_delay > 0.0:
2193
        wait_fn(current_delay)
2194

    
2195

    
2196
class FileLock(object):
2197
  """Utility class for file locks.
2198

2199
  """
2200
  def __init__(self, filename):
2201
    """Constructor for FileLock.
2202

2203
    This will open the file denoted by the I{filename} argument.
2204

2205
    @type filename: str
2206
    @param filename: path to the file to be locked
2207

2208
    """
2209
    self.filename = filename
2210
    self.fd = open(self.filename, "w")
2211

    
2212
  def __del__(self):
2213
    self.Close()
2214

    
2215
  def Close(self):
2216
    """Close the file and release the lock.
2217

2218
    """
2219
    if self.fd:
2220
      self.fd.close()
2221
      self.fd = None
2222

    
2223
  def _flock(self, flag, blocking, timeout, errmsg):
2224
    """Wrapper for fcntl.flock.
2225

2226
    @type flag: int
2227
    @param flag: operation flag
2228
    @type blocking: bool
2229
    @param blocking: whether the operation should be done in blocking mode.
2230
    @type timeout: None or float
2231
    @param timeout: for how long the operation should be retried (implies
2232
                    non-blocking mode).
2233
    @type errmsg: string
2234
    @param errmsg: error message in case operation fails.
2235

2236
    """
2237
    assert self.fd, "Lock was closed"
2238
    assert timeout is None or timeout >= 0, \
2239
      "If specified, timeout must be positive"
2240

    
2241
    if timeout is not None:
2242
      flag |= fcntl.LOCK_NB
2243
      timeout_end = time.time() + timeout
2244

    
2245
    # Blocking doesn't have effect with timeout
2246
    elif not blocking:
2247
      flag |= fcntl.LOCK_NB
2248
      timeout_end = None
2249

    
2250
    # TODO: Convert to utils.Retry
2251

    
2252
    retry = True
2253
    while retry:
2254
      try:
2255
        fcntl.flock(self.fd, flag)
2256
        retry = False
2257
      except IOError, err:
2258
        if err.errno in (errno.EAGAIN, ):
2259
          if timeout_end is not None and time.time() < timeout_end:
2260
            # Wait before trying again
2261
            time.sleep(max(0.1, min(1.0, timeout)))
2262
          else:
2263
            raise errors.LockError(errmsg)
2264
        else:
2265
          logging.exception("fcntl.flock failed")
2266
          raise
2267

    
2268
  def Exclusive(self, blocking=False, timeout=None):
2269
    """Locks the file in exclusive mode.
2270

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

2278
    """
2279
    self._flock(fcntl.LOCK_EX, blocking, timeout,
2280
                "Failed to lock %s in exclusive mode" % self.filename)
2281

    
2282
  def Shared(self, blocking=False, timeout=None):
2283
    """Locks the file in shared mode.
2284

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

2292
    """
2293
    self._flock(fcntl.LOCK_SH, blocking, timeout,
2294
                "Failed to lock %s in shared mode" % self.filename)
2295

    
2296
  def Unlock(self, blocking=True, timeout=None):
2297
    """Unlocks the file.
2298

2299
    According to C{flock(2)}, unlocking can also be a nonblocking
2300
    operation::
2301

2302
      To make a non-blocking request, include LOCK_NB with any of the above
2303
      operations.
2304

2305
    @type blocking: boolean
2306
    @param blocking: whether to block and wait until we
2307
        can lock the file or return immediately
2308
    @type timeout: int or None
2309
    @param timeout: if not None, the duration to wait for the lock
2310
        (in blocking mode)
2311

2312
    """
2313
    self._flock(fcntl.LOCK_UN, blocking, timeout,
2314
                "Failed to unlock %s" % self.filename)
2315

    
2316

    
2317
def SignalHandled(signums):
2318
  """Signal Handled decoration.
2319

2320
  This special decorator installs a signal handler and then calls the target
2321
  function. The function must accept a 'signal_handlers' keyword argument,
2322
  which will contain a dict indexed by signal number, with SignalHandler
2323
  objects as values.
2324

2325
  The decorator can be safely stacked with iself, to handle multiple signals
2326
  with different handlers.
2327

2328
  @type signums: list
2329
  @param signums: signals to intercept
2330

2331
  """
2332
  def wrap(fn):
2333
    def sig_function(*args, **kwargs):
2334
      assert 'signal_handlers' not in kwargs or \
2335
             kwargs['signal_handlers'] is None or \
2336
             isinstance(kwargs['signal_handlers'], dict), \
2337
             "Wrong signal_handlers parameter in original function call"
2338
      if 'signal_handlers' in kwargs and kwargs['signal_handlers'] is not None:
2339
        signal_handlers = kwargs['signal_handlers']
2340
      else:
2341
        signal_handlers = {}
2342
        kwargs['signal_handlers'] = signal_handlers
2343
      sighandler = SignalHandler(signums)
2344
      try:
2345
        for sig in signums:
2346
          signal_handlers[sig] = sighandler
2347
        return fn(*args, **kwargs)
2348
      finally:
2349
        sighandler.Reset()
2350
    return sig_function
2351
  return wrap
2352

    
2353

    
2354
class SignalHandler(object):
2355
  """Generic signal handler class.
2356

2357
  It automatically restores the original handler when deconstructed or
2358
  when L{Reset} is called. You can either pass your own handler
2359
  function in or query the L{called} attribute to detect whether the
2360
  signal was sent.
2361

2362
  @type signum: list
2363
  @ivar signum: the signals we handle
2364
  @type called: boolean
2365
  @ivar called: tracks whether any of the signals have been raised
2366

2367
  """
2368
  def __init__(self, signum):
2369
    """Constructs a new SignalHandler instance.
2370

2371
    @type signum: int or list of ints
2372
    @param signum: Single signal number or set of signal numbers
2373

2374
    """
2375
    self.signum = set(signum)
2376
    self.called = False
2377

    
2378
    self._previous = {}
2379
    try:
2380
      for signum in self.signum:
2381
        # Setup handler
2382
        prev_handler = signal.signal(signum, self._HandleSignal)
2383
        try:
2384
          self._previous[signum] = prev_handler
2385
        except:
2386
          # Restore previous handler
2387
          signal.signal(signum, prev_handler)
2388
          raise
2389
    except:
2390
      # Reset all handlers
2391
      self.Reset()
2392
      # Here we have a race condition: a handler may have already been called,
2393
      # but there's not much we can do about it at this point.
2394
      raise
2395

    
2396
  def __del__(self):
2397
    self.Reset()
2398

    
2399
  def Reset(self):
2400
    """Restore previous handler.
2401

2402
    This will reset all the signals to their previous handlers.
2403

2404
    """
2405
    for signum, prev_handler in self._previous.items():
2406
      signal.signal(signum, prev_handler)
2407
      # If successful, remove from dict
2408
      del self._previous[signum]
2409

    
2410
  def Clear(self):
2411
    """Unsets the L{called} flag.
2412

2413
    This function can be used in case a signal may arrive several times.
2414

2415
    """
2416
    self.called = False
2417

    
2418
  # we don't care about arguments, but we leave them named for the future
2419
  def _HandleSignal(self, signum, frame): # pylint: disable-msg=W0613
2420
    """Actual signal handling function.
2421

2422
    """
2423
    # This is not nice and not absolutely atomic, but it appears to be the only
2424
    # solution in Python -- there are no atomic types.
2425
    self.called = True
2426

    
2427

    
2428
class FieldSet(object):
2429
  """A simple field set.
2430

2431
  Among the features are:
2432
    - checking if a string is among a list of static string or regex objects
2433
    - checking if a whole list of string matches
2434
    - returning the matching groups from a regex match
2435

2436
  Internally, all fields are held as regular expression objects.
2437

2438
  """
2439
  def __init__(self, *items):
2440
    self.items = [re.compile("^%s$" % value) for value in items]
2441

    
2442
  def Extend(self, other_set):
2443
    """Extend the field set with the items from another one"""
2444
    self.items.extend(other_set.items)
2445

    
2446
  def Matches(self, field):
2447
    """Checks if a field matches the current set
2448

2449
    @type field: str
2450
    @param field: the string to match
2451
    @return: either None or a regular expression match object
2452

2453
    """
2454
    for m in itertools.ifilter(None, (val.match(field) for val in self.items)):
2455
      return m
2456
    return None
2457

    
2458
  def NonMatching(self, items):
2459
    """Returns the list of fields not matching the current set
2460

2461
    @type items: list
2462
    @param items: the list of fields to check
2463
    @rtype: list
2464
    @return: list of non-matching fields
2465

2466
    """
2467
    return [val for val in items if not self.Matches(val)]