Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ 4c32a8bd

History | View | Annotate | Download (85.2 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
import datetime
47
import calendar
48
import collections
49
import struct
50
import IN
51

    
52
from cStringIO import StringIO
53

    
54
try:
55
  from hashlib import sha1
56
except ImportError:
57
  import sha
58
  sha1 = sha.new
59

    
60
try:
61
  import ctypes
62
except ImportError:
63
  ctypes = None
64

    
65
from ganeti import errors
66
from ganeti import constants
67

    
68

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

    
72
debug_locks = False
73

    
74
#: when set to True, L{RunCmd} is disabled
75
no_fork = False
76

    
77
_RANDOM_UUID_FILE = "/proc/sys/kernel/random/uuid"
78

    
79
# Structure definition for getsockopt(SOL_SOCKET, SO_PEERCRED, ...):
80
# struct ucred { pid_t pid; uid_t uid; gid_t gid; };
81
#
82
# The GNU C Library defines gid_t and uid_t to be "unsigned int" and
83
# pid_t to "int".
84
#
85
# IEEE Std 1003.1-2008:
86
# "nlink_t, uid_t, gid_t, and id_t shall be integer types"
87
# "blksize_t, pid_t, and ssize_t shall be signed integer types"
88
_STRUCT_UCRED = "iII"
89
_STRUCT_UCRED_SIZE = struct.calcsize(_STRUCT_UCRED)
90

    
91
# Flags for mlockall() (from bits/mman.h)
92
_MCL_CURRENT = 1
93
_MCL_FUTURE = 2
94

    
95

    
96
class RunResult(object):
97
  """Holds the result of running external programs.
98

99
  @type exit_code: int
100
  @ivar exit_code: the exit code of the program, or None (if the program
101
      didn't exit())
102
  @type signal: int or None
103
  @ivar signal: the signal that caused the program to finish, or None
104
      (if the program wasn't terminated by a signal)
105
  @type stdout: str
106
  @ivar stdout: the standard output of the program
107
  @type stderr: str
108
  @ivar stderr: the standard error of the program
109
  @type failed: boolean
110
  @ivar failed: True in case the program was
111
      terminated by a signal or exited with a non-zero exit code
112
  @ivar fail_reason: a string detailing the termination reason
113

114
  """
115
  __slots__ = ["exit_code", "signal", "stdout", "stderr",
116
               "failed", "fail_reason", "cmd"]
117

    
118

    
119
  def __init__(self, exit_code, signal_, stdout, stderr, cmd):
120
    self.cmd = cmd
121
    self.exit_code = exit_code
122
    self.signal = signal_
123
    self.stdout = stdout
124
    self.stderr = stderr
125
    self.failed = (signal_ is not None or exit_code != 0)
126

    
127
    if self.signal is not None:
128
      self.fail_reason = "terminated by signal %s" % self.signal
129
    elif self.exit_code is not None:
130
      self.fail_reason = "exited with exit code %s" % self.exit_code
131
    else:
132
      self.fail_reason = "unable to determine termination reason"
133

    
134
    if self.failed:
135
      logging.debug("Command '%s' failed (%s); output: %s",
136
                    self.cmd, self.fail_reason, self.output)
137

    
138
  def _GetOutput(self):
139
    """Returns the combined stdout and stderr for easier usage.
140

141
    """
142
    return self.stdout + self.stderr
143

    
144
  output = property(_GetOutput, None, None, "Return full output")
145

    
146

    
147
def RunCmd(cmd, env=None, output=None, cwd='/', reset_env=False):
148
  """Execute a (shell) command.
149

150
  The command should not read from its standard input, as it will be
151
  closed.
152

153
  @type cmd: string or list
154
  @param cmd: Command to run
155
  @type env: dict
156
  @param env: Additional environment
157
  @type output: str
158
  @param output: if desired, the output of the command can be
159
      saved in a file instead of the RunResult instance; this
160
      parameter denotes the file name (if not None)
161
  @type cwd: string
162
  @param cwd: if specified, will be used as the working
163
      directory for the command; the default will be /
164
  @type reset_env: boolean
165
  @param reset_env: whether to reset or keep the default os environment
166
  @rtype: L{RunResult}
167
  @return: RunResult instance
168
  @raise errors.ProgrammerError: if we call this when forks are disabled
169

170
  """
171
  if no_fork:
172
    raise errors.ProgrammerError("utils.RunCmd() called with fork() disabled")
173

    
174
  if isinstance(cmd, list):
175
    cmd = [str(val) for val in cmd]
176
    strcmd = " ".join(cmd)
177
    shell = False
178
  else:
179
    strcmd = cmd
180
    shell = True
181
  logging.debug("RunCmd '%s'", strcmd)
182

    
183
  if not reset_env:
184
    cmd_env = os.environ.copy()
185
    cmd_env["LC_ALL"] = "C"
186
  else:
187
    cmd_env = {}
188

    
189
  if env is not None:
190
    cmd_env.update(env)
191

    
192
  try:
193
    if output is None:
194
      out, err, status = _RunCmdPipe(cmd, cmd_env, shell, cwd)
195
    else:
196
      status = _RunCmdFile(cmd, cmd_env, shell, output, cwd)
197
      out = err = ""
198
  except OSError, err:
199
    if err.errno == errno.ENOENT:
200
      raise errors.OpExecError("Can't execute '%s': not found (%s)" %
201
                               (strcmd, err))
202
    else:
203
      raise
204

    
205
  if status >= 0:
206
    exitcode = status
207
    signal_ = None
208
  else:
209
    exitcode = None
210
    signal_ = -status
211

    
212
  return RunResult(exitcode, signal_, out, err, strcmd)
213

    
214

    
215
def _RunCmdPipe(cmd, env, via_shell, cwd):
216
  """Run a command and return its output.
217

218
  @type  cmd: string or list
219
  @param cmd: Command to run
220
  @type env: dict
221
  @param env: The environment to use
222
  @type via_shell: bool
223
  @param via_shell: if we should run via the shell
224
  @type cwd: string
225
  @param cwd: the working directory for the program
226
  @rtype: tuple
227
  @return: (out, err, status)
228

229
  """
230
  poller = select.poll()
231
  child = subprocess.Popen(cmd, shell=via_shell,
232
                           stderr=subprocess.PIPE,
233
                           stdout=subprocess.PIPE,
234
                           stdin=subprocess.PIPE,
235
                           close_fds=True, env=env,
236
                           cwd=cwd)
237

    
238
  child.stdin.close()
239
  poller.register(child.stdout, select.POLLIN)
240
  poller.register(child.stderr, select.POLLIN)
241
  out = StringIO()
242
  err = StringIO()
243
  fdmap = {
244
    child.stdout.fileno(): (out, child.stdout),
245
    child.stderr.fileno(): (err, child.stderr),
246
    }
247
  for fd in fdmap:
248
    status = fcntl.fcntl(fd, fcntl.F_GETFL)
249
    fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
250

    
251
  while fdmap:
252
    try:
253
      pollresult = poller.poll()
254
    except EnvironmentError, eerr:
255
      if eerr.errno == errno.EINTR:
256
        continue
257
      raise
258
    except select.error, serr:
259
      if serr[0] == errno.EINTR:
260
        continue
261
      raise
262

    
263
    for fd, event in pollresult:
264
      if event & select.POLLIN or event & select.POLLPRI:
265
        data = fdmap[fd][1].read()
266
        # no data from read signifies EOF (the same as POLLHUP)
267
        if not data:
268
          poller.unregister(fd)
269
          del fdmap[fd]
270
          continue
271
        fdmap[fd][0].write(data)
272
      if (event & select.POLLNVAL or event & select.POLLHUP or
273
          event & select.POLLERR):
274
        poller.unregister(fd)
275
        del fdmap[fd]
276

    
277
  out = out.getvalue()
278
  err = err.getvalue()
279

    
280
  status = child.wait()
281
  return out, err, status
282

    
283

    
284
def _RunCmdFile(cmd, env, via_shell, output, cwd):
285
  """Run a command and save its output to a file.
286

287
  @type  cmd: string or list
288
  @param cmd: Command to run
289
  @type env: dict
290
  @param env: The environment to use
291
  @type via_shell: bool
292
  @param via_shell: if we should run via the shell
293
  @type output: str
294
  @param output: the filename in which to save the output
295
  @type cwd: string
296
  @param cwd: the working directory for the program
297
  @rtype: int
298
  @return: the exit status
299

300
  """
301
  fh = open(output, "a")
302
  try:
303
    child = subprocess.Popen(cmd, shell=via_shell,
304
                             stderr=subprocess.STDOUT,
305
                             stdout=fh,
306
                             stdin=subprocess.PIPE,
307
                             close_fds=True, env=env,
308
                             cwd=cwd)
309

    
310
    child.stdin.close()
311
    status = child.wait()
312
  finally:
313
    fh.close()
314
  return status
315

    
316

    
317
def RunParts(dir_name, env=None, reset_env=False):
318
  """Run Scripts or programs in a directory
319

320
  @type dir_name: string
321
  @param dir_name: absolute path to a directory
322
  @type env: dict
323
  @param env: The environment to use
324
  @type reset_env: boolean
325
  @param reset_env: whether to reset or keep the default os environment
326
  @rtype: list of tuples
327
  @return: list of (name, (one of RUNDIR_STATUS), RunResult)
328

329
  """
330
  rr = []
331

    
332
  try:
333
    dir_contents = ListVisibleFiles(dir_name)
334
  except OSError, err:
335
    logging.warning("RunParts: skipping %s (cannot list: %s)", dir_name, err)
336
    return rr
337

    
338
  for relname in sorted(dir_contents):
339
    fname = PathJoin(dir_name, relname)
340
    if not (os.path.isfile(fname) and os.access(fname, os.X_OK) and
341
            constants.EXT_PLUGIN_MASK.match(relname) is not None):
342
      rr.append((relname, constants.RUNPARTS_SKIP, None))
343
    else:
344
      try:
345
        result = RunCmd([fname], env=env, reset_env=reset_env)
346
      except Exception, err: # pylint: disable-msg=W0703
347
        rr.append((relname, constants.RUNPARTS_ERR, str(err)))
348
      else:
349
        rr.append((relname, constants.RUNPARTS_RUN, result))
350

    
351
  return rr
352

    
353

    
354
def GetSocketCredentials(sock):
355
  """Returns the credentials of the foreign process connected to a socket.
356

357
  @param sock: Unix socket
358
  @rtype: tuple; (number, number, number)
359
  @return: The PID, UID and GID of the connected foreign process.
360

361
  """
362
  peercred = sock.getsockopt(socket.SOL_SOCKET, IN.SO_PEERCRED,
363
                             _STRUCT_UCRED_SIZE)
364
  return struct.unpack(_STRUCT_UCRED, peercred)
365

    
366

    
367
def RemoveFile(filename):
368
  """Remove a file ignoring some errors.
369

370
  Remove a file, ignoring non-existing ones or directories. Other
371
  errors are passed.
372

373
  @type filename: str
374
  @param filename: the file to be removed
375

376
  """
377
  try:
378
    os.unlink(filename)
379
  except OSError, err:
380
    if err.errno not in (errno.ENOENT, errno.EISDIR):
381
      raise
382

    
383

    
384
def RemoveDir(dirname):
385
  """Remove an empty directory.
386

387
  Remove a directory, ignoring non-existing ones.
388
  Other errors are passed. This includes the case,
389
  where the directory is not empty, so it can't be removed.
390

391
  @type dirname: str
392
  @param dirname: the empty directory to be removed
393

394
  """
395
  try:
396
    os.rmdir(dirname)
397
  except OSError, err:
398
    if err.errno != errno.ENOENT:
399
      raise
400

    
401

    
402
def RenameFile(old, new, mkdir=False, mkdir_mode=0750):
403
  """Renames a file.
404

405
  @type old: string
406
  @param old: Original path
407
  @type new: string
408
  @param new: New path
409
  @type mkdir: bool
410
  @param mkdir: Whether to create target directory if it doesn't exist
411
  @type mkdir_mode: int
412
  @param mkdir_mode: Mode for newly created directories
413

414
  """
415
  try:
416
    return os.rename(old, new)
417
  except OSError, err:
418
    # In at least one use case of this function, the job queue, directory
419
    # creation is very rare. Checking for the directory before renaming is not
420
    # as efficient.
421
    if mkdir and err.errno == errno.ENOENT:
422
      # Create directory and try again
423
      Makedirs(os.path.dirname(new), mode=mkdir_mode)
424

    
425
      return os.rename(old, new)
426

    
427
    raise
428

    
429

    
430
def Makedirs(path, mode=0750):
431
  """Super-mkdir; create a leaf directory and all intermediate ones.
432

433
  This is a wrapper around C{os.makedirs} adding error handling not implemented
434
  before Python 2.5.
435

436
  """
437
  try:
438
    os.makedirs(path, mode)
439
  except OSError, err:
440
    # Ignore EEXIST. This is only handled in os.makedirs as included in
441
    # Python 2.5 and above.
442
    if err.errno != errno.EEXIST or not os.path.exists(path):
443
      raise
444

    
445

    
446
def ResetTempfileModule():
447
  """Resets the random name generator of the tempfile module.
448

449
  This function should be called after C{os.fork} in the child process to
450
  ensure it creates a newly seeded random generator. Otherwise it would
451
  generate the same random parts as the parent process. If several processes
452
  race for the creation of a temporary file, this could lead to one not getting
453
  a temporary name.
454

455
  """
456
  # pylint: disable-msg=W0212
457
  if hasattr(tempfile, "_once_lock") and hasattr(tempfile, "_name_sequence"):
458
    tempfile._once_lock.acquire()
459
    try:
460
      # Reset random name generator
461
      tempfile._name_sequence = None
462
    finally:
463
      tempfile._once_lock.release()
464
  else:
465
    logging.critical("The tempfile module misses at least one of the"
466
                     " '_once_lock' and '_name_sequence' attributes")
467

    
468

    
469
def _FingerprintFile(filename):
470
  """Compute the fingerprint of a file.
471

472
  If the file does not exist, a None will be returned
473
  instead.
474

475
  @type filename: str
476
  @param filename: the filename to checksum
477
  @rtype: str
478
  @return: the hex digest of the sha checksum of the contents
479
      of the file
480

481
  """
482
  if not (os.path.exists(filename) and os.path.isfile(filename)):
483
    return None
484

    
485
  f = open(filename)
486

    
487
  fp = sha1()
488
  while True:
489
    data = f.read(4096)
490
    if not data:
491
      break
492

    
493
    fp.update(data)
494

    
495
  return fp.hexdigest()
496

    
497

    
498
def FingerprintFiles(files):
499
  """Compute fingerprints for a list of files.
500

501
  @type files: list
502
  @param files: the list of filename to fingerprint
503
  @rtype: dict
504
  @return: a dictionary filename: fingerprint, holding only
505
      existing files
506

507
  """
508
  ret = {}
509

    
510
  for filename in files:
511
    cksum = _FingerprintFile(filename)
512
    if cksum:
513
      ret[filename] = cksum
514

    
515
  return ret
516

    
517

    
518
def ForceDictType(target, key_types, allowed_values=None):
519
  """Force the values of a dict to have certain types.
520

521
  @type target: dict
522
  @param target: the dict to update
523
  @type key_types: dict
524
  @param key_types: dict mapping target dict keys to types
525
                    in constants.ENFORCEABLE_TYPES
526
  @type allowed_values: list
527
  @keyword allowed_values: list of specially allowed values
528

529
  """
530
  if allowed_values is None:
531
    allowed_values = []
532

    
533
  if not isinstance(target, dict):
534
    msg = "Expected dictionary, got '%s'" % target
535
    raise errors.TypeEnforcementError(msg)
536

    
537
  for key in target:
538
    if key not in key_types:
539
      msg = "Unknown key '%s'" % key
540
      raise errors.TypeEnforcementError(msg)
541

    
542
    if target[key] in allowed_values:
543
      continue
544

    
545
    ktype = key_types[key]
546
    if ktype not in constants.ENFORCEABLE_TYPES:
547
      msg = "'%s' has non-enforceable type %s" % (key, ktype)
548
      raise errors.ProgrammerError(msg)
549

    
550
    if ktype == constants.VTYPE_STRING:
551
      if not isinstance(target[key], basestring):
552
        if isinstance(target[key], bool) and not target[key]:
553
          target[key] = ''
554
        else:
555
          msg = "'%s' (value %s) is not a valid string" % (key, target[key])
556
          raise errors.TypeEnforcementError(msg)
557
    elif ktype == constants.VTYPE_BOOL:
558
      if isinstance(target[key], basestring) and target[key]:
559
        if target[key].lower() == constants.VALUE_FALSE:
560
          target[key] = False
561
        elif target[key].lower() == constants.VALUE_TRUE:
562
          target[key] = True
563
        else:
564
          msg = "'%s' (value %s) is not a valid boolean" % (key, target[key])
565
          raise errors.TypeEnforcementError(msg)
566
      elif target[key]:
567
        target[key] = True
568
      else:
569
        target[key] = False
570
    elif ktype == constants.VTYPE_SIZE:
571
      try:
572
        target[key] = ParseUnit(target[key])
573
      except errors.UnitParseError, err:
574
        msg = "'%s' (value %s) is not a valid size. error: %s" % \
575
              (key, target[key], err)
576
        raise errors.TypeEnforcementError(msg)
577
    elif ktype == constants.VTYPE_INT:
578
      try:
579
        target[key] = int(target[key])
580
      except (ValueError, TypeError):
581
        msg = "'%s' (value %s) is not a valid integer" % (key, target[key])
582
        raise errors.TypeEnforcementError(msg)
583

    
584

    
585
def IsProcessAlive(pid):
586
  """Check if a given pid exists on the system.
587

588
  @note: zombie status is not handled, so zombie processes
589
      will be returned as alive
590
  @type pid: int
591
  @param pid: the process ID to check
592
  @rtype: boolean
593
  @return: True if the process exists
594

595
  """
596
  def _TryStat(name):
597
    try:
598
      os.stat(name)
599
      return True
600
    except EnvironmentError, err:
601
      if err.errno in (errno.ENOENT, errno.ENOTDIR):
602
        return False
603
      elif err.errno == errno.EINVAL:
604
        raise RetryAgain(err)
605
      raise
606

    
607
  assert isinstance(pid, int), "pid must be an integer"
608
  if pid <= 0:
609
    return False
610

    
611
  proc_entry = "/proc/%d/status" % pid
612
  # /proc in a multiprocessor environment can have strange behaviors.
613
  # Retry the os.stat a few times until we get a good result.
614
  try:
615
    return Retry(_TryStat, (0.01, 1.5, 0.1), 0.5, args=[proc_entry])
616
  except RetryTimeout, err:
617
    err.RaiseInner()
618

    
619

    
620
def ReadPidFile(pidfile):
621
  """Read a pid from a file.
622

623
  @type  pidfile: string
624
  @param pidfile: path to the file containing the pid
625
  @rtype: int
626
  @return: The process id, if the file exists and contains a valid PID,
627
           otherwise 0
628

629
  """
630
  try:
631
    raw_data = ReadOneLineFile(pidfile)
632
  except EnvironmentError, err:
633
    if err.errno != errno.ENOENT:
634
      logging.exception("Can't read pid file")
635
    return 0
636

    
637
  try:
638
    pid = int(raw_data)
639
  except (TypeError, ValueError), err:
640
    logging.info("Can't parse pid file contents", exc_info=True)
641
    return 0
642

    
643
  return pid
644

    
645

    
646
def MatchNameComponent(key, name_list, case_sensitive=True):
647
  """Try to match a name against a list.
648

649
  This function will try to match a name like test1 against a list
650
  like C{['test1.example.com', 'test2.example.com', ...]}. Against
651
  this list, I{'test1'} as well as I{'test1.example'} will match, but
652
  not I{'test1.ex'}. A multiple match will be considered as no match
653
  at all (e.g. I{'test1'} against C{['test1.example.com',
654
  'test1.example.org']}), except when the key fully matches an entry
655
  (e.g. I{'test1'} against C{['test1', 'test1.example.com']}).
656

657
  @type key: str
658
  @param key: the name to be searched
659
  @type name_list: list
660
  @param name_list: the list of strings against which to search the key
661
  @type case_sensitive: boolean
662
  @param case_sensitive: whether to provide a case-sensitive match
663

664
  @rtype: None or str
665
  @return: None if there is no match I{or} if there are multiple matches,
666
      otherwise the element from the list which matches
667

668
  """
669
  if key in name_list:
670
    return key
671

    
672
  re_flags = 0
673
  if not case_sensitive:
674
    re_flags |= re.IGNORECASE
675
    key = key.upper()
676
  mo = re.compile("^%s(\..*)?$" % re.escape(key), re_flags)
677
  names_filtered = []
678
  string_matches = []
679
  for name in name_list:
680
    if mo.match(name) is not None:
681
      names_filtered.append(name)
682
      if not case_sensitive and key == name.upper():
683
        string_matches.append(name)
684

    
685
  if len(string_matches) == 1:
686
    return string_matches[0]
687
  if len(names_filtered) == 1:
688
    return names_filtered[0]
689
  return None
690

    
691

    
692
class HostInfo:
693
  """Class implementing resolver and hostname functionality
694

695
  """
696
  _VALID_NAME_RE = re.compile("^[a-z0-9._-]{1,255}$")
697

    
698
  def __init__(self, name=None):
699
    """Initialize the host name object.
700

701
    If the name argument is not passed, it will use this system's
702
    name.
703

704
    """
705
    if name is None:
706
      name = self.SysName()
707

    
708
    self.query = name
709
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
710
    self.ip = self.ipaddrs[0]
711

    
712
  def ShortName(self):
713
    """Returns the hostname without domain.
714

715
    """
716
    return self.name.split('.')[0]
717

    
718
  @staticmethod
719
  def SysName():
720
    """Return the current system's name.
721

722
    This is simply a wrapper over C{socket.gethostname()}.
723

724
    """
725
    return socket.gethostname()
726

    
727
  @staticmethod
728
  def LookupHostname(hostname):
729
    """Look up hostname
730

731
    @type hostname: str
732
    @param hostname: hostname to look up
733

734
    @rtype: tuple
735
    @return: a tuple (name, aliases, ipaddrs) as returned by
736
        C{socket.gethostbyname_ex}
737
    @raise errors.ResolverError: in case of errors in resolving
738

739
    """
740
    try:
741
      result = socket.gethostbyname_ex(hostname)
742
    except (socket.gaierror, socket.herror, socket.error), err:
743
      # hostname not found in DNS, or other socket exception in the
744
      # (code, description format)
745
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
746

    
747
    return result
748

    
749
  @classmethod
750
  def NormalizeName(cls, hostname):
751
    """Validate and normalize the given hostname.
752

753
    @attention: the validation is a bit more relaxed than the standards
754
        require; most importantly, we allow underscores in names
755
    @raise errors.OpPrereqError: when the name is not valid
756

757
    """
758
    hostname = hostname.lower()
759
    if (not cls._VALID_NAME_RE.match(hostname) or
760
        # double-dots, meaning empty label
761
        ".." in hostname or
762
        # empty initial label
763
        hostname.startswith(".")):
764
      raise errors.OpPrereqError("Invalid hostname '%s'" % hostname,
765
                                 errors.ECODE_INVAL)
766
    if hostname.endswith("."):
767
      hostname = hostname.rstrip(".")
768
    return hostname
769

    
770

    
771
def GetHostInfo(name=None):
772
  """Lookup host name and raise an OpPrereqError for failures"""
773

    
774
  try:
775
    return HostInfo(name)
776
  except errors.ResolverError, err:
777
    raise errors.OpPrereqError("The given name (%s) does not resolve: %s" %
778
                               (err[0], err[2]), errors.ECODE_RESOLVER)
779

    
780

    
781
def ListVolumeGroups():
782
  """List volume groups and their size
783

784
  @rtype: dict
785
  @return:
786
       Dictionary with keys volume name and values
787
       the size of the volume
788

789
  """
790
  command = "vgs --noheadings --units m --nosuffix -o name,size"
791
  result = RunCmd(command)
792
  retval = {}
793
  if result.failed:
794
    return retval
795

    
796
  for line in result.stdout.splitlines():
797
    try:
798
      name, size = line.split()
799
      size = int(float(size))
800
    except (IndexError, ValueError), err:
801
      logging.error("Invalid output from vgs (%s): %s", err, line)
802
      continue
803

    
804
    retval[name] = size
805

    
806
  return retval
807

    
808

    
809
def BridgeExists(bridge):
810
  """Check whether the given bridge exists in the system
811

812
  @type bridge: str
813
  @param bridge: the bridge name to check
814
  @rtype: boolean
815
  @return: True if it does
816

817
  """
818
  return os.path.isdir("/sys/class/net/%s/bridge" % bridge)
819

    
820

    
821
def NiceSort(name_list):
822
  """Sort a list of strings based on digit and non-digit groupings.
823

824
  Given a list of names C{['a1', 'a10', 'a11', 'a2']} this function
825
  will sort the list in the logical order C{['a1', 'a2', 'a10',
826
  'a11']}.
827

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

832
  @type name_list: list
833
  @param name_list: the names to be sorted
834
  @rtype: list
835
  @return: a copy of the name list sorted with our algorithm
836

837
  """
838
  _SORTER_BASE = "(\D+|\d+)"
839
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
840
                                                  _SORTER_BASE, _SORTER_BASE,
841
                                                  _SORTER_BASE, _SORTER_BASE,
842
                                                  _SORTER_BASE, _SORTER_BASE)
