Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ ea34193f

History | View | Annotate | Download (67.3 kB)

1
#
2
#
3

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

    
21

    
22
"""Ganeti utility module.
23

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

27
"""
28

    
29

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

    
47
from cStringIO import StringIO
48

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

    
55
from ganeti import errors
56
from ganeti import constants
57

    
58

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

    
62
debug_locks = False
63

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

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

    
69

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

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

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

    
92

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

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

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

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

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

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

    
120

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

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

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

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

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

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

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

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

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

    
182

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

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

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

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

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

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

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

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

    
251

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

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

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

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

    
284

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

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

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

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

    
301

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

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

314
  """
315
  try:
316
    return os.rename(old, new)
317
  except OSError, err:
318
    # In at least one use case of this function, the job queue, directory
319
    # creation is very rare. Checking for the directory before renaming is not
320
    # as efficient.
321
    if mkdir and err.errno == errno.ENOENT:
322
      # Create directory and try again
323
      dirname = os.path.dirname(new)
324
      try:
325
        os.makedirs(dirname, mode=mkdir_mode)
326
      except OSError, err:
327
        # Ignore EEXIST. This is only handled in os.makedirs as included in
328
        # Python 2.5 and above.
329
        if err.errno != errno.EEXIST or not os.path.exists(dirname):
330
          raise
331

    
332
      return os.rename(old, new)
333

    
334
    raise
335

    
336

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

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

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

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

    
353
  f = open(filename)
354

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

    
361
    fp.update(data)
362

    
363
  return fp.hexdigest()
364

    
365

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

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

375
  """
376
  ret = {}
377

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

    
383
  return ret
384

    
385

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

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

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

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

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

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

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

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

    
452

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

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

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

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

    
475

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

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

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

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

    
499
  return pid
500

    
501

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

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

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

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

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

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

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

    
547

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    
600
    return result
601

    
602

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

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

    
612

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

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

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

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

    
636
    retval[name] = size
637

    
638
  return retval
639

    
640

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

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

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

    
652

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

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

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

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

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

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

    
689

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

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

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

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

    
711

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

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

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

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

    
729

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

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

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

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

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

    
748

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

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

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

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

    
770

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

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

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

    
789
  suffix = ''
790

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

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

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

    
806

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

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

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

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

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

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

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

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

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

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

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

    
849
  return value
850

    
851

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

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

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

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

    
880

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

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

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

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

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

    
913

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

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

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

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

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

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

    
960

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

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

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

    
972

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

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

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

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

    
1002
          out.write(line)
1003

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

    
1016

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

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

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

    
1030

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

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

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

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

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

    
1059
  return backup_name
1060

    
1061

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

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

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

    
1076

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

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

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

    
1088

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

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

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

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

    
1112
  success = False
1113

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

    
1121
  sock.settimeout(timeout)
1122

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

    
1132
  return success
1133

    
1134

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

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

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

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

    
1150

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

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

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

    
1164

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

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

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

    
1185

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

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

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

    
1196

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

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

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

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

    
1211

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

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

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

    
1229

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

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

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

    
1245

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

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

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

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

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

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

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

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

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

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

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

    
1342
  return result
1343

    
1344

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

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

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

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

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

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

    
1372

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

    
1379

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

    
1386

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

1390
  Element order is preserved.
1391

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

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

    
1401

    
1402
def NormalizeAndValidateMac(mac):
1403
  """Normalizes and check if a MAC address is valid.
1404

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

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

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

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

    
1421
  return mac.lower()
1422

    
1423

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

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

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

    
1438

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

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

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

    
1458

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

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

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

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

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

    
1490

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

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

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

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

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

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

    
1533

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

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

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

    
1546

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

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

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

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

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

    
1565

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

1569
  Any errors are ignored.
1570

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

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

    
1582

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

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

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

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

    
1614
  if not IsProcessAlive(pid):
1615
    return
1616

    
1617
  _helper(pid, signal_, waitpid)
1618

    
1619
  if timeout <= 0:
1620
    return
1621

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

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

    
1631
    if result_pid > 0:
1632
      return
1633

    
1634
    raise RetryAgain()
1635

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

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

    
1646

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

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

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

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

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

    
1679

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

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

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

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

    
1704

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

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

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

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

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

    
1722

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

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

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

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

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

    
1740

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

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

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

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

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

    
1762
  return port
1763

    
1764

    
1765
def SetupLogging(logfile, debug=0, stderr_logging=False, program="",
1766
                 multithreaded=False, syslog=constants.SYSLOG_USAGE):
1767
  """Configures the logging module.