843
  _SORTER_RE = re.compile(_SORTER_FULL)
844
  _SORTER_NODIGIT = re.compile("^\D*$")
845
  def _TryInt(val):
846
    """Attempts to convert a variable to integer."""
847
    if val is None or _SORTER_NODIGIT.match(val):
848
      return val
849
    rval = int(val)
850
    return rval
851

    
852
  to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
853
             for name in name_list]
854
  to_sort.sort()
855
  return [tup[1] for tup in to_sort]
856

    
857

    
858
def TryConvert(fn, val):
859
  """Try to convert a value ignoring errors.
860

861
  This function tries to apply function I{fn} to I{val}. If no
862
  C{ValueError} or C{TypeError} exceptions are raised, it will return
863
  the result, else it will return the original value. Any other
864
  exceptions are propagated to the caller.
865

866
  @type fn: callable
867
  @param fn: function to apply to the value
868
  @param val: the value to be converted
869
  @return: The converted value if the conversion was successful,
870
      otherwise the original value.
871

872
  """
873
  try:
874
    nv = fn(val)
875
  except (ValueError, TypeError):
876
    nv = val
877
  return nv
878

    
879

    
880
def IsValidIP(ip):
881
  """Verifies the syntax of an IPv4 address.
882

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

886
  @type ip: str
887
  @param ip: the address to be checked
888
  @rtype: a regular expression match object
889
  @return: a regular expression match object, or None if the
890
      address is not valid
891

892
  """
893
  unit = "(0|[1-9]\d{0,2})"
894
  #TODO: convert and return only boolean
895
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
896

    
897

    
898
def IsValidShellParam(word):
899
  """Verifies is the given word is safe from the shell's p.o.v.
900

901
  This means that we can pass this to a command via the shell and be
902
  sure that it doesn't alter the command line and is passed as such to
903
  the actual command.
904

905
  Note that we are overly restrictive here, in order to be on the safe
906
  side.
907

908
  @type word: str
909
  @param word: the word to check
910
  @rtype: boolean
911
  @return: True if the word is 'safe'
912

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

    
916

    
917
def BuildShellCmd(template, *args):
918
  """Build a safe shell command line from the given arguments.
919

920
  This function will check all arguments in the args list so that they
921
  are valid shell parameters (i.e. they don't contain shell
922
  metacharacters). If everything is ok, it will return the result of
923
  template % args.
924

925
  @type template: str
926
  @param template: the string holding the template for the
927
      string formatting
928
  @rtype: str
929
  @return: the expanded command line
930

931
  """
932
  for word in args:
933
    if not IsValidShellParam(word):
934
      raise errors.ProgrammerError("Shell argument '%s' contains"
935
                                   " invalid characters" % word)
936
  return template % args
937

    
938

    
939
def FormatUnit(value, units):
940
  """Formats an incoming number of MiB with the appropriate unit.
941

942
  @type value: int
943
  @param value: integer representing the value in MiB (1048576)
944
  @type units: char
945
  @param units: the type of formatting we should do:
946
      - 'h' for automatic scaling
947
      - 'm' for MiBs
948
      - 'g' for GiBs
949
      - 't' for TiBs
950
  @rtype: str
951
  @return: the formatted value (with suffix)
952

953
  """
954
  if units not in ('m', 'g', 't', 'h'):
955
    raise errors.ProgrammerError("Invalid unit specified '%s'" % str(units))
956

    
957
  suffix = ''
958

    
959
  if units == 'm' or (units == 'h' and value < 1024):
960
    if units == 'h':
961
      suffix = 'M'
962
    return "%d%s" % (round(value, 0), suffix)
963

    
964
  elif units == 'g' or (units == 'h' and value < (1024 * 1024)):
965
    if units == 'h':
966
      suffix = 'G'
967
    return "%0.1f%s" % (round(float(value) / 1024, 1), suffix)
968

    
969
  else:
970
    if units == 'h':
971
      suffix = 'T'
972
    return "%0.1f%s" % (round(float(value) / 1024 / 1024, 1), suffix)
973

    
974

    
975
def ParseUnit(input_string):
976
  """Tries to extract number and scale from the given string.
977

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

982
  """
983
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', str(input_string))
984
  if not m:
985
    raise errors.UnitParseError("Invalid format")
986

    
987
  value = float(m.groups()[0])
988

    
989
  unit = m.groups()[1]
990
  if unit:
991
    lcunit = unit.lower()
992
  else:
993
    lcunit = 'm'
994

    
995
  if lcunit in ('m', 'mb', 'mib'):
996
    # Value already in MiB
997
    pass
998

    
999
  elif lcunit in ('g', 'gb', 'gib'):
1000
    value *= 1024
1001

    
1002
  elif lcunit in ('t', 'tb', 'tib'):
1003
    value *= 1024 * 1024
1004

    
1005
  else:
1006
    raise errors.UnitParseError("Unknown unit: %s" % unit)
1007

    
1008
  # Make sure we round up
1009
  if int(value) < value:
1010
    value += 1
1011

    
1012
  # Round up to the next multiple of 4
1013
  value = int(value)
1014
  if value % 4:
1015
    value += 4 - value % 4
1016

    
1017
  return value
1018

    
1019

    
1020
def AddAuthorizedKey(file_name, key):
1021
  """Adds an SSH public key to an authorized_keys file.
1022

1023
  @type file_name: str
1024
  @param file_name: path to authorized_keys file
1025
  @type key: str
1026
  @param key: string containing key
1027

1028
  """
1029
  key_fields = key.split()
1030

    
1031
  f = open(file_name, 'a+')
1032
  try:
1033
    nl = True
1034
    for line in f:
1035
      # Ignore whitespace changes
1036
      if line.split() == key_fields:
1037
        break
1038
      nl = line.endswith('\n')
1039
    else:
1040
      if not nl:
1041
        f.write("\n")
1042
      f.write(key.rstrip('\r\n'))
1043
      f.write("\n")
1044
      f.flush()
1045
  finally:
1046
    f.close()
1047

    
1048

    
1049
def RemoveAuthorizedKey(file_name, key):
1050
  """Removes an SSH public key from an authorized_keys file.
1051

1052
  @type file_name: str
1053
  @param file_name: path to authorized_keys file
1054
  @type key: str
1055
  @param key: string containing key
1056

1057
  """
1058
  key_fields = key.split()
1059

    
1060
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
1061
  try:
1062
    out = os.fdopen(fd, 'w')
1063
    try:
1064
      f = open(file_name, 'r')
1065
      try:
1066
        for line in f:
1067
          # Ignore whitespace changes while comparing lines
1068
          if line.split() != key_fields:
1069
            out.write(line)
1070

    
1071
        out.flush()
1072
        os.rename(tmpname, file_name)
1073
      finally:
1074
        f.close()
1075
    finally:
1076
      out.close()
1077
  except:
1078
    RemoveFile(tmpname)
1079
    raise
1080

    
1081

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

1085
  @type file_name: str
1086
  @param file_name: path to the file to modify (usually C{/etc/hosts})
1087
  @type ip: str
1088
  @param ip: the IP address
1089
  @type hostname: str
1090
  @param hostname: the hostname to be added
1091
  @type aliases: list
1092
  @param aliases: the list of aliases to add for the hostname
1093

1094
  """
1095
  # FIXME: use WriteFile + fn rather than duplicating its efforts
1096
  # Ensure aliases are unique
1097
  aliases = UniqueSequence([hostname] + aliases)[1:]
1098

    
1099
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
1100
  try:
1101
    out = os.fdopen(fd, 'w')
1102
    try:
1103
      f = open(file_name, 'r')
1104
      try:
1105
        for line in f:
1106
          fields = line.split()
1107
          if fields and not fields[0].startswith('#') and ip == fields[0]:
1108
            continue
1109
          out.write(line)
1110

    
1111
        out.write("%s\t%s" % (ip, hostname))
1112
        if aliases:
1113
          out.write(" %s" % ' '.join(aliases))
1114
        out.write('\n')
1115

    
1116
        out.flush()
1117
        os.fsync(out)
1118
        os.chmod(tmpname, 0644)
1119
        os.rename(tmpname, file_name)
1120
      finally:
1121
        f.close()
1122
    finally:
1123
      out.close()
1124
  except:
1125
    RemoveFile(tmpname)
1126
    raise
1127

    
1128

    
1129
def AddHostToEtcHosts(hostname):
1130
  """Wrapper around SetEtcHostsEntry.
1131

1132
  @type hostname: str
1133
  @param hostname: a hostname that will be resolved and added to
1134
      L{constants.ETC_HOSTS}
1135

1136
  """
1137
  hi = HostInfo(name=hostname)
1138
  SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])
1139

    
1140

    
1141
def RemoveEtcHostsEntry(file_name, hostname):
1142
  """Removes a hostname from /etc/hosts.
1143

1144
  IP addresses without names are removed from the file.
1145

1146
  @type file_name: str
1147
  @param file_name: path to the file to modify (usually C{/etc/hosts})
1148
  @type hostname: str
1149
  @param hostname: the hostname to be removed
1150

1151
  """
1152
  # FIXME: use WriteFile + fn rather than duplicating its efforts
1153
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
1154
  try:
1155
    out = os.fdopen(fd, 'w')
1156
    try:
1157
      f = open(file_name, 'r')
1158
      try:
1159
        for line in f:
1160
          fields = line.split()
1161
          if len(fields) > 1 and not fields[0].startswith('#'):
1162
            names = fields[1:]
1163
            if hostname in names:
1164
              while hostname in names:
1165
                names.remove(hostname)
1166
              if names:
1167
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
1168
              continue
1169

    
1170
          out.write(line)
1171

    
1172
        out.flush()
1173
        os.fsync(out)
1174
        os.chmod(tmpname, 0644)
1175
        os.rename(tmpname, file_name)
1176
      finally:
1177
        f.close()
1178
    finally:
1179
      out.close()
1180
  except:
1181
    RemoveFile(tmpname)
1182
    raise
1183

    
1184

    
1185
def RemoveHostFromEtcHosts(hostname):
1186
  """Wrapper around RemoveEtcHostsEntry.
1187

1188
  @type hostname: str
1189
  @param hostname: hostname that will be resolved and its
1190
      full and shot name will be removed from
1191
      L{constants.ETC_HOSTS}
1192

1193
  """
1194
  hi = HostInfo(name=hostname)
1195
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.name)
1196
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.ShortName())
1197

    
1198

    
1199
def TimestampForFilename():
1200
  """Returns the current time formatted for filenames.
1201

1202
  The format doesn't contain colons as some shells and applications them as
1203
  separators.
1204

1205
  """
1206
  return time.strftime("%Y-%m-%d_%H_%M_%S")
1207

    
1208

    
1209
def CreateBackup(file_name):
1210
  """Creates a backup of a file.
1211

1212
  @type file_name: str
1213
  @param file_name: file to be backed up
1214
  @rtype: str
1215
  @return: the path to the newly created backup
1216
  @raise errors.ProgrammerError: for invalid file names
1217

1218
  """
1219
  if not os.path.isfile(file_name):
1220
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
1221
                                file_name)
1222

    
1223
  prefix = ("%s.backup-%s." %
1224
            (os.path.basename(file_name), TimestampForFilename()))
1225
  dir_name = os.path.dirname(file_name)
1226

    
1227
  fsrc = open(file_name, 'rb')
1228
  try:
1229
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
1230
    fdst = os.fdopen(fd, 'wb')
1231
    try:
1232
      logging.debug("Backing up %s at %s", file_name, backup_name)
1233
      shutil.copyfileobj(fsrc, fdst)
1234
    finally:
1235
      fdst.close()
1236
  finally:
1237
    fsrc.close()
1238

    
1239
  return backup_name
1240

    
1241

    
1242
def ShellQuote(value):
1243
  """Quotes shell argument according to POSIX.
1244

1245
  @type value: str
1246
  @param value: the argument to be quoted
1247
  @rtype: str
1248
  @return: the quoted value
1249

1250
  """
1251
  if _re_shell_unquoted.match(value):
1252
    return value
1253
  else:
1254
    return "'%s'" % value.replace("'", "'\\''")
1255

    
1256

    
1257
def ShellQuoteArgs(args):
1258
  """Quotes a list of shell arguments.
1259

1260
  @type args: list
1261
  @param args: list of arguments to be quoted
1262
  @rtype: str
1263
  @return: the quoted arguments concatenated with spaces
1264

1265
  """