1768

1769
  @type logfile: str
1770
  @param logfile: the filename to which we should log
1771
  @type debug: integer
1772
  @param debug: if greater than zero, enable debug messages, otherwise
1773
      only those at C{INFO} and above level
1774
  @type stderr_logging: boolean
1775
  @param stderr_logging: whether we should also log to the standard error
1776
  @type program: str
1777
  @param program: the name under which we should log messages
1778
  @type multithreaded: boolean
1779
  @param multithreaded: if True, will add the thread name to the log file
1780
  @type syslog: string
1781
  @param syslog: one of 'no', 'yes', 'only':
1782
      - if no, syslog is not used
1783
      - if yes, syslog is used (in addition to file-logging)
1784
      - if only, only syslog is used
1785
  @raise EnvironmentError: if we can't open the log file and
1786
      syslog/stderr logging is disabled
1787

1788
  """
1789
  fmt = "%(asctime)s: " + program + " pid=%(process)d"
1790
  sft = program + "[%(process)d]:"
1791
  if multithreaded:
1792
    fmt += "/%(threadName)s"
1793
    sft += " (%(threadName)s)"
1794
  if debug:
1795
    fmt += " %(module)s:%(lineno)s"
1796
    # no debug info for syslog loggers
1797
  fmt += " %(levelname)s %(message)s"
1798
  # yes, we do want the textual level, as remote syslog will probably
1799
  # lose the error level, and it's easier to grep for it
1800
  sft += " %(levelname)s %(message)s"
1801
  formatter = logging.Formatter(fmt)
1802
  sys_fmt = logging.Formatter(sft)
1803

    
1804
  root_logger = logging.getLogger("")
1805
  root_logger.setLevel(logging.NOTSET)
1806

    
1807
  # Remove all previously setup handlers
1808
  for handler in root_logger.handlers:
1809
    handler.close()
1810
    root_logger.removeHandler(handler)
1811

    
1812
  if stderr_logging:
1813
    stderr_handler = logging.StreamHandler()
1814
    stderr_handler.setFormatter(formatter)
1815
    if debug:
1816
      stderr_handler.setLevel(logging.NOTSET)
1817
    else:
1818
      stderr_handler.setLevel(logging.CRITICAL)
1819
    root_logger.addHandler(stderr_handler)
1820

    
1821
  if syslog in (constants.SYSLOG_YES, constants.SYSLOG_ONLY):
1822
    facility = logging.handlers.SysLogHandler.LOG_DAEMON
1823
    syslog_handler = logging.handlers.SysLogHandler(constants.SYSLOG_SOCKET,
1824
                                                    facility)
1825
    syslog_handler.setFormatter(sys_fmt)
1826
    # Never enable debug over syslog
1827
    syslog_handler.setLevel(logging.INFO)
1828
    root_logger.addHandler(syslog_handler)
1829

    
1830
  if syslog != constants.SYSLOG_ONLY:
1831
    # this can fail, if the logging directories are not setup or we have
1832
    # a permisssion problem; in this case, it's best to log but ignore
1833
    # the error if stderr_logging is True, and if false we re-raise the
1834
    # exception since otherwise we could run but without any logs at all
1835
    try:
1836
      logfile_handler = logging.FileHandler(logfile)
1837
      logfile_handler.setFormatter(formatter)
1838
      if debug:
1839
        logfile_handler.setLevel(logging.DEBUG)
1840
      else:
1841
        logfile_handler.setLevel(logging.INFO)
1842
      root_logger.addHandler(logfile_handler)
1843
    except EnvironmentError:
1844
      if stderr_logging or syslog == constants.SYSLOG_YES:
1845
        logging.exception("Failed to enable logging to file '%s'", logfile)
1846
      else:
1847
        # we need to re-raise the exception
1848
        raise
1849

    
1850

    
1851
def IsNormAbsPath(path):
1852
  """Check whether a path is absolute and also normalized
1853

1854
  This avoids things like /dir/../../other/path to be valid.
1855