1266
  return ' '.join([ShellQuote(i) for i in args])
1267

    
1268

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

1272
  Check if the given IP is reachable by doing attempting a TCP connect
1273
  to it.
1274

1275
  @type target: str
1276
  @param target: the IP or hostname to ping
1277
  @type port: int
1278
  @param port: the port to connect to
1279
  @type timeout: int
1280
  @param timeout: the timeout on the connection attempt
1281
  @type live_port_needed: boolean
1282
  @param live_port_needed: whether a closed port will cause the
1283
      function to return failure, as if there was a timeout
1284
  @type source: str or None
1285
  @param source: if specified, will cause the connect to be made
1286
      from this specific source address; failures to bind other
1287
      than C{EADDRNOTAVAIL} will be ignored
1288

1289
  """
1290
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
1291

    
1292
  success = False
1293

    
1294
  if source is not None:
1295
    try:
1296
      sock.bind((source, 0))
1297
    except socket.error, (errcode, _):
1298
      if errcode == errno.EADDRNOTAVAIL:
1299
        success = False
1300

    
1301
  sock.settimeout(timeout)
1302

    
1303
  try:
1304
    sock.connect((target, port))
1305
    sock.close()
1306
    success = True
1307
  except socket.timeout:
1308
    success = False
1309
  except socket.error, (errcode, _):
1310
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
1311

    
1312
  return success
1313

    
1314

    
1315
def OwnIpAddress(address):
1316
  """Check if the current host has the the given IP address.
1317

1318
  Currently this is done by TCP-pinging the address from the loopback
1319
  address.
1320

1321
  @type address: string
1322
  @param address: the address to check
1323
  @rtype: bool
1324
  @return: True if we own the address
1325

1326
  """
1327
  return TcpPing(address, constants.DEFAULT_NODED_PORT,
1328
                 source=constants.LOCALHOST_IP_ADDRESS)
1329

    
1330

    
1331
def ListVisibleFiles(path):
1332
  """Returns a list of visible files in a directory.
1333

1334
  @type path: str
1335
  @param path: the directory to enumerate
1336
  @rtype: list
1337
  @return: the list of all files not starting with a dot
1338
  @raise ProgrammerError: if L{path} is not an absolue and normalized path
1339

1340
  """
1341
  if not IsNormAbsPath(path):
1342
    raise errors.ProgrammerError("Path passed to ListVisibleFiles is not"
1343
                                 " absolute/normalized: '%s'" % path)
1344
  files = [i for i in os.listdir(path) if not i.startswith(".")]
1345
  files.sort()
1346
  return files
1347

    
1348

    
1349
def GetHomeDir(user, default=None):
1350
  """Try to get the homedir of the given user.
1351

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

1356
  """
1357
  try:
1358
    if isinstance(user, basestring):
1359
      result = pwd.getpwnam(user)
1360
    elif isinstance(user, (int, long)):
1361
      result = pwd.getpwuid(user)
1362
    else:
1363
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
1364
                                   type(user))
1365
  except KeyError:
1366
    return default
1367
  return result.pw_dir
1368

    
1369

    
1370
def NewUUID():
1371
  """Returns a random UUID.
1372

1373
  @note: This is a Linux-specific method as it uses the /proc
1374
      filesystem.
1375
  @rtype: str
1376

1377
  """
1378
  return ReadFile(_RANDOM_UUID_FILE, size=128).rstrip("\n")
1379

    
1380

    
1381
def GenerateSecret(numbytes=20):
1382
  """Generates a random secret.
1383

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

1387
  @param numbytes: the number of bytes which will be represented by the returned
1388
      string (defaulting to 20, the length of a SHA1 hash)
1389
  @rtype: str
1390
  @return: an hex representation of the pseudo-random sequence
1391

1392
  """
1393
  return os.urandom(numbytes).encode('hex')
1394

    
1395

    
1396
def EnsureDirs(dirs):
1397
  """Make required directories, if they don't exist.
1398

1399
  @param dirs: list of tuples (dir_name, dir_mode)
1400
  @type dirs: list of (string, integer)
1401

1402
  """
1403
  for dir_name, dir_mode in dirs:
1404
    try:
1405
      os.mkdir(dir_name, dir_mode)
1406
    except EnvironmentError, err:
1407
      if err.errno != errno.EEXIST:
1408
        raise errors.GenericError("Cannot create needed directory"
1409
                                  " '%s': %s" % (dir_name, err))
1410
    try:
1411
      os.chmod(dir_name, dir_mode)
1412
    except EnvironmentError, err:
1413
      raise errors.GenericError("Cannot change directory permissions on"
1414
                                " '%s': %s" % (dir_name, err))
1415
    if not os.path.isdir(dir_name):
1416
      raise errors.GenericError("%s is not a directory" % dir_name)
1417

    
1418

    
1419
def ReadFile(file_name, size=-1):
1420
  """Reads a file.
1421

1422
  @type size: int
1423
  @param size: Read at most size bytes (if negative, entire file)
1424
  @rtype: str
1425
  @return: the (possibly partial) content of the file
1426

1427
  """
1428
  f = open(file_name, "r")
1429
  try:
1430
    return f.read(size)
1431
  finally:
1432
    f.close()
1433

    
1434

    
1435
def WriteFile(file_name, fn=None, data=None,
1436
              mode=None, uid=-1, gid=-1,
1437
              atime=None, mtime=None, close=True,
1438
              dry_run=False, backup=False,
1439
              prewrite=None, postwrite=None):
1440
  """(Over)write a file atomically.
1441

1442
  The file_name and either fn (a function taking one argument, the
1443
  file descriptor, and which should write the data to it) or data (the
1444
  contents of the file) must be passed. The other arguments are
1445
  optional and allow setting the file mode, owner and group, and the
1446
  mtime/atime of the file.
1447

1448
  If the function doesn't raise an exception, it has succeeded and the
1449
  target file has the new contents. If the function has raised an
1450
  exception, an existing target file should be unmodified and the
1451
  temporary file should be removed.
1452

1453
  @type file_name: str
1454
  @param file_name: the target filename
1455
  @type fn: callable
1456
  @param fn: content writing function, called with
1457
      file descriptor as parameter
1458
  @type data: str
1459
  @param data: contents of the file
1460
  @type mode: int
1461
  @param mode: file mode
1462
  @type uid: int
1463
  @param uid: the owner of the file
1464
  @type gid: int
1465
  @param gid: the group of the file
1466
  @type atime: int
1467
  @param atime: a custom access time to be set on the file
1468
  @type mtime: int
1469
  @param mtime: a custom modification time to be set on the file
1470
  @type close: boolean
1471
  @param close: whether to close file after writing it
1472
  @type prewrite: callable
1473
  @param prewrite: function to be called before writing content
1474
  @type postwrite: callable
1475
  @param postwrite: function to be called after writing content
1476

1477
  @rtype: None or int
1478
  @return: None if the 'close' parameter evaluates to True,
1479
      otherwise the file descriptor
1480

1481
  @raise errors.ProgrammerError: if any of the arguments are not valid
1482

1483
  """
1484
  if not os.path.isabs(file_name):
1485
    raise errors.ProgrammerError("Path passed to WriteFile is not"
1486
                                 " absolute: '%s'" % file_name)
1487

    
1488
  if [fn, data].count(None) != 1:
1489
    raise errors.ProgrammerError("fn or data required")
1490

    
1491
  if [atime, mtime].count(None) == 1:
1492
    raise errors.ProgrammerError("Both atime and mtime must be either"
1493
                                 " set or None")
1494

    
1495
  if backup and not dry_run and os.path.isfile(file_name):
1496
    CreateBackup(file_name)
1497

    
1498
  dir_name, base_name = os.path.split(file_name)
1499
  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
1500
  do_remove = True
1501
  # here we need to make sure we remove the temp file, if any error
1502
  # leaves it in place
1503
  try:
1504
    if uid != -1 or gid != -1:
1505
      os.chown(new_name, uid, gid)
1506
    if mode:
1507
      os.chmod(new_name, mode)
1508
    if callable(prewrite):
1509
      prewrite(fd)
1510
    if data is not None:
1511
      os.write(fd, data)
1512
    else:
1513
      fn(fd)
1514
    if callable(postwrite):
1515
      postwrite(fd)
1516
    os.fsync(fd)
1517
    if atime is not None and mtime is not None:
1518
      os.utime(new_name, (atime, mtime))
1519
    if not dry_run:
1520
      os.rename(new_name, file_name)
1521
      do_remove = False
1522
  finally:
1523
    if close:
1524
      os.close(fd)
1525
      result = None
1526
    else:
1527
      result = fd
1528
    if do_remove:
1529
      RemoveFile(new_name)
1530

    
1531
  return result
1532

    
1533

    
1534
def ReadOneLineFile(file_name, strict=False):
1535
  """Return the first non-empty line from a file.
1536

1537
  @type strict: boolean
1538
  @param strict: if True, abort if the file has more than one
1539
      non-empty line
1540

1541
  """
1542
  file_lines = ReadFile(file_name).splitlines()
1543
  full_lines = filter(bool, file_lines)
1544
  if not file_lines or not full_lines:
1545
    raise errors.GenericError("No data in one-liner file %s" % file_name)
1546
  elif strict and len(full_lines) > 1:
1547
    raise errors.GenericError("Too many lines in one-liner file %s" %
1548
                              file_name)
1549
  return full_lines[0]
1550

    
1551

    
1552
def FirstFree(seq, base=0):
1553
  """Returns the first non-existing integer from seq.
1554

1555
  The seq argument should be a sorted list of positive integers. The
1556
  first time the index of an element is smaller than the element
1557
  value, the index will be returned.
1558

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

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

1564
  @type seq: sequence
1565
  @param seq: the sequence to be analyzed.
1566
  @type base: int
1567
  @param base: use this value as the base index of the sequence
1568
  @rtype: int
1569
  @return: the first non-used index in the sequence
1570

1571
  """
1572
  for idx, elem in enumerate(seq):
1573
    assert elem >= base, "Passed element is higher than base offset"
1574
    if elem > idx + base:
1575
      # idx is not used
1576
      return idx + base
1577
  return None
1578

    
1579

    
1580
def SingleWaitForFdCondition(fdobj, event, timeout):
1581
  """Waits for a condition to occur on the socket.
1582

1583
  Immediately returns at the first interruption.
1584

1585
  @type fdobj: integer or object supporting a fileno() method
1586
  @param fdobj: entity to wait for events on
1587
  @type event: integer
1588
  @param event: ORed condition (see select module)
1589
  @type timeout: float or None
1590
  @param timeout: Timeout in seconds
1591
  @rtype: int or None
1592
  @return: None for timeout, otherwise occured conditions
1593

1594
  """
1595
  check = (event | select.POLLPRI |
1596
           select.POLLNVAL | select.POLLHUP | select.POLLERR)
1597

    
1598
  if timeout is not None:
1599
    # Poller object expects milliseconds
1600
    timeout *= 1000
1601

    
1602
  poller = select.poll()
1603
  poller.register(fdobj, event)
1604
  try:
1605
    # TODO: If the main thread receives a signal and we have no timeout, we
1606
    # could wait forever. This should check a global "quit" flag or something
1607
    # every so often.
1608
    io_events = poller.poll(timeout)
1609
  except select.error, err:
1610
    if err[0] != errno.EINTR:
1611
      raise
1612
    io_events = []
1613
  if io_events and io_events[0][1] & check:
1614
    return io_events[0][1]
1615
  else:
1616
    return None
1617

    
1618

    
1619
class FdConditionWaiterHelper(object):
1620
  """Retry helper for WaitForFdCondition.
1621

1622
  This class contains the retried and wait functions that make sure
1623
  WaitForFdCondition can continue waiting until the timeout is actually
1624
  expired.
1625

1626
  """
1627

    
1628
  def __init__(self, timeout):
1629
    self.timeout = timeout
1630

    
1631
  def Poll(self, fdobj, event):
1632
    result = SingleWaitForFdCondition(fdobj, event, self.timeout)
1633
    if result is None:
1634
      raise RetryAgain()
1635
    else:
1636
      return result
1637

    
1638
  def UpdateTimeout(self, timeout):
1639
    self.timeout = timeout
1640

    
1641

    
1642
def WaitForFdCondition(fdobj, event, timeout):
1643
  """Waits for a condition to occur on the socket.
1644

1645
  Retries until the timeout is expired, even if interrupted.
1646

1647
  @type fdobj: integer or object supporting a fileno() method
1648
  @param fdobj: entity to wait for events on
1649
  @type event: integer
1650
  @param event: ORed condition (see select module)
1651
  @type timeout: float or None
1652
  @param timeout: Timeout in seconds
1653
  @rtype: int or None
1654
  @return: None for timeout, otherwise occured conditions
1655

1656
  """
1657
  if timeout is not None:
1658
    retrywaiter = FdConditionWaiterHelper(timeout)
1659
    try:
1660
      result = Retry(retrywaiter.Poll, RETRY_REMAINING_TIME, timeout,
1661
                     args=(fdobj, event), wait_fn=retrywaiter.UpdateTimeout)
1662
    except RetryTimeout:
1663
      result = None
1664
  else:
1665
    result = None
1666
    while result is None:
1667
      result = SingleWaitForFdCondition(fdobj, event, timeout)
1668
  return result
1669

    
1670

    
1671
def UniqueSequence(seq):
1672
  """Returns a list with unique elements.
1673

1674
  Element order is preserved.
1675

1676
  @type seq: sequence
1677
  @param seq: the sequence with the source elements
1678
  @rtype: list
1679
  @return: list of unique elements from seq
1680

1681
  """
1682
  seen = set()
1683
  return [i for i in seq if i not in seen and not seen.add(i)]
1684

    
1685

    
1686
def NormalizeAndValidateMac(mac):
1687
  """Normalizes and check if a MAC address is valid.
1688

1689
  Checks whether the supplied MAC address is formally correct, only
1690
  accepts colon separated format. Normalize it to all lower.
1691

1692
  @type mac: str
1693
  @param mac: the MAC to be validated
1694
  @rtype: str
1695
  @return: returns the normalized and validated MAC.
1696

1697
  @raise errors.OpPrereqError: If the MAC isn't valid
1698

1699
  """
1700
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$", re.I)
1701
  if not mac_check.match(mac):
1702
    raise errors.OpPrereqError("Invalid MAC address specified: %s" %
1703
                               mac, errors.ECODE_INVAL)
1704

    
1705
  return mac.lower()
1706

    
1707

    
1708
def TestDelay(duration):
1709
  """Sleep for a fixed amount of time.
1710

1711
  @type duration: float
1712
  @param duration: the sleep duration
1713
  @rtype: boolean
1714
  @return: False for negative value, True otherwise
1715

1716
  """
1717
  if duration < 0:
1718
    return False, "Invalid sleep duration"
1719
  time.sleep(duration)
1720
  return True, None
1721

    
1722

    
1723
def _CloseFDNoErr(fd, retries=5):
1724
  """Close a file descriptor ignoring errors.
1725

1726
  @type fd: int
1727
  @param fd: the file descriptor
1728
  @type retries: int
1729
  @param retries: how many retries to make, in case we get any
1730
      other error than EBADF
1731

1732
  """
1733
  try:
1734
    os.close(fd)
1735
  except OSError, err:
1736
    if err.errno != errno.EBADF:
1737
      if retries > 0:
1738
        _CloseFDNoErr(fd, retries - 1)
1739
    # else either it's closed already or we're out of retries, so we
1740
    # ignore this and go on
1741

    
1742

    
1743
def CloseFDs(noclose_fds=None):
1744
  """Close file descriptors.
1745

1746
  This closes all file descriptors above 2 (i.e. except
1747
  stdin/out/err).
1748

1749
  @type noclose_fds: list or None
1750
  @param noclose_fds: if given, it denotes a list of file descriptor
1751
      that should not be closed
1752

1753
  """
1754
  # Default maximum for the number of available file descriptors.
1755
  if 'SC_OPEN_MAX' in os.sysconf_names:
1756
    try:
1757
      MAXFD = os.sysconf('SC_OPEN_MAX')
1758
      if MAXFD < 0:
1759
        MAXFD = 1024
1760
    except OSError:
1761
      MAXFD = 1024
1762
  else:
1763
    MAXFD = 1024
1764
  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1765
  if (maxfd == resource.RLIM_INFINITY):
1766
    maxfd = MAXFD
1767

    
1768
  # Iterate through and close all file descriptors (except the standard ones)
1769
  for fd in range(3, maxfd):
1770
    if noclose_fds and fd in noclose_fds:
1771
      continue
1772
    _CloseFDNoErr(fd)
1773

    
1774

    
1775
def Mlockall(_ctypes=ctypes):
1776
  """Lock current process' virtual address space into RAM.
1777

1778
  This is equivalent to the C call mlockall(MCL_CURRENT|MCL_FUTURE),
1779
  see mlock(2) for more details. This function requires ctypes module.
1780

1781
  @raises errors.NoCtypesError: if ctypes module is not found
1782

1783
  """
1784
  if _ctypes is None:
1785
    raise errors.NoCtypesError()
1786

    
1787
  libc = _ctypes.cdll.LoadLibrary("libc.so.6")
1788
  if libc is None:
1789
    logging.error("Cannot set memory lock, ctypes cannot load libc")
1790
    return
1791

    
1792
  # Some older version of the ctypes module don't have built-in functionality
1793
  # to access the errno global variable, where function error codes are stored.
1794
  # By declaring this variable as a pointer to an integer we can then access
1795
  # its value correctly, should the mlockall call fail, in order to see what
1796
  # the actual error code was.
1797
  # pylint: disable-msg=W0212
1798
  libc.__errno_location.restype = _ctypes.POINTER(_ctypes.c_int)
1799

    
1800
  if libc.mlockall(_MCL_CURRENT | _MCL_FUTURE):
1801
    # pylint: disable-msg=W0212
1802
    logging.error("Cannot set memory lock: %s",
1803
                  os.strerror(libc.__errno_location().contents.value))
1804
    return
1805

    
1806
  logging.debug("Memory lock set")
1807

    
1808

    
1809
def Daemonize(logfile):
1810
  """Daemonize the current process.
1811

1812
  This detaches the current process from the controlling terminal and
1813
  runs it in the background as a daemon.
1814

1815
  @type logfile: str
1816
  @param logfile: the logfile to which we should redirect stdout/stderr
1817
  @rtype: int
1818
  @return: the value zero
1819

1820
  """
1821
  # pylint: disable-msg=W0212
1822
  # yes, we really want os._exit
1823
  UMASK = 077
1824
  WORKDIR = "/"
1825

    
1826
  # this might fail
1827
  pid = os.fork()
1828
  if (pid == 0):  # The first child.
1829
    os.setsid()
1830
    # this might fail
1831
    pid = os.fork() # Fork a second child.
1832
    if (pid == 0):  # The second child.
1833
      os.chdir(WORKDIR)
1834
      os.umask(UMASK)
1835
    else:
1836
      # exit() or _exit()?  See below.
1837
      os._exit(0) # Exit parent (the first child) of the second child.
1838
  else:
1839
    os._exit(0) # Exit parent of the first child.
1840

    
1841
  for fd in range(3):
1842
    _CloseFDNoErr(fd)
1843
  i = os.open("/dev/null", os.O_RDONLY) # stdin
1844
  assert i == 0, "Can't close/reopen stdin"
1845
  i = os.open(logfile, os.O_WRONLY|os.O_CREAT|os.O_APPEND, 0600) # stdout
1846
  assert i == 1, "Can't close/reopen stdout"
1847
  # Duplicate standard output to standard error.
1848
  os.dup2(1, 2)
1849
  return 0
1850

    
1851

    
1852
def DaemonPidFileName(name):
1853
  """Compute a ganeti pid file absolute path
1854

1855
  @type name: str
1856
  @param name: the daemon name
1857
  @rtype: str
1858
  @return: the full path to the pidfile corresponding to the given
1859
      daemon name
1860

1861
  """
1862
  return PathJoin(constants.RUN_GANETI_DIR, "%s.pid" % name)
1863

    
1864

    
1865
def EnsureDaemon(name):
1866
  """Check for and start daemon if not alive.
1867

1868
  """
1869
  result = RunCmd([constants.DAEMON_UTIL, "check-and-start", name])
1870
  if result.failed:
1871
    logging.error("Can't start daemon '%s', failure %s, output: %s",
1872
                  name, result.fail_reason, result.output)
1873
    return False
1874

    
1875
  return True
1876

    
1877

    
1878
def WritePidFile(name):
1879
  """Write the current process pidfile.
1880

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

1883
  @type name: str
1884
  @param name: the daemon name to use
1885
  @raise errors.GenericError: if the pid file already exists and
1886
      points to a live process
1887

1888
  """
1889
  pid = os.getpid()
1890
  pidfilename = DaemonPidFileName(name)
1891
  if IsProcessAlive(ReadPidFile(pidfilename)):
1892
    raise errors.GenericError("%s contains a live process" % pidfilename)
1893

    
1894
  WriteFile(pidfilename, data="%d\n" % pid)
1895

    
1896

    
1897
def RemovePidFile(name):
1898
  """Remove the current process pidfile.
1899

1900
  Any errors are ignored.
1901

1902
  @type name: str
1903
  @param name: the daemon name used to derive the pidfile name
1904

1905
  """
1906
  pidfilename = DaemonPidFileName(name)
1907
  # TODO: we could check here that the file contains our pid
1908
  try:
1909
    RemoveFile(pidfilename)
1910
  except: # pylint: disable-msg=W0702
1911
    pass
1912

    
1913

    
1914
def KillProcess(pid, signal_=signal.SIGTERM, timeout=30,
1915
                waitpid=False):
1916
  """Kill a process given by its pid.
1917

1918
  @type pid: int
1919
  @param pid: The PID to terminate.
1920
  @type signal_: int
1921
  @param signal_: The signal to send, by default SIGTERM
1922
  @type timeout: int
1923
  @param timeout: The timeout after which, if the process is still alive,
1924
                  a SIGKILL will be sent. If not positive, no such checking
1925
                  will be done
1926
  @type waitpid: boolean
1927
  @param waitpid: If true, we should waitpid on this process after
1928
      sending signals, since it's our own child and otherwise it
1929
      would remain as zombie
1930

1931
  """
1932
  def _helper(pid, signal_, wait):
1933
    """Simple helper to encapsulate the kill/waitpid sequence"""
1934
    os.kill(pid, signal_)
1935
    if wait:
1936
      try:
1937
        os.waitpid(pid, os.WNOHANG)
1938
      except OSError:
1939
        pass
1940

    
1941
  if pid <= 0:
1942
    # kill with pid=0 == suicide
1943
    raise errors.ProgrammerError("Invalid pid given '%s'" % pid)
1944

    
1945
  if not IsProcessAlive(pid):
1946
    return
1947

    
1948
  _helper(pid, signal_, waitpid)
1949

    
1950
  if timeout <= 0:
1951
    return
1952

    
1953
  def _CheckProcess():
1954
    if not IsProcessAlive(pid):
1955
      return
1956

    
1957
    try:
1958
      (result_pid, _) = os.waitpid(pid, os.WNOHANG)
1959
    except OSError:
1960
      raise RetryAgain()
1961

    
1962
    if result_pid > 0:
1963
      return
1964

    
1965
    raise RetryAgain()
1966

    
1967
  try:
1968
    # Wait up to $timeout seconds
1969
    Retry(_CheckProcess, (0.01, 1.5, 0.1), timeout)
1970
  except RetryTimeout:
1971
    pass
1972

    
1973
  if IsProcessAlive(pid):
1974
    # Kill process if it's still alive
1975
    _helper(pid, signal.SIGKILL, waitpid)
1976

    
1977

    
1978
def FindFile(name, search_path, test=os.path.exists):
1979
  """Look for a filesystem object in a given path.
1980

1981
  This is an abstract method to search for filesystem object (files,
1982
  dirs) under a given search path.
1983

1984
  @type name: str
1985
  @param name: the name to look for
1986
  @type search_path: str
1987
  @param search_path: location to start at
1988
  @type test: callable
1989
  @param test: a function taking one argument that should return True
1990
      if the a given object is valid; the default value is
1991
      os.path.exists, causing only existing files to be returned
1992
  @rtype: str or None
1993
  @return: full path to the object if found, None otherwise
1994

1995
  """
1996
  # validate the filename mask
1997
  if constants.EXT_PLUGIN_MASK.match(name) is None:
1998
    logging.critical("Invalid value passed for external script name: '%s'",
1999
                     name)
2000
    return None
2001

    
2002
  for dir_name in search_path:
2003
    # FIXME: investigate switch to PathJoin
2004
    item_name = os.path.sep.join([dir_name, name])
2005
    # check the user test and that we're indeed resolving to the given
2006
    # basename
2007
    if test(item_name) and os.path.basename(item_name) == name:
2008
      return item_name
2009
  return None
2010

    
2011

    
2012
def CheckVolumeGroupSize(vglist, vgname, minsize):
2013
  """Checks if the volume group list is valid.
2014

2015
  The function will check if a given volume group is in the list of