1856
  """
1857
  return os.path.normpath(path) == path and os.path.isabs(path)
1858

    
1859

    
1860
def TailFile(fname, lines=20):
1861
  """Return the last lines from a file.
1862

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

1867
  @param fname: the file name
1868
  @type lines: int
1869
  @param lines: the (maximum) number of lines to return
1870

1871
  """
1872
  fd = open(fname, "r")
1873
  try:
1874
    fd.seek(0, 2)
1875
    pos = fd.tell()
1876
    pos = max(0, pos-4096)
1877
    fd.seek(pos, 0)
1878
    raw_data = fd.read()
1879
  finally:
1880
    fd.close()
1881

    
1882
  rows = raw_data.splitlines()
1883
  return rows[-lines:]
1884

    
1885

    
1886
def SafeEncode(text):
1887
  """Return a 'safe' version of a source string.
1888

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

1898
  @type text: str or unicode
1899
  @param text: input data
1900
  @rtype: str
1901
  @return: a safe version of text
1902

1903
  """
1904
  if isinstance(text, unicode):
1905
    # only if unicode; if str already, we handle it below
1906
    text = text.encode('ascii', 'backslashreplace')
1907
  resu = ""
1908
  for char in text:
1909
    c = ord(char)
1910
    if char  == '\t':
1911
      resu += r'\t'
1912
    elif char == '\n':
1913
      resu += r'\n'
1914
    elif char == '\r':
1915
      resu += r'\'r'
1916
    elif c < 32 or c >= 127: # non-printable
1917
      resu += "\\x%02x" % (c & 0xff)
1918
    else:
1919
      resu += char
1920
  return resu
1921

    
1922

    
1923
def UnescapeAndSplit(text, sep=","):
1924
  """Split and unescape a string based on a given separator.
1925

1926
  This function splits a string based on a separator where the
1927
  separator itself can be escape in order to be an element of the
1928
  elements. The escaping rules are (assuming coma being the
1929
  separator):
1930
    - a plain , separates the elements
1931
    - a sequence \\\\, (double backslash plus comma) is handled as a
1932
      backslash plus a separator comma
1933
    - a sequence \, (backslash plus comma) is handled as a
1934
      non-separator comma
1935

1936
  @type text: string
1937
  @param text: the string to split
1938
  @type sep: string
1939
  @param text: the separator
1940
  @rtype: string
1941
  @return: a list of strings
1942

1943
  """
1944
  # we split the list by sep (with no escaping at this stage)
1945
  slist = text.split(sep)
1946
  # next, we revisit the elements and if any of them ended with an odd
1947
  # number of backslashes, then we join it with the next
1948
  rlist = []
1949
  while slist:
1950
    e1 = slist.pop(0)
1951
    if e1.endswith("\\"):
1952
      num_b = len(e1) - len(e1.rstrip("\\"))
1953
      if num_b % 2 == 1:
1954
        e2 = slist.pop(0)
1955
        # here the backslashes remain (all), and will be reduced in
1956
        # the next step
1957
        rlist.append(e1 + sep + e2)
1958
        continue
1959
    rlist.append(e1)
1960
  # finally, replace backslash-something with something
1961
  rlist = [re.sub(r"\\(.)", r"\1", v) for v in rlist]
1962
  return rlist
1963

    
1964

    
1965
def CommaJoin(names):
1966
  """Nicely join a set of identifiers.
1967

1968
  @param names: set, list or tuple
1969
  @return: a string with the formatted results
1970

1971
  """
1972
  return ", ".join([str(val) for val in names])
1973

    
1974

    
1975
def BytesToMebibyte(value):
1976
  """Converts bytes to mebibytes.
1977

1978
  @type value: int
1979
  @param value: Value in bytes
1980
  @rtype: int
1981
  @return: Value in mebibytes
1982

1983
  """
1984
  return int(round(value / (1024.0 * 1024.0), 0))
1985

    
1986

    
1987
def CalculateDirectorySize(path):
1988
  """Calculates the size of a directory recursively.
1989

1990
  @type path: string
1991
  @param path: Path to directory
1992
  @rtype: int
1993
  @return: Size in mebibytes
1994