2016
  volume groups and has a minimum size.
2017

2018
  @type vglist: dict
2019
  @param vglist: dictionary of volume group names and their size
2020
  @type vgname: str
2021
  @param vgname: the volume group we should check
2022
  @type minsize: int
2023
  @param minsize: the minimum size we accept
2024
  @rtype: None or str
2025
  @return: None for success, otherwise the error message
2026

2027
  """
2028
  vgsize = vglist.get(vgname, None)
2029
  if vgsize is None:
2030
    return "volume group '%s' missing" % vgname
2031
  elif vgsize < minsize:
2032
    return ("volume group '%s' too small (%s MiB required, %d MiB found)" %
2033
            (vgname, minsize, vgsize))
2034
  return None
2035

    
2036

    
2037
def SplitTime(value):
2038
  """Splits time as floating point number into a tuple.
2039

2040
  @param value: Time in seconds
2041
  @type value: int or float
2042
  @return: Tuple containing (seconds, microseconds)
2043

2044
  """
2045
  (seconds, microseconds) = divmod(int(value * 1000000), 1000000)
2046

    
2047
  assert 0 <= seconds, \
2048
    "Seconds must be larger than or equal to 0, but are %s" % seconds
2049
  assert 0 <= microseconds <= 999999, \
2050
    "Microseconds must be 0-999999, but are %s" % microseconds
2051

    
2052
  return (int(seconds), int(microseconds))
2053

    
2054

    
2055
def MergeTime(timetuple):
2056
  """Merges a tuple into time as a floating point number.
2057

2058
  @param timetuple: Time as tuple, (seconds, microseconds)
2059
  @type timetuple: tuple
2060
  @return: Time as a floating point number expressed in seconds
2061

2062
  """
2063
  (seconds, microseconds) = timetuple
2064

    
2065
  assert 0 <= seconds, \
2066
    "Seconds must be larger than or equal to 0, but are %s" % seconds
2067
  assert 0 <= microseconds <= 999999, \
2068
    "Microseconds must be 0-999999, but are %s" % microseconds
2069

    
2070
  return float(seconds) + (float(microseconds) * 0.000001)
2071

    
2072

    
2073
def GetDaemonPort(daemon_name):
2074
  """Get the daemon port for this cluster.
2075

2076
  Note that this routine does not read a ganeti-specific file, but
2077
  instead uses C{socket.getservbyname} to allow pre-customization of
2078
  this parameter outside of Ganeti.
2079

2080
  @type daemon_name: string
2081
  @param daemon_name: daemon name (in constants.DAEMONS_PORTS)
2082
  @rtype: int
2083

2084
  """
2085
  if daemon_name not in constants.DAEMONS_PORTS:
2086
    raise errors.ProgrammerError("Unknown daemon: %s" % daemon_name)
2087

    
2088
  (proto, default_port) = constants.DAEMONS_PORTS[daemon_name]
2089
  try:
2090
    port = socket.getservbyname(daemon_name, proto)
2091
  except socket.error:
2092
    port = default_port
2093

    
2094
  return port
2095

    
2096

    
2097
class LogFileHandler(logging.FileHandler):
2098
  """Log handler that doesn't fallback to stderr.
2099

2100
  When an error occurs while writing on the logfile, logging.FileHandler tries
2101
  to log on stderr. This doesn't work in ganeti since stderr is redirected to
2102
  the logfile. This class avoids failures reporting errors to /dev/console.
2103

2104
  """
2105
  def __init__(self, filename, mode="a", encoding=None):
2106
    """Open the specified file and use it as the stream for logging.
2107

2108
    Also open /dev/console to report errors while logging.
2109

2110
    """
2111
    logging.FileHandler.__init__(self, filename, mode, encoding)
2112
    self.console = open(constants.DEV_CONSOLE, "a")
2113

    
2114
  def handleError(self, record): # pylint: disable-msg=C0103
2115
    """Handle errors which occur during an emit() call.
2116

2117
    Try to handle errors with FileHandler method, if it fails write to
2118
    /dev/console.
2119

2120
    """
2121
    try:
2122
      logging.FileHandler.handleError(self, record)
2123
    except Exception: # pylint: disable-msg=W0703
2124
      try:
2125
        self.console.write("Cannot log message:\n%s\n" % self.format(record))
2126
      except Exception: # pylint: disable-msg=W0703
2127
        # Log handler tried everything it could, now just give up
2128
        pass
2129

    
2130

    
2131
def SetupLogging(logfile, debug=0, stderr_logging=False, program="",
2132
                 multithreaded=False, syslog=constants.SYSLOG_USAGE,
2133
                 console_logging=False):
2134
  """Configures the logging module.
2135

2136
  @type logfile: str
2137
  @param logfile: the filename to which we should log
2138
  @type debug: integer
2139
  @param debug: if greater than zero, enable debug messages, otherwise
2140
      only those at C{INFO} and above level
2141
  @type stderr_logging: boolean
2142
  @param stderr_logging: whether we should also log to the standard error
2143
  @type program: str
2144
  @param program: the name under which we should log messages
2145
  @type multithreaded: boolean
2146
  @param multithreaded: if True, will add the thread name to the log file
2147
  @type syslog: string
2148
  @param syslog: one of 'no', 'yes', 'only':
2149
      - if no, syslog is not used
2150
      - if yes, syslog is used (in addition to file-logging)
2151
      - if only, only syslog is used
2152
  @type console_logging: boolean
2153
  @param console_logging: if True, will use a FileHandler which falls back to
2154
      the system console if logging fails
2155
  @raise EnvironmentError: if we can't open the log file and
2156
      syslog/stderr logging is disabled
2157

2158
  """
2159
  fmt = "%(asctime)s: " + program + " pid=%(process)d"
2160
  sft = program + "[%(process)d]:"
2161
  if multithreaded:
2162
    fmt += "/%(threadName)s"
2163
    sft += " (%(threadName)s)"
2164
  if debug:
2165
    fmt += " %(module)s:%(lineno)s"
2166
    # no debug info for syslog loggers
2167
  fmt += " %(levelname)s %(message)s"
2168
  # yes, we do want the textual level, as remote syslog will probably
2169
  # lose the error level, and it's easier to grep for it
2170
  sft += " %(levelname)s %(message)s"
2171
  formatter = logging.Formatter(fmt)
2172
  sys_fmt = logging.Formatter(sft)
2173

    
2174
  root_logger = logging.getLogger("")
2175
  root_logger.setLevel(logging.NOTSET)
2176

    
2177
  # Remove all previously setup handlers
2178
  for handler in root_logger.handlers:
2179
    handler.close()
2180
    root_logger.removeHandler(handler)
2181

    
2182
  if stderr_logging:
2183
    stderr_handler = logging.StreamHandler()
2184
    stderr_handler.setFormatter(formatter)
2185
    if debug:
2186
      stderr_handler.setLevel(logging.NOTSET)
2187
    else:
2188
      stderr_handler.setLevel(logging.CRITICAL)
2189
    root_logger.addHandler(stderr_handler)
2190

    
2191
  if syslog in (constants.SYSLOG_YES, constants.SYSLOG_ONLY):
2192
    facility = logging.handlers.SysLogHandler.LOG_DAEMON
2193
    syslog_handler = logging.handlers.SysLogHandler(constants.SYSLOG_SOCKET,
2194
                                                    facility)
2195
    syslog_handler.setFormatter(sys_fmt)
2196
    # Never enable debug over syslog
2197
    syslog_handler.setLevel(logging.INFO)
2198
    root_logger.addHandler(syslog_handler)
2199

    
2200
  if syslog != constants.SYSLOG_ONLY:
2201
    # this can fail, if the logging directories are not setup or we have
2202
    # a permisssion problem; in this case, it's best to log but ignore
2203
    # the error if stderr_logging is True, and if false we re-raise the
2204
    # exception since otherwise we could run but without any logs at all
2205
    try:
2206
      if console_logging:
2207
        logfile_handler = LogFileHandler(logfile)
2208
      else:
2209
        logfile_handler = logging.FileHandler(logfile)
2210
      logfile_handler.setFormatter(formatter)
2211
      if debug:
2212
        logfile_handler.setLevel(logging.DEBUG)
2213
      else:
2214
        logfile_handler.setLevel(logging.INFO)
2215
      root_logger.addHandler(logfile_handler)
2216
    except EnvironmentError:
2217
      if stderr_logging or syslog == constants.SYSLOG_YES:
2218
        logging.exception("Failed to enable logging to file '%s'", logfile)
2219
      else:
2220
        # we need to re-raise the exception
2221
        raise
2222

    
2223

    
2224
def IsNormAbsPath(path):
2225
  """Check whether a path is absolute and also normalized
2226

2227
  This avoids things like /dir/../../other/path to be valid.
2228

2229
  """
2230
  return os.path.normpath(path) == path and os.path.isabs(path)
2231

    
2232

    
2233
def PathJoin(*args):
2234
  """Safe-join a list of path components.
2235

2236
  Requirements:
2237
      - the first argument must be an absolute path
2238
      - no component in the path must have backtracking (e.g. /../),
2239
        since we check for normalization at the end
2240

2241
  @param args: the path components to be joined
2242
  @raise ValueError: for invalid paths
2243

2244
  """
2245
  # ensure we're having at least one path passed in
2246
  assert args
2247
  # ensure the first component is an absolute and normalized path name
2248
  root = args[0]
2249
  if not IsNormAbsPath(root):
2250
    raise ValueError("Invalid parameter to PathJoin: '%s'" % str(args[0]))
2251
  result = os.path.join(*args)
2252
  # ensure that the whole path is normalized
2253
  if not IsNormAbsPath(result):
2254
    raise ValueError("Invalid parameters to PathJoin: '%s'" % str(args))
2255
  # check that we're still under the original prefix
2256
  prefix = os.path.commonprefix([root, result])
2257
  if prefix != root:
2258
    raise ValueError("Error: path joining resulted in different prefix"
2259
                     " (%s != %s)" % (prefix, root))
2260
  return result
2261

    
2262

    
2263
def TailFile(fname, lines=20):
2264
  """Return the last lines from a file.
2265

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

2270
  @param fname: the file name
2271
  @type lines: int
2272
  @param lines: the (maximum) number of lines to return
2273

2274
  """
2275
  fd = open(fname, "r")
2276
  try:
2277
    fd.seek(0, 2)
2278
    pos = fd.tell()
2279
    pos = max(0, pos-4096)
2280
    fd.seek(pos, 0)
2281
    raw_data = fd.read()
2282
  finally:
2283
    fd.close()
2284

    
2285
  rows = raw_data.splitlines()
2286
  return rows[-lines:]
2287

    
2288

    
2289
def _ParseAsn1Generalizedtime(value):
2290
  """Parses an ASN1 GENERALIZEDTIME timestamp as used by pyOpenSSL.
2291

2292
  @type value: string
2293
  @param value: ASN1 GENERALIZEDTIME timestamp
2294

2295
  """
2296
  m = re.match(r"^(\d+)([-+]\d\d)(\d\d)$", value)
2297
  if m:
2298
    # We have an offset
2299
    asn1time = m.group(1)
2300
    hours = int(m.group(2))
2301
    minutes = int(m.group(3))
2302
    utcoffset = (60 * hours) + minutes
2303
  else:
2304
    if not value.endswith("Z"):
2305
      raise ValueError("Missing timezone")
2306
    asn1time = value[:-1]
2307
    utcoffset = 0
2308

    
2309
  parsed = time.strptime(asn1time, "%Y%m%d%H%M%S")
2310

    
2311
  tt = datetime.datetime(*(parsed[:7])) - datetime.timedelta(minutes=utcoffset)
2312

    
2313
  return calendar.timegm(tt.utctimetuple())
2314

    
2315

    
2316
def GetX509CertValidity(cert):
2317
  """Returns the validity period of the certificate.
2318

2319
  @type cert: OpenSSL.crypto.X509
2320
  @param cert: X509 certificate object
2321

2322
  """
2323
  # The get_notBefore and get_notAfter functions are only supported in
2324
  # pyOpenSSL 0.7 and above.
2325
  try:
2326
    get_notbefore_fn = cert.get_notBefore
2327
  except AttributeError:
2328
    not_before = None
2329
  else:
2330
    not_before_asn1 = get_notbefore_fn()
2331

    
2332
    if not_before_asn1 is None:
2333
      not_before = None
2334
    else:
2335
      not_before = _ParseAsn1Generalizedtime(not_before_asn1)
2336

    
2337
  try:
2338
    get_notafter_fn = cert.get_notAfter
2339
  except AttributeError:
2340
    not_after = None
2341
  else:
2342
    not_after_asn1 = get_notafter_fn()
2343

    
2344
    if not_after_asn1 is None:
2345
      not_after = None
2346
    else:
2347
      not_after = _ParseAsn1Generalizedtime(not_after_asn1)
2348

    
2349
  return (not_before, not_after)
2350

    
2351

    
2352
def SafeEncode(text):
2353
  """Return a 'safe' version of a source string.
2354

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

2364
  @type text: str or unicode
2365
  @param text: input data
2366
  @rtype: str
2367
  @return: a safe version of text
2368

2369
  """
2370
  if isinstance(text, unicode):
2371
    # only if unicode; if str already, we handle it below
2372
    text = text.encode('ascii', 'backslashreplace')
2373
  resu = ""
2374
  for char in text:
2375
    c = ord(char)
2376
    if char  == '\t':
2377
      resu += r'\t'
2378
    elif char == '\n':
2379
      resu += r'\n'
2380
    elif char == '\r':
2381
      resu += r'\'r'
2382
    elif c < 32 or c >= 127: # non-printable
2383
      resu += "\\x%02x" % (c & 0xff)
2384
    else:
2385
      resu += char
2386
  return resu
2387

    
2388

    
2389
def UnescapeAndSplit(text, sep=","):
2390
  """Split and unescape a string based on a given separator.
2391

2392
  This function splits a string based on a separator where the
2393
  separator itself can be escape in order to be an element of the
2394
  elements. The escaping rules are (assuming coma being the
2395
  separator):
2396
    - a plain , separates the elements
2397
    - a sequence \\\\, (double backslash plus comma) is handled as a
2398
      backslash plus a separator comma
2399
    - a sequence \, (backslash plus comma) is handled as a
2400
      non-separator comma
2401

2402
  @type text: string
2403
  @param text: the string to split
2404
  @type sep: string
2405
  @param text: the separator
2406
  @rtype: string
2407
  @return: a list of strings
2408

2409
  """
2410
  # we split the list by sep (with no escaping at this stage)
2411
  slist = text.split(sep)
2412
  # next, we revisit the elements and if any of them ended with an odd
2413
  # number of backslashes, then we join it with the next
2414
  rlist = []
2415
  while slist:
2416
    e1 = slist.pop(0)
2417
    if e1.endswith("\\"):
2418
      num_b = len(e1) - len(e1.rstrip("\\"))
2419
      if num_b % 2 == 1:
2420
        e2 = slist.pop(0)
2421
        # here the backslashes remain (all), and will be reduced in
2422
        # the next step
2423
        rlist.append(e1 + sep + e2)
2424
        continue
2425
    rlist.append(e1)
2426
  # finally, replace backslash-something with something
2427
  rlist = [re.sub(r"\\(.)", r"\1", v) for v in rlist]
2428
  return rlist
2429

    
2430

    
2431
def CommaJoin(names):
2432
  """Nicely join a set of identifiers.
2433

2434
  @param names: set, list or tuple
2435
  @return: a string with the formatted results
2436

2437
  """
2438
  return ", ".join([str(val) for val in names])
2439

    
2440

    
2441
def BytesToMebibyte(value):
2442
  """Converts bytes to mebibytes.
2443

2444
  @type value: int
2445
  @param value: Value in bytes
2446
  @rtype: int
2447
  @return: Value in mebibytes
2448

2449
  """
2450
  return int(round(value / (1024.0 * 1024.0), 0))
2451

    
2452

    
2453
def CalculateDirectorySize(path):
2454
  """Calculates the size of a directory recursively.
2455

2456
  @type path: string
2457
  @param path: Path to directory
2458
  @rtype: int
2459
  @return: Size in mebibytes
2460

2461
  """
2462
  size = 0
2463

    
2464
  for (curpath, _, files) in os.walk(path):
2465
    for filename in files:
2466
      st = os.lstat(PathJoin(curpath, filename))
2467
      size += st.st_size
2468

    
2469
  return BytesToMebibyte(size)
2470

    
2471

    
2472
def GetFilesystemStats(path):
2473
  """Returns the total and free space on a filesystem.
2474

2475
  @type path: string
2476
  @param path: Path on filesystem to be examined
2477
  @rtype: int
2478
  @return: tuple of (Total space, Free space) in mebibytes
2479

2480
  """
2481
  st = os.statvfs(path)
2482

    
2483
  fsize = BytesToMebibyte(st.f_bavail * st.f_frsize)
2484
  tsize = BytesToMebibyte(st.f_blocks * st.f_frsize)
2485
  return (tsize, fsize)
2486

    
2487

    
2488
def RunInSeparateProcess(fn, *args):
2489
  """Runs a function in a separate process.
2490

2491
  Note: Only boolean return values are supported.
2492

2493
  @type fn: callable
2494
  @param fn: Function to be called
2495
  @rtype: bool
2496
  @return: Function's result
2497

2498
  """
2499
  pid = os.fork()
2500
  if pid == 0:
2501
    # Child process
2502
    try:
2503
      # In case the function uses temporary files
2504
      ResetTempfileModule()
2505

    
2506
      # Call function
2507
      result = int(bool(fn(*args)))
2508
      assert result in (0, 1)
2509
    except: # pylint: disable-msg=W0702
2510
      logging.exception("Error while calling function in separate process")
2511
      # 0 and 1 are reserved for the return value
2512
      result = 33
2513

    
2514
    os._exit(result) # pylint: disable-msg=W0212
2515

    
2516
  # Parent process
2517

    
2518
  # Avoid zombies and check exit code
2519
  (_, status) = os.waitpid(pid, 0)
2520

    
2521
  if os.WIFSIGNALED(status):
2522
    exitcode = None
2523
    signum = os.WTERMSIG(status)
2524
  else:
2525
    exitcode = os.WEXITSTATUS(status)
2526
    signum = None
2527

    
2528
  if not (exitcode in (0, 1) and signum is None):
2529
    raise errors.GenericError("Child program failed (code=%s, signal=%s)" %
2530
                              (exitcode, signum))
2531

    
2532
  return bool(exitcode)
2533

    
2534

    
2535
def IgnoreSignals(fn, *args, **kwargs):
2536
  """Tries to call a function ignoring failures due to EINTR.
2537

2538
  """
2539
  try:
2540
    return fn(*args, **kwargs)
2541
  except EnvironmentError, err:
2542
    if err.errno == errno.EINTR:
2543
      return None
2544
    else:
2545
      raise
2546
  except (select.error, socket.error), err:
2547
    # In python 2.6 and above select.error is an IOError, so it's handled
2548
    # above, in 2.5 and below it's not, and it's handled here.
2549
    if err.args and err.args[0] == errno.EINTR:
2550
      return None
2551
    else:
2552
      raise
2553

    
2554

    
2555
def LockedMethod(fn):
2556
  """Synchronized object access decorator.
2557

2558
  This decorator is intended to protect access to an object using the
2559
  object's own lock which is hardcoded to '_lock'.
2560

2561
  """
2562
  def _LockDebug(*args, **kwargs):
2563
    if debug_locks:
2564
      logging.debug(*args, **kwargs)
2565

    
2566
  def wrapper(self, *args, **kwargs):
2567
    # pylint: disable-msg=W0212
2568
    assert hasattr(self, '_lock')
2569
    lock = self._lock
2570
    _LockDebug("Waiting for %s", lock)
2571
    lock.acquire()
2572
    try:
2573
      _LockDebug("Acquired %s", lock)
2574
      result = fn(self, *args, **kwargs)
2575
    finally:
2576
      _LockDebug("Releasing %s", lock)
2577
      lock.release()
2578
      _LockDebug("Released %s", lock)
2579
    return result
2580
  return wrapper
2581

    
2582

    
2583
def LockFile(fd):
2584
  """Locks a file using POSIX locks.
2585

2586
  @type fd: int
2587
  @param fd: the file descriptor we need to lock
2588

2589
  """
2590
  try:
2591
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
2592
  except IOError, err:
2593
    if err.errno == errno.EAGAIN:
2594
      raise errors.LockError("File already locked")
2595
    raise
2596

    
2597

    
2598
def FormatTime(val):
2599
  """Formats a time value.
2600

2601
  @type val: float or None
2602
  @param val: the timestamp as returned by time.time()
2603
  @return: a string value or N/A if we don't have a valid timestamp
2604

2605
  """
2606
  if val is None or not isinstance(val, (int, float)):
2607
    return "N/A"
2608
  # these two codes works on Linux, but they are not guaranteed on all
2609
  # platforms
2610
  return time.strftime("%F %T", time.localtime(val))
2611

    
2612

    
2613
def ReadWatcherPauseFile(filename, now=None, remove_after=3600):
2614
  """Reads the watcher pause file.
2615

2616
  @type filename: string
2617
  @param filename: Path to watcher pause file
2618
  @type now: None, float or int
2619
  @param now: Current time as Unix timestamp
2620
  @type remove_after: int
2621
  @param remove_after: Remove watcher pause file after specified amount of
2622
    seconds past the pause end time
2623

2624
  """
2625
  if now is None:
2626
    now = time.time()
2627

    
2628
  try:
2629
    value = ReadFile(filename)
2630
  except IOError, err:
2631
    if err.errno != errno.ENOENT:
2632
      raise
2633
    value = None
2634

    
2635
  if value is not None:
2636
    try:
2637
      value = int(value)
2638
    except ValueError:
2639
      logging.warning(("Watcher pause file (%s) contains invalid value,"
2640
                       " removing it"), filename)
2641
      RemoveFile(filename)
2642
      value = None
2643

    
2644
    if value is not None:
2645
      # Remove file if it's outdated
2646
      if now > (value + remove_after):
2647
        RemoveFile(filename)
2648
        value = None
2649

    
2650
      elif now > value:
2651
        value = None
2652

    
2653
  return value
2654

    
2655

    
2656
class RetryTimeout(Exception):
2657
  """Retry loop timed out.
2658

2659
  Any arguments which was passed by the retried function to RetryAgain will be
2660
  preserved in RetryTimeout, if it is raised. If such argument was an exception
2661
  the RaiseInner helper method will reraise it.
2662

2663
  """
2664
  def RaiseInner(self):
2665
    if self.args and isinstance(self.args[0], Exception):
2666
      raise self.args[0]
2667
    else:
2668
      raise RetryTimeout(*self.args)
2669

    
2670

    
2671
class RetryAgain(Exception):
2672
  """Retry again.
2673

2674
  Any arguments passed to RetryAgain will be preserved, if a timeout occurs, as
2675
  arguments to RetryTimeout. If an exception is passed, the RaiseInner() method
2676
  of the RetryTimeout() method can be used to reraise it.
2677

2678
  """
2679

    
2680

    
2681
class _RetryDelayCalculator(object):
2682
  """Calculator for increasing delays.
2683

2684
  """
2685
  __slots__ = [
2686
    "_factor",
2687
    "_limit",
2688
    "_next",
2689
    "_start",
2690
    ]
2691

    
2692
  def __init__(self, start, factor, limit):
2693
    """Initializes this class.
2694

2695
    @type start: float
2696
    @param start: Initial delay
2697
    @type factor: float
2698
    @param factor: Factor for delay increase
2699
    @type limit: float or None
2700
    @param limit: Upper limit for delay or None for no limit
2701

2702
    """
2703
    assert start > 0.0
2704
    assert factor >= 1.0
2705
    assert limit is None or limit >= 0.0
2706

    
2707
    self._start = start
2708
    self._factor = factor
2709
    self._limit = limit
2710

    
2711
    self._next = start
2712

    
2713
  def __call__(self):
2714
    """Returns current delay and calculates the next one.
2715

2716
    """
2717
    current = self._next
2718

    
2719
    # Update for next run
2720
    if self._limit is None or self._next < self._limit:
2721
      self._next = min(self._limit, self._next * self._factor)
2722

    
2723
    return current
2724

    
2725

    
2726
#: Special delay to specify whole remaining timeout
2727
RETRY_REMAINING_TIME = object()
2728

    
2729

    
2730
def Retry(fn, delay, timeout, args=None, wait_fn=time.sleep,
2731
          _time_fn=time.time):
2732
  """Call a function repeatedly until it succeeds.
2733

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

2738
  C{delay} can be one of the following:
2739
    - callable returning the delay length as a float
2740
    - Tuple of (start, factor, limit)
2741
    - L{RETRY_REMAINING_TIME} to sleep until the timeout expires (this is
2742
      useful when overriding L{wait_fn} to wait for an external event)
2743
    - A static delay as a number (int or float)
2744

2745
  @type fn: callable
2746
  @param fn: Function to be called
2747
  @param delay: Either a callable (returning the delay), a tuple of (start,
2748
                factor, limit) (see L{_RetryDelayCalculator}),
2749
                L{RETRY_REMAINING_TIME} or a number (int or float)
2750
  @type timeout: float