1995
  """
1996
  size = 0
1997

    
1998
  for (curpath, _, files) in os.walk(path):
1999
    for filename in files:
2000
      st = os.lstat(os.path.join(curpath, filename))
2001
      size += st.st_size
2002

    
2003
  return BytesToMebibyte(size)
2004

    
2005

    
2006
def GetFilesystemStats(path):
2007
  """Returns the total and free space on a filesystem.
2008

2009
  @type path: string
2010
  @param path: Path on filesystem to be examined
2011
  @rtype: int
2012
  @return: tuple of (Total space, Free space) in mebibytes
2013

2014
  """
2015
  st = os.statvfs(path)
2016

    
2017
  fsize = BytesToMebibyte(st.f_bavail * st.f_frsize)
2018
  tsize = BytesToMebibyte(st.f_blocks * st.f_frsize)
2019
  return (tsize, fsize)
2020

    
2021

    
2022
def LockedMethod(fn):
2023
  """Synchronized object access decorator.
2024

2025
  This decorator is intended to protect access to an object using the
2026
  object's own lock which is hardcoded to '_lock'.
2027

2028
  """
2029
  def _LockDebug(*args, **kwargs):
2030
    if debug_locks:
2031
      logging.debug(*args, **kwargs)
2032

    
2033
  def wrapper(self, *args, **kwargs):
2034
    # pylint: disable-msg=W0212
2035
    assert hasattr(self, '_lock')
2036
    lock = self._lock
2037
    _LockDebug("Waiting for %s", lock)
2038
    lock.acquire()
2039
    try:
2040
      _LockDebug("Acquired %s", lock)
2041
      result = fn(self, *args, **kwargs)
2042
    finally:
2043
      _LockDebug("Releasing %s", lock)
2044
      lock.release()
2045
      _LockDebug("Released %s", lock)
2046
    return result
2047
  return wrapper
2048

    
2049

    
2050
def LockFile(fd):
2051
  """Locks a file using POSIX locks.
2052

2053
  @type fd: int
2054
  @param fd: the file descriptor we need to lock
2055

2056
  """
2057
  try:
2058
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
2059
  except IOError, err:
2060
    if err.errno == errno.EAGAIN:
2061
      raise errors.LockError("File already locked")
2062
    raise
2063

    
2064

    
2065
def FormatTime(val):
2066
  """Formats a time value.
2067

2068
  @type val: float or None
2069
  @param val: the timestamp as returned by time.time()
2070
  @return: a string value or N/A if we don't have a valid timestamp
2071

2072
  """
2073
  if val is None or not isinstance(val, (int, float)):
2074
    return "N/A"
2075
  # these two codes works on Linux, but they are not guaranteed on all
2076
  # platforms
2077
  return time.strftime("%F %T", time.localtime(val))
2078

    
2079

    
2080
def ReadWatcherPauseFile(filename, now=None, remove_after=3600):
2081
  """Reads the watcher pause file.
2082

2083
  @type filename: string
2084
  @param filename: Path to watcher pause file
2085
  @type now: None, float or int
2086
  @param now: Current time as Unix timestamp
2087
  @type remove_after: int
2088
  @param remove_after: Remove watcher pause file after specified amount of
2089
    seconds past the pause end time
2090

2091
  """
2092
  if now is None:
2093
    now = time.time()
2094

    
2095
  try:
2096
    value = ReadFile(filename)
2097
  except IOError, err:
2098
    if err.errno != errno.ENOENT:
2099
      raise
2100
    value = None
2101

    
2102
  if value is not None:
2103
    try:
2104
      value = int(value)
2105
    except ValueError:
2106
      logging.warning(("Watcher pause file (%s) contains invalid value,"
2107
                       " removing it"), filename)
2108
      RemoveFile(filename)
2109
      value = None
2110

    
2111
    if value is not None:
2112
      # Remove file if it's outdated
2113
      if now > (value + remove_after):
2114
        RemoveFile(filename)
2115
        value = None
2116

    
2117
      elif now > value:
2118
        value = None
2119

    
2120
  return value
2121

    
2122

    
2123
class RetryTimeout(Exception):
2124
  """Retry loop timed out.
2125

2126
  """
2127

    
2128

    
2129
class RetryAgain(Exception):
2130
  """Retry again.
2131

2132
  """
2133

    
2134

    
2135
class _RetryDelayCalculator(object):
2136
  """Calculator for increasing delays.
2137

2138
  """
2139
  __slots__ = [
2140
    "_factor",
2141
    "_limit",
2142
    "_next",
2143
    "_start",
2144
    ]
2145

    
2146
  def __init__(self, start, factor, limit):
2147
    """Initializes this class.
2148

2149
    @type start: float
2150
    @param start: Initial delay
2151
    @type factor: float
2152
    @param factor: Factor for delay increase
2153
    @type limit: float or None
2154
    @param limit: Upper limit for delay or None for no limit
2155

2156
    """
2157
    assert start > 0.0
2158
    assert factor >= 1.0
2159
    assert limit is None or limit >= 0.0
2160

    
2161
    self._start = start
2162
    self._factor = factor
2163
    self._limit = limit
2164

    
2165
    self._next = start
2166

    
2167
  def __call__(self):
2168
    """Returns current delay and calculates the next one.
2169

2170
    """
2171
    current = self._next
2172

    
2173
    # Update for next run
2174
    if self._limit is None or self._next < self._limit:
2175
      self._next = min(self._limit, self._next * self._factor)
2176

    
2177
    return current
2178

    
2179

    
2180
#: Special delay to specify whole remaining timeout
2181
RETRY_REMAINING_TIME = object()
2182

    
2183

    
2184
def Retry(fn, delay, timeout, args=None, wait_fn=time.sleep,
2185
          _time_fn=time.time):
2186
  """Call a function repeatedly until it succeeds.
2187

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

2192
  C{delay} can be one of the following:
2193
    - callable returning the delay length as a float
2194
    - Tuple of (start, factor, limit)
2195
    - L{RETRY_REMAINING_TIME} to sleep until the timeout expires (this is
2196
      useful when overriding L{wait_fn} to wait for an external event)
2197
    - A static delay as a number (int or float)
2198

2199
  @type fn: callable
2200
  @param fn: Function to be called
2201
  @param delay: Either a callable (returning the delay), a tuple of (start,
2202
                factor, limit) (see L{_RetryDelayCalculator}),
2203
                L{RETRY_REMAINING_TIME} or a number (int or float)
2204
  @type timeout: float
2205
  @param timeout: Total timeout
2206
  @type wait_fn: callable
2207
  @param wait_fn: Waiting function
2208
  @return: Return value of function
2209

2210
  """
2211
  assert callable(fn)
2212
  assert callable(wait_fn)
2213
  assert callable(_time_fn)
2214

    
2215
  if args is None:
2216
    args = []
2217

    
2218
  end_time = _time_fn() + timeout
2219

    
2220
  if callable(delay):
2221
    # External function to calculate delay
2222
    calc_delay = delay
2223

    
2224
  elif isinstance(delay, (tuple, list)):
2225
    # Increasing delay with optional upper boundary
2226
    (start, factor, limit) = delay
2227
    calc_delay = _RetryDelayCalculator(start, factor, limit)
2228

    
2229
  elif delay is RETRY_REMAINING_TIME:
2230
    # Always use the remaining time
2231
    calc_delay = None
2232

    
2233
  else:
2234
    # Static delay
2235
    calc_delay = lambda: delay
2236

    
2237
  assert calc_delay is None or callable(calc_delay)
2238

    
2239
  while True:
2240
    try:
2241
      # pylint: disable-msg=W0142
2242
      return fn(*args)
2243
    except RetryAgain:
2244
      pass
2245

    
2246
    remaining_time = end_time - _time_fn()
2247

    
2248
    if remaining_time < 0.0:
2249
      raise RetryTimeout()
2250

    
2251
    assert remaining_time >= 0.0
2252

    
2253
    if calc_delay is None:
2254
      wait_fn(remaining_time)
2255
    else:
2256
      current_delay = calc_delay()
2257
      if current_delay > 0.0:
2258
        wait_fn(current_delay)
2259

    
2260

    
2261
class FileLock(object):
2262
  """Utility class for file locks.
2263

2264
  """
2265
  def __init__(self, filename):
2266
    """Constructor for FileLock.
2267

2268
    This will open the file denoted by the I{filename} argument.
2269

2270
    @type filename: str
2271
    @param filename: path to the file to be locked
2272

2273
    """
2274
    self.filename = filename
2275
    self.fd = open(self.filename, "w")
2276

    
2277
  def __del__(self):
2278
    self.Close()
2279

    
2280
  def Close(self):
2281
    """Close the file and release the lock.