2751
  @param timeout: Total timeout
2752
  @type wait_fn: callable
2753
  @param wait_fn: Waiting function
2754
  @return: Return value of function
2755

2756
  """
2757
  assert callable(fn)
2758
  assert callable(wait_fn)
2759
  assert callable(_time_fn)
2760

    
2761
  if args is None:
2762
    args = []
2763

    
2764
  end_time = _time_fn() + timeout
2765

    
2766
  if callable(delay):
2767
    # External function to calculate delay
2768
    calc_delay = delay
2769

    
2770
  elif isinstance(delay, (tuple, list)):
2771
    # Increasing delay with optional upper boundary
2772
    (start, factor, limit) = delay
2773
    calc_delay = _RetryDelayCalculator(start, factor, limit)
2774

    
2775
  elif delay is RETRY_REMAINING_TIME:
2776
    # Always use the remaining time
2777
    calc_delay = None
2778

    
2779
  else:
2780
    # Static delay
2781
    calc_delay = lambda: delay
2782

    
2783
  assert calc_delay is None or callable(calc_delay)
2784

    
2785
  while True:
2786
    retry_args = []
2787
    try:
2788
      # pylint: disable-msg=W0142
2789
      return fn(*args)
2790
    except RetryAgain, err:
2791
      retry_args = err.args
2792
    except RetryTimeout:
2793
      raise errors.ProgrammerError("Nested retry loop detected that didn't"
2794
                                   " handle RetryTimeout")
2795

    
2796
    remaining_time = end_time - _time_fn()
2797

    
2798
    if remaining_time < 0.0:
2799
      # pylint: disable-msg=W0142
2800
      raise RetryTimeout(*retry_args)
2801

    
2802
    assert remaining_time >= 0.0
2803

    
2804
    if calc_delay is None:
2805
      wait_fn(remaining_time)
2806
    else:
2807
      current_delay = calc_delay()
2808
      if current_delay > 0.0:
2809
        wait_fn(current_delay)
2810

    
2811

    
2812
class FileLock(object):
2813
  """Utility class for file locks.
2814

2815
  """
2816
  def __init__(self, fd, filename):
2817
    """Constructor for FileLock.
2818

2819
    @type fd: file
2820
    @param fd: File object
2821
    @type filename: str
2822
    @param filename: Path of the file opened at I{fd}
2823

2824
    """
2825
    self.fd = fd
2826
    self.filename = filename
2827

    
2828
  @classmethod
2829
  def Open(cls, filename):
2830
    """Creates and opens a file to be used as a file-based lock.
2831

2832
    @type filename: string
2833
    @param filename: path to the file to be locked
2834

2835
    """
2836
    # Using "os.open" is necessary to allow both opening existing file
2837
    # read/write and creating if not existing. Vanilla "open" will truncate an
2838
    # existing file -or- allow creating if not existing.
2839
    return cls(os.fdopen(os.open(filename, os.O_RDWR | os.O_CREAT), "w+"),
2840
               filename)
2841

    
2842
  def __del__(self):
2843
    self.Close()
2844

    
2845
  def Close(self):
2846
    """Close the file and release the lock.
2847

2848
    """
2849
    if hasattr(self, "fd") and self.fd:
2850
      self.fd.close()
2851
      self.fd = None
2852

    
2853
  def _flock(self, flag, blocking, timeout, errmsg):
2854
    """Wrapper for fcntl.flock.
2855

2856
    @type flag: int
2857
    @param flag: operation flag
2858
    @type blocking: bool
2859
    @param blocking: whether the operation should be done in blocking mode.
2860
    @type timeout: None or float
2861
    @param timeout: for how long the operation should be retried (implies
2862
                    non-blocking mode).
2863
    @type errmsg: string
2864
    @param errmsg: error message in case operation fails.
2865

2866
    """
2867
    assert self.fd, "Lock was closed"
2868
    assert timeout is None or timeout >= 0, \
2869
      "If specified, timeout must be positive"
2870
    assert not (flag & fcntl.LOCK_NB), "LOCK_NB must not be set"
2871

    
2872
    # When a timeout is used, LOCK_NB must always be set
2873
    if not (timeout is None and blocking):
2874
      flag |= fcntl.LOCK_NB
2875

    
2876
    if timeout is None:
2877
      self._Lock(self.fd, flag, timeout)
2878
    else:
2879
      try:
2880
        Retry(self._Lock, (0.1, 1.2, 1.0), timeout,
2881
              args=(self.fd, flag, timeout))
2882
      except RetryTimeout:
2883
        raise errors.LockError(errmsg)
2884

    
2885
  @staticmethod
2886
  def _Lock(fd, flag, timeout):
2887
    try:
2888
      fcntl.flock(fd, flag)
2889
    except IOError, err:
2890
      if timeout is not None and err.errno == errno.EAGAIN:
2891
        raise RetryAgain()
2892

    
2893
      logging.exception("fcntl.flock failed")
2894
      raise
2895

    
2896
  def Exclusive(self, blocking=False, timeout=None):
2897
    """Locks the file in exclusive mode.
2898

2899
    @type blocking: boolean
2900
    @param blocking: whether to block and wait until we
2901
        can lock the file or return immediately
2902
    @type timeout: int or None
2903
    @param timeout: if not None, the duration to wait for the lock
2904
        (in blocking mode)
2905

2906
    """
2907
    self._flock(fcntl.LOCK_EX, blocking, timeout,
2908
                "Failed to lock %s in exclusive mode" % self.filename)
2909

    
2910
  def Shared(self, blocking=False, timeout=None):
2911
    """Locks the file in shared mode.
2912

2913
    @type blocking: boolean
2914
    @param blocking: whether to block and wait until we
2915
        can lock the file or return immediately
2916
    @type timeout: int or None
2917
    @param timeout: if not None, the duration to wait for the lock
2918
        (in blocking mode)
2919

2920
    """
2921
    self._flock(fcntl.LOCK_SH, blocking, timeout,
2922
                "Failed to lock %s in shared mode" % self.filename)
2923

    
2924
  def Unlock(self, blocking=True, timeout=None):
2925
    """Unlocks the file.
2926

2927
    According to C{flock(2)}, unlocking can also be a nonblocking
2928
    operation::
2929

2930
      To make a non-blocking request, include LOCK_NB with any of the above
2931
      operations.
2932

2933
    @type blocking: boolean
2934
    @param blocking: whether to block and wait until we
2935
        can lock the file or return immediately
2936
    @type timeout: int or None
2937
    @param timeout: if not None, the duration to wait for the lock
2938
        (in blocking mode)
2939

2940
    """
2941
    self._flock(fcntl.LOCK_UN, blocking, timeout,
2942
                "Failed to unlock %s" % self.filename)
2943

    
2944

    
2945
class LineSplitter:
2946
  """Splits data chunks into lines separated by newline.
2947

2948
  Instances provide a file-like interface.
2949

2950
  """
2951
  def __init__(self, line_fn, *args):
2952
    """Initializes this class.
2953

2954
    @type line_fn: callable
2955
    @param line_fn: Function called for each line, first parameter is line
2956
    @param args: Extra arguments for L{line_fn}
2957

2958
    """
2959
    assert callable(line_fn)
2960

    
2961
    if args:
2962
      # Python 2.4 doesn't have functools.partial yet
2963
      self._line_fn = \
2964
        lambda line: line_fn(line, *args) # pylint: disable-msg=W0142
2965
    else:
2966
      self._line_fn = line_fn
2967

    
2968
    self._lines = collections.deque()
2969
    self._buffer = ""
2970

    
2971
  def write(self, data):
2972
    parts = (self._buffer + data).split("\n")
2973
    self._buffer = parts.pop()
2974
    self._lines.extend(parts)
2975

    
2976
  def flush(self):
2977
    while self._lines:
2978
      self._line_fn(self._lines.popleft().rstrip("\r\n"))
2979

    
2980
  def close(self):
2981
    self.flush()
2982
    if self._buffer:
2983
      self._line_fn(self._buffer)
2984

    
2985

    
2986
def SignalHandled(signums):
2987
  """Signal Handled decoration.
2988

2989
  This special decorator installs a signal handler and then calls the target
2990
  function. The function must accept a 'signal_handlers' keyword argument,
2991
  which will contain a dict indexed by signal number, with SignalHandler
2992
  objects as values.
2993

2994
  The decorator can be safely stacked with iself, to handle multiple signals
2995
  with different handlers.
2996

2997
  @type signums: list
2998
  @param signums: signals to intercept
2999

3000
  """
3001
  def wrap(fn):
3002
    def sig_function(*args, **kwargs):
3003
      assert 'signal_handlers' not in kwargs or \
3004
             kwargs['signal_handlers'] is None or \
3005
             isinstance(kwargs['signal_handlers'], dict), \
3006
             "Wrong signal_handlers parameter in original function call"
3007
      if 'signal_handlers' in kwargs and kwargs['signal_handlers'] is not None:
3008
        signal_handlers = kwargs['signal_handlers']
3009
      else:
3010
        signal_handlers = {}
3011
        kwargs['signal_handlers'] = signal_handlers
3012
      sighandler = SignalHandler(signums)
3013
      try:
3014
        for sig in signums:
3015
          signal_handlers[sig] = sighandler
3016
        return fn(*args, **kwargs)
3017
      finally:
3018
        sighandler.Reset()
3019
    return sig_function
3020
  return wrap
3021

    
3022

    
3023
class SignalHandler(object):
3024
  """Generic signal handler class.
3025

3026
  It automatically restores the original handler when deconstructed or
3027
  when L{Reset} is called. You can either pass your own handler
3028
  function in or query the L{called} attribute to detect whether the
3029
  signal was sent.
3030

3031
  @type signum: list
3032
  @ivar signum: the signals we handle
3033
  @type called: boolean
3034
  @ivar called: tracks whether any of the signals have been raised
3035

3036
  """
3037
  def __init__(self, signum):
3038
    """Constructs a new SignalHandler instance.
3039

3040
    @type signum: int or list of ints
3041
    @param signum: Single signal number or set of signal numbers
3042

3043
    """
3044
    self.signum = set(signum)
3045
    self.called = False
3046

    
3047
    self._previous = {}
3048
    try:
3049
      for signum in self.signum:
3050
        # Setup handler
3051
        prev_handler = signal.signal(signum, self._HandleSignal)
3052
        try:
3053
          self._previous[signum] = prev_handler
3054
        except:
3055
          # Restore previous handler
3056
          signal.signal(signum, prev_handler)
3057
          raise
3058
    except:
3059
      # Reset all handlers
3060
      self.Reset()
3061
      # Here we have a race condition: a handler may have already been called,
3062
      # but there's not much we can do about it at this point.
3063
      raise
3064

    
3065
  def __del__(self):
3066
    self.Reset()
3067

    
3068
  def Reset(self):
3069
    """Restore previous handler.
3070

3071
    This will reset all the signals to their previous handlers.
3072

3073
    """
3074
    for signum, prev_handler in self._previous.items():
3075
      signal.signal(signum, prev_handler)
3076
      # If successful, remove from dict
3077
      del self._previous[signum]
3078

    
3079
  def Clear(self):
3080
    """Unsets the L{called} flag.
3081

3082
    This function can be used in case a signal may arrive several times.
3083

3084
    """
3085
    self.called = False
3086

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

3091
    """
3092
    # This is not nice and not absolutely atomic, but it appears to be the only
3093
    # solution in Python -- there are no atomic types.
3094
    self.called = True
3095

    
3096

    
3097
class FieldSet(object):
3098
  """A simple field set.
3099

3100
  Among the features are:
3101
    - checking if a string is among a list of static string or regex objects
3102
    - checking if a whole list of string matches
3103
    - returning the matching groups from a regex match
3104

3105
  Internally, all fields are held as regular expression objects.
3106

3107
  """
3108
  def __init__(self, *items):
3109
    self.items = [re.compile("^%s$" % value) for value in items]
3110

    
3111
  def Extend(self, other_set):
3112
    """Extend the field set with the items from another one"""
3113
    self.items.extend(other_set.items)
3114

    
3115
  def Matches(self, field):
3116
    """Checks if a field matches the current set
3117

3118
    @type field: str
3119
    @param field: the string to match
3120
    @return: either None or a regular expression match object
3121

3122
    """
3123
    for m in itertools.ifilter(None, (val.match(field) for val in self.items)):
3124
      return m
3125
    return None
3126

    
3127
  def NonMatching(self, items):
3128
    """Returns the list of fields not matching the current set
3129

3130
    @type items: list
3131
    @param items: the list of fields to check
3132
    @rtype: list
3133
    @return: list of non-matching fields
3134

3135
    """
3136
    return [val for val in items if not self.Matches(val)]