2282

2283
    """
2284
    if hasattr(self, "fd") and self.fd:
2285
      self.fd.close()
2286
      self.fd = None
2287

    
2288
  def _flock(self, flag, blocking, timeout, errmsg):
2289
    """Wrapper for fcntl.flock.
2290

2291
    @type flag: int
2292
    @param flag: operation flag
2293
    @type blocking: bool
2294
    @param blocking: whether the operation should be done in blocking mode.
2295
    @type timeout: None or float
2296
    @param timeout: for how long the operation should be retried (implies
2297
                    non-blocking mode).
2298
    @type errmsg: string
2299
    @param errmsg: error message in case operation fails.
2300

2301
    """
2302
    assert self.fd, "Lock was closed"
2303
    assert timeout is None or timeout >= 0, \
2304
      "If specified, timeout must be positive"
2305

    
2306
    if timeout is not None:
2307
      flag |= fcntl.LOCK_NB
2308
      timeout_end = time.time() + timeout
2309

    
2310
    # Blocking doesn't have effect with timeout
2311
    elif not blocking:
2312
      flag |= fcntl.LOCK_NB
2313
      timeout_end = None
2314

    
2315
    # TODO: Convert to utils.Retry
2316

    
2317
    retry = True
2318
    while retry:
2319
      try:
2320
        fcntl.flock(self.fd, flag)
2321
        retry = False
2322
      except IOError, err:
2323
        if err.errno in (errno.EAGAIN, ):
2324
          if timeout_end is not None and time.time() < timeout_end:
2325
            # Wait before trying again
2326
            time.sleep(max(0.1, min(1.0, timeout)))
2327
          else:
2328
            raise errors.LockError(errmsg)
2329
        else:
2330
          logging.exception("fcntl.flock failed")
2331
          raise
2332

    
2333
  def Exclusive(self, blocking=False, timeout=None):
2334
    """Locks the file in exclusive mode.
2335

2336
    @type blocking: boolean
2337
    @param blocking: whether to block and wait until we
2338
        can lock the file or return immediately
2339
    @type timeout: int or None
2340
    @param timeout: if not None, the duration to wait for the lock
2341
        (in blocking mode)
2342

2343
    """
2344
    self._flock(fcntl.LOCK_EX, blocking, timeout,
2345
                "Failed to lock %s in exclusive mode" % self.filename)
2346

    
2347
  def Shared(self, blocking=False, timeout=None):
2348
    """Locks the file in shared mode.
2349

2350
    @type blocking: boolean
2351
    @param blocking: whether to block and wait until we
2352
        can lock the file or return immediately
2353
    @type timeout: int or None
2354
    @param timeout: if not None, the duration to wait for the lock
2355
        (in blocking mode)
2356

2357
    """
2358
    self._flock(fcntl.LOCK_SH, blocking, timeout,
2359
                "Failed to lock %s in shared mode" % self.filename)
2360

    
2361
  def Unlock(self, blocking=True, timeout=None):
2362
    """Unlocks the file.
2363

2364
    According to C{flock(2)}, unlocking can also be a nonblocking
2365
    operation::
2366

2367
      To make a non-blocking request, include LOCK_NB with any of the above
2368
      operations.
2369

2370
    @type blocking: boolean
2371
    @param blocking: whether to block and wait until we
2372
        can lock the file or return immediately
2373
    @type timeout: int or None
2374
    @param timeout: if not None, the duration to wait for the lock
2375
        (in blocking mode)
2376

2377
    """
2378
    self._flock(fcntl.LOCK_UN, blocking, timeout,
2379
                "Failed to unlock %s" % self.filename)
2380

    
2381

    
2382
def SignalHandled(signums):
2383
  """Signal Handled decoration.
2384

2385
  This special decorator installs a signal handler and then calls the target
2386
  function. The function must accept a 'signal_handlers' keyword argument,
2387
  which will contain a dict indexed by signal number, with SignalHandler
2388
  objects as values.
2389

2390
  The decorator can be safely stacked with iself, to handle multiple signals
2391
  with different handlers.
2392

2393
  @type signums: list
2394
  @param signums: signals to intercept
2395

2396
  """
2397
  def wrap(fn):
2398
    def sig_function(*args, **kwargs):
2399
      assert 'signal_handlers' not in kwargs or \
2400
             kwargs['signal_handlers'] is None or \
2401
             isinstance(kwargs['signal_handlers'], dict), \
2402
             "Wrong signal_handlers parameter in original function call"
2403
      if 'signal_handlers' in kwargs and kwargs['signal_handlers'] is not None:
2404
        signal_handlers = kwargs['signal_handlers']
2405
      else:
2406
        signal_handlers = {}
2407
        kwargs['signal_handlers'] = signal_handlers
2408
      sighandler = SignalHandler(signums)
2409
      try:
2410
        for sig in signums:
2411
          signal_handlers[sig] = sighandler
2412
        return fn(*args, **kwargs)
2413
      finally:
2414
        sighandler.Reset()
2415
    return sig_function
2416
  return wrap
2417

    
2418

    
2419
class SignalHandler(object):
2420
  """Generic signal handler class.
2421

2422
  It automatically restores the original handler when deconstructed or
2423
  when L{Reset} is called. You can either pass your own handler
2424
  function in or query the L{called} attribute to detect whether the
2425
  signal was sent.
2426

2427
  @type signum: list
2428
  @ivar signum: the signals we handle
2429
  @type called: boolean
2430
  @ivar called: tracks whether any of the signals have been raised
2431

2432
  """
2433
  def __init__(self, signum):
2434
    """Constructs a new SignalHandler instance.
2435

2436
    @type signum: int or list of ints
2437
    @param signum: Single signal number or set of signal numbers
2438

2439
    """
2440
    self.signum = set(signum)
2441
    self.called = False
2442

    
2443
    self._previous = {}
2444
    try:
2445
      for signum in self.signum:
2446
        # Setup handler
2447
        prev_handler = signal.signal(signum, self._HandleSignal)
2448
        try:
2449
          self._previous[signum] = prev_handler
2450
        except:
2451
          # Restore previous handler
2452
          signal.signal(signum, prev_handler)
2453
          raise
2454
    except:
2455
      # Reset all handlers
2456
      self.Reset()
2457
      # Here we have a race condition: a handler may have already been called,
2458
      # but there's not much we can do about it at this point.
2459
      raise
2460

    
2461
  def __del__(self):
2462
    self.Reset()
2463

    
2464
  def Reset(self):
2465
    """Restore previous handler.
2466

2467
    This will reset all the signals to their previous handlers.
2468

2469
    """
2470
    for signum, prev_handler in self._previous.items():
2471
      signal.signal(signum, prev_handler)
2472
      # If successful, remove from dict
2473
      del self._previous[signum]
2474

    
2475
  def Clear(self):
2476
    """Unsets the L{called} flag.
2477

2478
    This function can be used in case a signal may arrive several times.
2479

2480
    """
2481
    self.called = False
2482

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

2487
    """
2488
    # This is not nice and not absolutely atomic, but it appears to be the only
2489
    # solution in Python -- there are no atomic types.
2490
    self.called = True
2491

    
2492

    
2493
class FieldSet(object):
2494
  """A simple field set.
2495

2496
  Among the features are:
2497
    - checking if a string is among a list of static string or regex objects
2498
    - checking if a whole list of string matches
2499
    - returning the matching groups from a regex match
2500

2501
  Internally, all fields are held as regular expression objects.
2502

2503
  """
2504
  def __init__(self, *items):
2505
    self.items = [re.compile("^%s$" % value) for value in items]
2506

    
2507
  def Extend(self, other_set):
2508
    """Extend the field set with the items from another one"""
2509
    self.items.extend(other_set.items)
2510

    
2511
  def Matches(self, field):
2512
    """Checks if a field matches the current set
2513

2514
    @type field: str
2515
    @param field: the string to match
2516
    @return: either None or a regular expression match object
2517

2518
    """
2519
    for m in itertools.ifilter(None, (val.match(field) for val in self.items)):
2520
      return m
2521
    return None
2522

    
2523
  def NonMatching(self, items):
2524
    """Returns the list of fields not matching the current set
2525

2526
    @type items: list
2527
    @param items: the list of fields to check
2528
    @rtype: list
2529
    @return: list of non-matching fields
2530

2531
    """
2532
    return [val for val in items if not self.Matches(val)]