Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ f95c81bf

History | View | Annotate | Download (56.5 kB)

1
#
2
#
3

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

    
21

    
22
"""Ganeti utility module.
23

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

27
"""
28

    
29

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

    
46
from cStringIO import StringIO
47

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

    
54
from ganeti import errors
55
from ganeti import constants
56

    
57

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

    
61
debug_locks = False
62

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

    
66

    
67
class RunResult(object):
68
  """Holds the result of running external programs.
69

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

85
  """
86
  __slots__ = ["exit_code", "signal", "stdout", "stderr",
87
               "failed", "fail_reason", "cmd"]
88

    
89

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

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

    
105
    if self.failed:
106
      logging.debug("Command '%s' failed (%s); output: %s",
107
                    self.cmd, self.fail_reason, self.output)
108

    
109
  def _GetOutput(self):
110
    """Returns the combined stdout and stderr for easier usage.
111

112
    """
113
    return self.stdout + self.stderr
114

    
115
  output = property(_GetOutput, None, None, "Return full output")
116

    
117

    
118
def RunCmd(cmd, env=None, output=None, cwd='/'):
119
  """Execute a (shell) command.
120

121
  The command should not read from its standard input, as it will be
122
  closed.
123

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

139
  """
140
  if no_fork:
141
    raise errors.ProgrammerError("utils.RunCmd() called with fork() disabled")
142

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

    
152
  cmd_env = os.environ.copy()
153
  cmd_env["LC_ALL"] = "C"
154
  if env is not None:
155
    cmd_env.update(env)
156

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

    
170
  if status >= 0:
171
    exitcode = status
172
    signal_ = None
173
  else:
174
    exitcode = None
175
    signal_ = -status
176

    
177
  return RunResult(exitcode, signal_, out, err, strcmd)
178

    
179

    
180
def _RunCmdPipe(cmd, env, via_shell, cwd):
181
  """Run a command and return its output.
182

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

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

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

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

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

    
242
  out = out.getvalue()
243
  err = err.getvalue()
244

    
245
  status = child.wait()
246
  return out, err, status
247

    
248

    
249
def _RunCmdFile(cmd, env, via_shell, output, cwd):
250
  """Run a command and save its output to a file.
251

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

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

    
275
    child.stdin.close()
276
    status = child.wait()
277
  finally:
278
    fh.close()
279
  return status
280

    
281

    
282
def RemoveFile(filename):
283
  """Remove a file ignoring some errors.
284

285
  Remove a file, ignoring non-existing ones or directories. Other
286
  errors are passed.
287

288
  @type filename: str
289
  @param filename: the file to be removed
290

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

    
298

    
299
def RenameFile(old, new, mkdir=False, mkdir_mode=0750):
300
  """Renames a file.
301

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

311
  """
312
  try:
313
    return os.rename(old, new)
314
  except OSError, err:
315
    # In at least one use case of this function, the job queue, directory
316
    # creation is very rare. Checking for the directory before renaming is not
317
    # as efficient.
318
    if mkdir and err.errno == errno.ENOENT:
319
      # Create directory and try again
320
      os.makedirs(os.path.dirname(new), mkdir_mode)
321
      return os.rename(old, new)
322
    raise
323

    
324

    
325
def _FingerprintFile(filename):
326
  """Compute the fingerprint of a file.
327

328
  If the file does not exist, a None will be returned
329
  instead.
330

331
  @type filename: str
332
  @param filename: the filename to checksum
333
  @rtype: str
334
  @return: the hex digest of the sha checksum of the contents
335
      of the file
336

337
  """
338
  if not (os.path.exists(filename) and os.path.isfile(filename)):
339
    return None
340

    
341
  f = open(filename)
342

    
343
  fp = sha1()
344
  while True:
345
    data = f.read(4096)
346
    if not data:
347
      break
348

    
349
    fp.update(data)
350

    
351
  return fp.hexdigest()
352

    
353

    
354
def FingerprintFiles(files):
355
  """Compute fingerprints for a list of files.
356

357
  @type files: list
358
  @param files: the list of filename to fingerprint
359
  @rtype: dict
360
  @return: a dictionary filename: fingerprint, holding only
361
      existing files
362

363
  """
364
  ret = {}
365

    
366
  for filename in files:
367
    cksum = _FingerprintFile(filename)
368
    if cksum:
369
      ret[filename] = cksum
370

    
371
  return ret
372

    
373

    
374
def CheckDict(target, template, logname=None):
375
  """Ensure a dictionary has a required set of keys.
376

377
  For the given dictionaries I{target} and I{template}, ensure
378
  I{target} has all the keys from I{template}. Missing keys are added
379
  with values from template.
380

381
  @type target: dict
382
  @param target: the dictionary to update
383
  @type template: dict
384
  @param template: the dictionary holding the default values
385
  @type logname: str or None
386
  @param logname: if not None, causes the missing keys to be
387
      logged with this name
388

389
  """
390
  missing = []
391
  for k in template:
392
    if k not in target:
393
      missing.append(k)
394
      target[k] = template[k]
395

    
396
  if missing and logname:
397
    logging.warning('%s missing keys %s', logname, ', '.join(missing))
398

    
399

    
400
def ForceDictType(target, key_types, allowed_values=None):
401
  """Force the values of a dict to have certain types.
402

403
  @type target: dict
404
  @param target: the dict to update
405
  @type key_types: dict
406
  @param key_types: dict mapping target dict keys to types
407
                    in constants.ENFORCEABLE_TYPES
408
  @type allowed_values: list
409
  @keyword allowed_values: list of specially allowed values
410

411
  """
412
  if allowed_values is None:
413
    allowed_values = []
414

    
415
  for key in target:
416
    if key not in key_types:
417
      msg = "Unknown key '%s'" % key
418
      raise errors.TypeEnforcementError(msg)
419

    
420
    if target[key] in allowed_values:
421
      continue
422

    
423
    type = key_types[key]
424
    if type not in constants.ENFORCEABLE_TYPES:
425
      msg = "'%s' has non-enforceable type %s" % (key, type)
426
      raise errors.ProgrammerError(msg)
427

    
428
    if type == constants.VTYPE_STRING:
429
      if not isinstance(target[key], basestring):
430
        if isinstance(target[key], bool) and not target[key]:
431
          target[key] = ''
432
        else:
433
          msg = "'%s' (value %s) is not a valid string" % (key, target[key])
434
          raise errors.TypeEnforcementError(msg)
435
    elif type == constants.VTYPE_BOOL:
436
      if isinstance(target[key], basestring) and target[key]:
437
        if target[key].lower() == constants.VALUE_FALSE:
438
          target[key] = False
439
        elif target[key].lower() == constants.VALUE_TRUE:
440
          target[key] = True
441
        else:
442
          msg = "'%s' (value %s) is not a valid boolean" % (key, target[key])
443
          raise errors.TypeEnforcementError(msg)
444
      elif target[key]:
445
        target[key] = True
446
      else:
447
        target[key] = False
448
    elif type == constants.VTYPE_SIZE:
449
      try:
450
        target[key] = ParseUnit(target[key])
451
      except errors.UnitParseError, err:
452
        msg = "'%s' (value %s) is not a valid size. error: %s" % \
453
              (key, target[key], err)
454
        raise errors.TypeEnforcementError(msg)
455
    elif type == constants.VTYPE_INT:
456
      try:
457
        target[key] = int(target[key])
458
      except (ValueError, TypeError):
459
        msg = "'%s' (value %s) is not a valid integer" % (key, target[key])
460
        raise errors.TypeEnforcementError(msg)
461

    
462

    
463
def IsProcessAlive(pid):
464
  """Check if a given pid exists on the system.
465

466
  @note: zombie status is not handled, so zombie processes
467
      will be returned as alive
468
  @type pid: int
469
  @param pid: the process ID to check
470
  @rtype: boolean
471
  @return: True if the process exists
472

473
  """
474
  if pid <= 0:
475
    return False
476

    
477
  try:
478
    os.stat("/proc/%d/status" % pid)
479
    return True
480
  except EnvironmentError, err:
481
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
482
      return False
483
    raise
484

    
485

    
486
def ReadPidFile(pidfile):
487
  """Read a pid from a file.
488

489
  @type  pidfile: string
490
  @param pidfile: path to the file containing the pid
491
  @rtype: int
492
  @return: The process id, if the file exists and contains a valid PID,
493
           otherwise 0
494

495
  """
496
  try:
497
    pf = open(pidfile, 'r')
498
  except EnvironmentError, err:
499
    if err.errno != errno.ENOENT:
500
      logging.exception("Can't read pid file?!")
501
    return 0
502

    
503
  try:
504
    pid = int(pf.read())
505
  except ValueError, err:
506
    logging.info("Can't parse pid file contents", exc_info=True)
507
    return 0
508

    
509
  return pid
510

    
511

    
512
def MatchNameComponent(key, name_list):
513
  """Try to match a name against a list.
514

515
  This function will try to match a name like test1 against a list
516
  like C{['test1.example.com', 'test2.example.com', ...]}. Against
517
  this list, I{'test1'} as well as I{'test1.example'} will match, but
518
  not I{'test1.ex'}. A multiple match will be considered as no match
519
  at all (e.g. I{'test1'} against C{['test1.example.com',
520
  'test1.example.org']}).
521

522
  @type key: str
523
  @param key: the name to be searched
524
  @type name_list: list
525
  @param name_list: the list of strings against which to search the key
526

527
  @rtype: None or str
528
  @return: None if there is no match I{or} if there are multiple matches,
529
      otherwise the element from the list which matches
530

531
  """
532
  mo = re.compile("^%s(\..*)?$" % re.escape(key))
533
  names_filtered = [name for name in name_list if mo.match(name) is not None]
534
  if len(names_filtered) != 1:
535
    return None
536
  return names_filtered[0]
537

    
538

    
539
class HostInfo:
540
  """Class implementing resolver and hostname functionality
541

542
  """
543
  def __init__(self, name=None):
544
    """Initialize the host name object.
545

546
    If the name argument is not passed, it will use this system's
547
    name.
548

549
    """
550
    if name is None:
551
      name = self.SysName()
552

    
553
    self.query = name
554
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
555
    self.ip = self.ipaddrs[0]
556

    
557
  def ShortName(self):
558
    """Returns the hostname without domain.
559

560
    """
561
    return self.name.split('.')[0]
562

    
563
  @staticmethod
564
  def SysName():
565
    """Return the current system's name.
566

567
    This is simply a wrapper over C{socket.gethostname()}.
568

569
    """
570
    return socket.gethostname()
571

    
572
  @staticmethod
573
  def LookupHostname(hostname):
574
    """Look up hostname
575

576
    @type hostname: str
577
    @param hostname: hostname to look up
578

579
    @rtype: tuple
580
    @return: a tuple (name, aliases, ipaddrs) as returned by
581
        C{socket.gethostbyname_ex}
582
    @raise errors.ResolverError: in case of errors in resolving
583

584
    """
585
    try:
586
      result = socket.gethostbyname_ex(hostname)
587
    except socket.gaierror, err:
588
      # hostname not found in DNS
589
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
590

    
591
    return result
592

    
593

    
594
def ListVolumeGroups():
595
  """List volume groups and their size
596

597
  @rtype: dict
598
  @return:
599
       Dictionary with keys volume name and values
600
       the size of the volume
601

602
  """
603
  command = "vgs --noheadings --units m --nosuffix -o name,size"
604
  result = RunCmd(command)
605
  retval = {}
606
  if result.failed:
607
    return retval
608

    
609
  for line in result.stdout.splitlines():
610
    try:
611
      name, size = line.split()
612
      size = int(float(size))
613
    except (IndexError, ValueError), err:
614
      logging.error("Invalid output from vgs (%s): %s", err, line)
615
      continue
616

    
617
    retval[name] = size
618

    
619
  return retval
620

    
621

    
622
def BridgeExists(bridge):
623
  """Check whether the given bridge exists in the system
624

625
  @type bridge: str
626
  @param bridge: the bridge name to check
627
  @rtype: boolean
628
  @return: True if it does
629

630
  """
631
  return os.path.isdir("/sys/class/net/%s/bridge" % bridge)
632

    
633

    
634
def NiceSort(name_list):
635
  """Sort a list of strings based on digit and non-digit groupings.
636

637
  Given a list of names C{['a1', 'a10', 'a11', 'a2']} this function
638
  will sort the list in the logical order C{['a1', 'a2', 'a10',
639
  'a11']}.
640

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

645
  @type name_list: list
646
  @param name_list: the names to be sorted
647
  @rtype: list
648
  @return: a copy of the name list sorted with our algorithm
649

650
  """
651
  _SORTER_BASE = "(\D+|\d+)"
652
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
653
                                                  _SORTER_BASE, _SORTER_BASE,
654
                                                  _SORTER_BASE, _SORTER_BASE,
655
                                                  _SORTER_BASE, _SORTER_BASE)
656
  _SORTER_RE = re.compile(_SORTER_FULL)
657
  _SORTER_NODIGIT = re.compile("^\D*$")
658
  def _TryInt(val):
659
    """Attempts to convert a variable to integer."""
660
    if val is None or _SORTER_NODIGIT.match(val):
661
      return val
662
    rval = int(val)
663
    return rval
664

    
665
  to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
666
             for name in name_list]
667
  to_sort.sort()
668
  return [tup[1] for tup in to_sort]
669

    
670

    
671
def TryConvert(fn, val):
672
  """Try to convert a value ignoring errors.
673

674
  This function tries to apply function I{fn} to I{val}. If no
675
  C{ValueError} or C{TypeError} exceptions are raised, it will return
676
  the result, else it will return the original value. Any other
677
  exceptions are propagated to the caller.
678

679
  @type fn: callable
680
  @param fn: function to apply to the value
681
  @param val: the value to be converted
682
  @return: The converted value if the conversion was successful,
683
      otherwise the original value.
684

685
  """
686
  try:
687
    nv = fn(val)
688
  except (ValueError, TypeError):
689
    nv = val
690
  return nv
691

    
692

    
693
def IsValidIP(ip):
694
  """Verifies the syntax of an IPv4 address.
695

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

699
  @type ip: str
700
  @param ip: the address to be checked
701
  @rtype: a regular expression match object
702
  @return: a regular expression match object, or None if the
703
      address is not valid
704

705
  """
706
  unit = "(0|[1-9]\d{0,2})"
707
  #TODO: convert and return only boolean
708
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
709

    
710

    
711
def IsValidShellParam(word):
712
  """Verifies is the given word is safe from the shell's p.o.v.
713

714
  This means that we can pass this to a command via the shell and be
715
  sure that it doesn't alter the command line and is passed as such to
716
  the actual command.
717

718
  Note that we are overly restrictive here, in order to be on the safe
719
  side.
720

721
  @type word: str
722
  @param word: the word to check
723
  @rtype: boolean
724
  @return: True if the word is 'safe'
725

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

    
729

    
730
def BuildShellCmd(template, *args):
731
  """Build a safe shell command line from the given arguments.
732

733
  This function will check all arguments in the args list so that they
734
  are valid shell parameters (i.e. they don't contain shell
735
  metacharacters). If everything is ok, it will return the result of
736
  template % args.
737

738
  @type template: str
739
  @param template: the string holding the template for the
740
      string formatting
741
  @rtype: str
742
  @return: the expanded command line
743

744
  """
745
  for word in args:
746
    if not IsValidShellParam(word):
747
      raise errors.ProgrammerError("Shell argument '%s' contains"
748
                                   " invalid characters" % word)
749
  return template % args
750

    
751

    
752
def FormatUnit(value, units):
753
  """Formats an incoming number of MiB with the appropriate unit.
754

755
  @type value: int
756
  @param value: integer representing the value in MiB (1048576)
757
  @type units: char
758
  @param units: the type of formatting we should do:
759
      - 'h' for automatic scaling
760
      - 'm' for MiBs
761
      - 'g' for GiBs
762
      - 't' for TiBs
763
  @rtype: str
764
  @return: the formatted value (with suffix)
765

766
  """
767
  if units not in ('m', 'g', 't', 'h'):
768
    raise errors.ProgrammerError("Invalid unit specified '%s'" % str(units))
769

    
770
  suffix = ''
771

    
772
  if units == 'm' or (units == 'h' and value < 1024):
773
    if units == 'h':
774
      suffix = 'M'
775
    return "%d%s" % (round(value, 0), suffix)
776

    
777
  elif units == 'g' or (units == 'h' and value < (1024 * 1024)):
778
    if units == 'h':
779
      suffix = 'G'
780
    return "%0.1f%s" % (round(float(value) / 1024, 1), suffix)
781

    
782
  else:
783
    if units == 'h':
784
      suffix = 'T'
785
    return "%0.1f%s" % (round(float(value) / 1024 / 1024, 1), suffix)
786

    
787

    
788
def ParseUnit(input_string):
789
  """Tries to extract number and scale from the given string.
790

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

795
  """
796
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', str(input_string))
797
  if not m:
798
    raise errors.UnitParseError("Invalid format")
799

    
800
  value = float(m.groups()[0])
801

    
802
  unit = m.groups()[1]
803
  if unit:
804
    lcunit = unit.lower()
805
  else:
806
    lcunit = 'm'
807

    
808
  if lcunit in ('m', 'mb', 'mib'):
809
    # Value already in MiB
810
    pass
811

    
812
  elif lcunit in ('g', 'gb', 'gib'):
813
    value *= 1024
814

    
815
  elif lcunit in ('t', 'tb', 'tib'):
816
    value *= 1024 * 1024
817

    
818
  else:
819
    raise errors.UnitParseError("Unknown unit: %s" % unit)
820

    
821
  # Make sure we round up
822
  if int(value) < value:
823
    value += 1
824

    
825
  # Round up to the next multiple of 4
826
  value = int(value)
827
  if value % 4:
828
    value += 4 - value % 4
829

    
830
  return value
831

    
832

    
833
def AddAuthorizedKey(file_name, key):
834
  """Adds an SSH public key to an authorized_keys file.
835

836
  @type file_name: str
837
  @param file_name: path to authorized_keys file
838
  @type key: str
839
  @param key: string containing key
840

841
  """
842
  key_fields = key.split()
843

    
844
  f = open(file_name, 'a+')
845
  try:
846
    nl = True
847
    for line in f:
848
      # Ignore whitespace changes
849
      if line.split() == key_fields:
850
        break
851
      nl = line.endswith('\n')
852
    else:
853
      if not nl:
854
        f.write("\n")
855
      f.write(key.rstrip('\r\n'))
856
      f.write("\n")
857
      f.flush()
858
  finally:
859
    f.close()
860

    
861

    
862
def RemoveAuthorizedKey(file_name, key):
863
  """Removes an SSH public key from an authorized_keys file.
864

865
  @type file_name: str
866
  @param file_name: path to authorized_keys file
867
  @type key: str
868
  @param key: string containing key
869

870
  """
871
  key_fields = key.split()
872

    
873
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
874
  try:
875
    out = os.fdopen(fd, 'w')
876
    try:
877
      f = open(file_name, 'r')
878
      try:
879
        for line in f:
880
          # Ignore whitespace changes while comparing lines
881
          if line.split() != key_fields:
882
            out.write(line)
883

    
884
        out.flush()
885
        os.rename(tmpname, file_name)
886
      finally:
887
        f.close()
888
    finally:
889
      out.close()
890
  except:
891
    RemoveFile(tmpname)
892
    raise
893

    
894

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

898
  @type file_name: str
899
  @param file_name: path to the file to modify (usually C{/etc/hosts})
900
  @type ip: str
901
  @param ip: the IP address
902
  @type hostname: str
903
  @param hostname: the hostname to be added
904
  @type aliases: list
905
  @param aliases: the list of aliases to add for the hostname
906

907
  """
908
  # FIXME: use WriteFile + fn rather than duplicating its efforts
909
  # Ensure aliases are unique
910
  aliases = UniqueSequence([hostname] + aliases)[1:]
911

    
912
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
913
  try:
914
    out = os.fdopen(fd, 'w')
915
    try:
916
      f = open(file_name, 'r')
917
      try:
918
        for line in f:
919
          fields = line.split()
920
          if fields and not fields[0].startswith('#') and ip == fields[0]:
921
            continue
922
          out.write(line)
923

    
924
        out.write("%s\t%s" % (ip, hostname))
925
        if aliases:
926
          out.write(" %s" % ' '.join(aliases))
927
        out.write('\n')
928

    
929
        out.flush()
930
        os.fsync(out)
931
        os.chmod(tmpname, 0644)
932
        os.rename(tmpname, file_name)
933
      finally:
934
        f.close()
935
    finally:
936
      out.close()
937
  except:
938
    RemoveFile(tmpname)
939
    raise
940

    
941

    
942
def AddHostToEtcHosts(hostname):
943
  """Wrapper around SetEtcHostsEntry.
944

945
  @type hostname: str
946
  @param hostname: a hostname that will be resolved and added to
947
      L{constants.ETC_HOSTS}
948

949
  """
950
  hi = HostInfo(name=hostname)
951
  SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])
952

    
953

    
954
def RemoveEtcHostsEntry(file_name, hostname):
955
  """Removes a hostname from /etc/hosts.
956

957
  IP addresses without names are removed from the file.
958

959
  @type file_name: str
960
  @param file_name: path to the file to modify (usually C{/etc/hosts})
961
  @type hostname: str
962
  @param hostname: the hostname to be removed
963

964
  """
965
  # FIXME: use WriteFile + fn rather than duplicating its efforts
966
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
967
  try:
968
    out = os.fdopen(fd, 'w')
969
    try:
970
      f = open(file_name, 'r')
971
      try:
972
        for line in f:
973
          fields = line.split()
974
          if len(fields) > 1 and not fields[0].startswith('#'):
975
            names = fields[1:]
976
            if hostname in names:
977
              while hostname in names:
978
                names.remove(hostname)
979
              if names:
980
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
981
              continue
982

    
983
          out.write(line)
984

    
985
        out.flush()
986
        os.fsync(out)
987
        os.chmod(tmpname, 0644)
988
        os.rename(tmpname, file_name)
989
      finally:
990
        f.close()
991
    finally:
992
      out.close()
993
  except:
994
    RemoveFile(tmpname)
995
    raise
996

    
997

    
998
def RemoveHostFromEtcHosts(hostname):
999
  """Wrapper around RemoveEtcHostsEntry.
1000

1001
  @type hostname: str
1002
  @param hostname: hostname that will be resolved and its
1003
      full and shot name will be removed from
1004
      L{constants.ETC_HOSTS}
1005

1006
  """
1007
  hi = HostInfo(name=hostname)
1008
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.name)
1009
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.ShortName())
1010

    
1011

    
1012
def CreateBackup(file_name):
1013
  """Creates a backup of a file.
1014

1015
  @type file_name: str
1016
  @param file_name: file to be backed up
1017
  @rtype: str
1018
  @return: the path to the newly created backup
1019
  @raise errors.ProgrammerError: for invalid file names
1020

1021
  """
1022
  if not os.path.isfile(file_name):
1023
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
1024
                                file_name)
1025

    
1026
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
1027
  dir_name = os.path.dirname(file_name)
1028

    
1029
  fsrc = open(file_name, 'rb')
1030
  try:
1031
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
1032
    fdst = os.fdopen(fd, 'wb')
1033
    try:
1034
      shutil.copyfileobj(fsrc, fdst)
1035
    finally:
1036
      fdst.close()
1037
  finally:
1038
    fsrc.close()
1039

    
1040
  return backup_name
1041

    
1042

    
1043
def ShellQuote(value):
1044
  """Quotes shell argument according to POSIX.
1045

1046
  @type value: str
1047
  @param value: the argument to be quoted
1048
  @rtype: str
1049
  @return: the quoted value
1050

1051
  """
1052
  if _re_shell_unquoted.match(value):
1053
    return value
1054
  else:
1055
    return "'%s'" % value.replace("'", "'\\''")
1056

    
1057

    
1058
def ShellQuoteArgs(args):
1059
  """Quotes a list of shell arguments.
1060

1061
  @type args: list
1062
  @param args: list of arguments to be quoted
1063
  @rtype: str
1064
  @return: the quoted arguments concatenated with spaces
1065

1066
  """
1067
  return ' '.join([ShellQuote(i) for i in args])
1068

    
1069

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

1073
  Check if the given IP is reachable by doing attempting a TCP connect
1074
  to it.
1075

1076
  @type target: str
1077
  @param target: the IP or hostname to ping
1078
  @type port: int
1079
  @param port: the port to connect to
1080
  @type timeout: int
1081
  @param timeout: the timeout on the connection attempt
1082
  @type live_port_needed: boolean
1083
  @param live_port_needed: whether a closed port will cause the
1084
      function to return failure, as if there was a timeout
1085
  @type source: str or None
1086
  @param source: if specified, will cause the connect to be made
1087
      from this specific source address; failures to bind other
1088
      than C{EADDRNOTAVAIL} will be ignored
1089

1090
  """
1091
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
1092

    
1093
  success = False
1094

    
1095
  if source is not None:
1096
    try:
1097
      sock.bind((source, 0))
1098
    except socket.error, (errcode, _):
1099
      if errcode == errno.EADDRNOTAVAIL:
1100
        success = False
1101

    
1102
  sock.settimeout(timeout)
1103

    
1104
  try:
1105
    sock.connect((target, port))
1106
    sock.close()
1107
    success = True
1108
  except socket.timeout:
1109
    success = False
1110
  except socket.error, (errcode, errstring):
1111
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
1112

    
1113
  return success
1114

    
1115

    
1116
def OwnIpAddress(address):
1117
  """Check if the current host has the the given IP address.
1118

1119
  Currently this is done by TCP-pinging the address from the loopback
1120
  address.
1121

1122
  @type address: string
1123
  @param address: the address to check
1124
  @rtype: bool
1125
  @return: True if we own the address
1126

1127
  """
1128
  return TcpPing(address, constants.DEFAULT_NODED_PORT,
1129
                 source=constants.LOCALHOST_IP_ADDRESS)
1130

    
1131

    
1132
def ListVisibleFiles(path):
1133
  """Returns a list of visible files in a directory.
1134

1135
  @type path: str
1136
  @param path: the directory to enumerate
1137
  @rtype: list
1138
  @return: the list of all files not starting with a dot
1139

1140
  """
1141
  files = [i for i in os.listdir(path) if not i.startswith(".")]
1142
  files.sort()
1143
  return files
1144

    
1145

    
1146
def GetHomeDir(user, default=None):
1147
  """Try to get the homedir of the given user.
1148

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

1153
  """
1154
  try:
1155
    if isinstance(user, basestring):
1156
      result = pwd.getpwnam(user)
1157
    elif isinstance(user, (int, long)):
1158
      result = pwd.getpwuid(user)
1159
    else:
1160
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
1161
                                   type(user))
1162
  except KeyError:
1163
    return default
1164
  return result.pw_dir
1165

    
1166

    
1167
def NewUUID():
1168
  """Returns a random UUID.
1169

1170
  @note: This is a Linux-specific method as it uses the /proc
1171
      filesystem.
1172
  @rtype: str
1173

1174
  """
1175
  f = open("/proc/sys/kernel/random/uuid", "r")
1176
  try:
1177
    return f.read(128).rstrip("\n")
1178
  finally:
1179
    f.close()
1180

    
1181

    
1182
def GenerateSecret():
1183
  """Generates a random secret.
1184

1185
  This will generate a pseudo-random secret, and return its sha digest
1186
  (so that it can be used where an ASCII string is needed).
1187

1188
  @rtype: str
1189
  @return: a sha1 hexdigest of a block of 64 random bytes
1190

1191
  """
1192
  return sha1(os.urandom(64)).hexdigest()
1193

    
1194

    
1195
def EnsureDirs(dirs):
1196
  """Make required directories, if they don't exist.
1197

1198
  @param dirs: list of tuples (dir_name, dir_mode)
1199
  @type dirs: list of (string, integer)
1200

1201
  """
1202
  for dir_name, dir_mode in dirs:
1203
    try:
1204
      os.mkdir(dir_name, dir_mode)
1205
    except EnvironmentError, err:
1206
      if err.errno != errno.EEXIST:
1207
        raise errors.GenericError("Cannot create needed directory"
1208
                                  " '%s': %s" % (dir_name, err))
1209
    if not os.path.isdir(dir_name):
1210
      raise errors.GenericError("%s is not a directory" % dir_name)
1211

    
1212

    
1213
def ReadFile(file_name, size=None):
1214
  """Reads a file.
1215

1216
  @type size: None or int
1217
  @param size: Read at most size bytes
1218
  @rtype: str
1219
  @return: the (possibly partial) content of the file
1220

1221
  """
1222
  f = open(file_name, "r")
1223
  try:
1224
    if size is None:
1225
      return f.read()
1226
    else:
1227
      return f.read(size)
1228
  finally:
1229
    f.close()
1230

    
1231

    
1232
def WriteFile(file_name, fn=None, data=None,
1233
              mode=None, uid=-1, gid=-1,
1234
              atime=None, mtime=None, close=True,
1235
              dry_run=False, backup=False,
1236
              prewrite=None, postwrite=None):
1237
  """(Over)write a file atomically.
1238

1239
  The file_name and either fn (a function taking one argument, the
1240
  file descriptor, and which should write the data to it) or data (the
1241
  contents of the file) must be passed. The other arguments are
1242
  optional and allow setting the file mode, owner and group, and the
1243
  mtime/atime of the file.
1244

1245
  If the function doesn't raise an exception, it has succeeded and the
1246
  target file has the new contents. If the function has raised an
1247
  exception, an existing target file should be unmodified and the
1248
  temporary file should be removed.
1249

1250
  @type file_name: str
1251
  @param file_name: the target filename
1252
  @type fn: callable
1253
  @param fn: content writing function, called with
1254
      file descriptor as parameter
1255
  @type data: str
1256
  @param data: contents of the file
1257
  @type mode: int
1258
  @param mode: file mode
1259
  @type uid: int
1260
  @param uid: the owner of the file
1261
  @type gid: int
1262
  @param gid: the group of the file
1263
  @type atime: int
1264
  @param atime: a custom access time to be set on the file
1265
  @type mtime: int
1266
  @param mtime: a custom modification time to be set on the file
1267
  @type close: boolean
1268
  @param close: whether to close file after writing it
1269
  @type prewrite: callable
1270
  @param prewrite: function to be called before writing content
1271
  @type postwrite: callable
1272
  @param postwrite: function to be called after writing content
1273

1274
  @rtype: None or int
1275
  @return: None if the 'close' parameter evaluates to True,
1276
      otherwise the file descriptor
1277

1278
  @raise errors.ProgrammerError: if any of the arguments are not valid
1279

1280
  """
1281
  if not os.path.isabs(file_name):
1282
    raise errors.ProgrammerError("Path passed to WriteFile is not"
1283
                                 " absolute: '%s'" % file_name)
1284

    
1285
  if [fn, data].count(None) != 1:
1286
    raise errors.ProgrammerError("fn or data required")
1287

    
1288
  if [atime, mtime].count(None) == 1:
1289
    raise errors.ProgrammerError("Both atime and mtime must be either"
1290
                                 " set or None")
1291

    
1292
  if backup and not dry_run and os.path.isfile(file_name):
1293
    CreateBackup(file_name)
1294

    
1295
  dir_name, base_name = os.path.split(file_name)
1296
  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
1297
  do_remove = True
1298
  # here we need to make sure we remove the temp file, if any error
1299
  # leaves it in place
1300
  try:
1301
    if uid != -1 or gid != -1:
1302
      os.chown(new_name, uid, gid)
1303
    if mode:
1304
      os.chmod(new_name, mode)
1305
    if callable(prewrite):
1306
      prewrite(fd)
1307
    if data is not None:
1308
      os.write(fd, data)
1309
    else:
1310
      fn(fd)
1311
    if callable(postwrite):
1312
      postwrite(fd)
1313
    os.fsync(fd)
1314
    if atime is not None and mtime is not None:
1315
      os.utime(new_name, (atime, mtime))
1316
    if not dry_run:
1317
      os.rename(new_name, file_name)
1318
      do_remove = False
1319
  finally:
1320
    if close:
1321
      os.close(fd)
1322
      result = None
1323
    else:
1324
      result = fd
1325
    if do_remove:
1326
      RemoveFile(new_name)
1327

    
1328
  return result
1329

    
1330

    
1331
def FirstFree(seq, base=0):
1332
  """Returns the first non-existing integer from seq.
1333

1334
  The seq argument should be a sorted list of positive integers. The
1335
  first time the index of an element is smaller than the element
1336
  value, the index will be returned.
1337

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

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

1343
  @type seq: sequence
1344
  @param seq: the sequence to be analyzed.
1345
  @type base: int
1346
  @param base: use this value as the base index of the sequence
1347
  @rtype: int
1348
  @return: the first non-used index in the sequence
1349

1350
  """
1351
  for idx, elem in enumerate(seq):
1352
    assert elem >= base, "Passed element is higher than base offset"
1353
    if elem > idx + base:
1354
      # idx is not used
1355
      return idx + base
1356
  return None
1357

    
1358

    
1359
def all(seq, pred=bool):
1360
  "Returns True if pred(x) is True for every element in the iterable"
1361
  for _ in itertools.ifilterfalse(pred, seq):
1362
    return False
1363
  return True
1364

    
1365

    
1366
def any(seq, pred=bool):
1367
  "Returns True if pred(x) is True for at least one element in the iterable"
1368
  for _ in itertools.ifilter(pred, seq):
1369
    return True
1370
  return False
1371

    
1372

    
1373
def UniqueSequence(seq):
1374
  """Returns a list with unique elements.
1375

1376
  Element order is preserved.
1377

1378
  @type seq: sequence
1379
  @param seq: the sequence with the source elements
1380
  @rtype: list
1381
  @return: list of unique elements from seq
1382

1383
  """
1384
  seen = set()
1385
  return [i for i in seq if i not in seen and not seen.add(i)]
1386

    
1387

    
1388
def IsValidMac(mac):
1389
  """Predicate to check if a MAC address is valid.
1390

1391
  Checks whether the supplied MAC address is formally correct, only
1392
  accepts colon separated format.
1393

1394
  @type mac: str
1395
  @param mac: the MAC to be validated
1396
  @rtype: boolean
1397
  @return: True is the MAC seems valid
1398

1399
  """
1400
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
1401
  return mac_check.match(mac) is not None
1402

    
1403

    
1404
def TestDelay(duration):
1405
  """Sleep for a fixed amount of time.
1406

1407
  @type duration: float
1408
  @param duration: the sleep duration
1409
  @rtype: boolean
1410
  @return: False for negative value, True otherwise
1411

1412
  """
1413
  if duration < 0:
1414
    return False
1415
  time.sleep(duration)
1416
  return True
1417

    
1418

    
1419
def _CloseFDNoErr(fd, retries=5):
1420
  """Close a file descriptor ignoring errors.
1421

1422
  @type fd: int
1423
  @param fd: the file descriptor
1424
  @type retries: int
1425
  @param retries: how many retries to make, in case we get any
1426
      other error than EBADF
1427

1428
  """
1429
  try:
1430
    os.close(fd)
1431
  except OSError, err:
1432
    if err.errno != errno.EBADF:
1433
      if retries > 0:
1434
        _CloseFDNoErr(fd, retries - 1)
1435
    # else either it's closed already or we're out of retries, so we
1436
    # ignore this and go on
1437

    
1438

    
1439
def CloseFDs(noclose_fds=None):
1440
  """Close file descriptors.
1441

1442
  This closes all file descriptors above 2 (i.e. except
1443
  stdin/out/err).
1444

1445
  @type noclose_fds: list or None
1446
  @param noclose_fds: if given, it denotes a list of file descriptor
1447
      that should not be closed
1448

1449
  """
1450
  # Default maximum for the number of available file descriptors.
1451
  if 'SC_OPEN_MAX' in os.sysconf_names:
1452
    try:
1453
      MAXFD = os.sysconf('SC_OPEN_MAX')
1454
      if MAXFD < 0:
1455
        MAXFD = 1024
1456
    except OSError:
1457
      MAXFD = 1024
1458
  else:
1459
    MAXFD = 1024
1460
  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1461
  if (maxfd == resource.RLIM_INFINITY):
1462
    maxfd = MAXFD
1463

    
1464
  # Iterate through and close all file descriptors (except the standard ones)
1465
  for fd in range(3, maxfd):
1466
    if noclose_fds and fd in noclose_fds:
1467
      continue
1468
    _CloseFDNoErr(fd)
1469

    
1470

    
1471
def Daemonize(logfile):
1472
  """Daemonize the current process.
1473

1474
  This detaches the current process from the controlling terminal and
1475
  runs it in the background as a daemon.
1476

1477
  @type logfile: str
1478
  @param logfile: the logfile to which we should redirect stdout/stderr
1479
  @rtype: int
1480
  @return: the value zero
1481

1482
  """
1483
  UMASK = 077
1484
  WORKDIR = "/"
1485

    
1486
  # this might fail
1487
  pid = os.fork()
1488
  if (pid == 0):  # The first child.
1489
    os.setsid()
1490
    # this might fail
1491
    pid = os.fork() # Fork a second child.
1492
    if (pid == 0):  # The second child.
1493
      os.chdir(WORKDIR)
1494
      os.umask(UMASK)
1495
    else:
1496
      # exit() or _exit()?  See below.
1497
      os._exit(0) # Exit parent (the first child) of the second child.
1498
  else:
1499
    os._exit(0) # Exit parent of the first child.
1500

    
1501
  for fd in range(3):
1502
    _CloseFDNoErr(fd)
1503
  i = os.open("/dev/null", os.O_RDONLY) # stdin
1504
  assert i == 0, "Can't close/reopen stdin"
1505
  i = os.open(logfile, os.O_WRONLY|os.O_CREAT|os.O_APPEND, 0600) # stdout
1506
  assert i == 1, "Can't close/reopen stdout"
1507
  # Duplicate standard output to standard error.
1508
  os.dup2(1, 2)
1509
  return 0
1510

    
1511

    
1512
def DaemonPidFileName(name):
1513
  """Compute a ganeti pid file absolute path
1514

1515
  @type name: str
1516
  @param name: the daemon name
1517
  @rtype: str
1518
  @return: the full path to the pidfile corresponding to the given
1519
      daemon name
1520

1521
  """
1522
  return os.path.join(constants.RUN_GANETI_DIR, "%s.pid" % name)
1523

    
1524

    
1525
def WritePidFile(name):
1526
  """Write the current process pidfile.
1527

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

1530
  @type name: str
1531
  @param name: the daemon name to use
1532
  @raise errors.GenericError: if the pid file already exists and
1533
      points to a live process
1534

1535
  """
1536
  pid = os.getpid()
1537
  pidfilename = DaemonPidFileName(name)
1538
  if IsProcessAlive(ReadPidFile(pidfilename)):
1539
    raise errors.GenericError("%s contains a live process" % pidfilename)
1540

    
1541
  WriteFile(pidfilename, data="%d\n" % pid)
1542

    
1543

    
1544
def RemovePidFile(name):
1545
  """Remove the current process pidfile.
1546

1547
  Any errors are ignored.
1548

1549
  @type name: str
1550
  @param name: the daemon name used to derive the pidfile name
1551

1552
  """
1553
  pidfilename = DaemonPidFileName(name)
1554
  # TODO: we could check here that the file contains our pid
1555
  try:
1556
    RemoveFile(pidfilename)
1557
  except:
1558
    pass
1559

    
1560

    
1561
def KillProcess(pid, signal_=signal.SIGTERM, timeout=30,
1562
                waitpid=False):
1563
  """Kill a process given by its pid.
1564

1565
  @type pid: int
1566
  @param pid: The PID to terminate.
1567
  @type signal_: int
1568
  @param signal_: The signal to send, by default SIGTERM
1569
  @type timeout: int
1570
  @param timeout: The timeout after which, if the process is still alive,
1571
                  a SIGKILL will be sent. If not positive, no such checking
1572
                  will be done
1573
  @type waitpid: boolean
1574
  @param waitpid: If true, we should waitpid on this process after
1575
      sending signals, since it's our own child and otherwise it
1576
      would remain as zombie
1577

1578
  """
1579
  def _helper(pid, signal_, wait):
1580
    """Simple helper to encapsulate the kill/waitpid sequence"""
1581
    os.kill(pid, signal_)
1582
    if wait:
1583
      try:
1584
        os.waitpid(pid, os.WNOHANG)
1585
      except OSError:
1586
        pass
1587

    
1588
  if pid <= 0:
1589
    # kill with pid=0 == suicide
1590
    raise errors.ProgrammerError("Invalid pid given '%s'" % pid)
1591

    
1592
  if not IsProcessAlive(pid):
1593
    return
1594
  _helper(pid, signal_, waitpid)
1595
  if timeout <= 0:
1596
    return
1597

    
1598
  # Wait up to $timeout seconds
1599
  end = time.time() + timeout
1600
  wait = 0.01
1601
  while time.time() < end and IsProcessAlive(pid):
1602
    try:
1603
      (result_pid, _) = os.waitpid(pid, os.WNOHANG)
1604
      if result_pid > 0:
1605
        break
1606
    except OSError:
1607
      pass
1608
    time.sleep(wait)
1609
    # Make wait time longer for next try
1610
    if wait < 0.1:
1611
      wait *= 1.5
1612

    
1613
  if IsProcessAlive(pid):
1614
    # Kill process if it's still alive
1615
    _helper(pid, signal.SIGKILL, waitpid)
1616

    
1617

    
1618
def FindFile(name, search_path, test=os.path.exists):
1619
  """Look for a filesystem object in a given path.
1620

1621
  This is an abstract method to search for filesystem object (files,
1622
  dirs) under a given search path.
1623

1624
  @type name: str
1625
  @param name: the name to look for
1626
  @type search_path: str
1627
  @param search_path: location to start at
1628
  @type test: callable
1629
  @param test: a function taking one argument that should return True
1630
      if the a given object is valid; the default value is
1631
      os.path.exists, causing only existing files to be returned
1632
  @rtype: str or None
1633
  @return: full path to the object if found, None otherwise
1634

1635
  """
1636
  # validate the filename mask
1637
  if constants.EXT_PLUGIN_MASK.match(name) is None:
1638
    logging.critical("Invalid value passed for external script name: '%s'",
1639
                     name)
1640
    return None
1641

    
1642
  for dir_name in search_path:
1643
    item_name = os.path.sep.join([dir_name, name])
1644
    # check the user test and that we're indeed resolving to the given
1645
    # basename
1646
    if test(item_name) and os.path.basename(item_name) == name:
1647
      return item_name
1648
  return None
1649

    
1650

    
1651
def CheckVolumeGroupSize(vglist, vgname, minsize):
1652
  """Checks if the volume group list is valid.
1653

1654
  The function will check if a given volume group is in the list of
1655
  volume groups and has a minimum size.
1656

1657
  @type vglist: dict
1658
  @param vglist: dictionary of volume group names and their size
1659
  @type vgname: str
1660
  @param vgname: the volume group we should check
1661
  @type minsize: int
1662
  @param minsize: the minimum size we accept
1663
  @rtype: None or str
1664
  @return: None for success, otherwise the error message
1665

1666
  """
1667
  vgsize = vglist.get(vgname, None)
1668
  if vgsize is None:
1669
    return "volume group '%s' missing" % vgname
1670
  elif vgsize < minsize:
1671
    return ("volume group '%s' too small (%s MiB required, %d MiB found)" %
1672
            (vgname, minsize, vgsize))
1673
  return None
1674

    
1675

    
1676
def SplitTime(value):
1677
  """Splits time as floating point number into a tuple.
1678

1679
  @param value: Time in seconds
1680
  @type value: int or float
1681
  @return: Tuple containing (seconds, microseconds)
1682

1683
  """
1684
  (seconds, microseconds) = divmod(int(value * 1000000), 1000000)
1685

    
1686
  assert 0 <= seconds, \
1687
    "Seconds must be larger than or equal to 0, but are %s" % seconds
1688
  assert 0 <= microseconds <= 999999, \
1689
    "Microseconds must be 0-999999, but are %s" % microseconds
1690

    
1691
  return (int(seconds), int(microseconds))
1692

    
1693

    
1694
def MergeTime(timetuple):
1695
  """Merges a tuple into time as a floating point number.
1696

1697
  @param timetuple: Time as tuple, (seconds, microseconds)
1698
  @type timetuple: tuple
1699
  @return: Time as a floating point number expressed in seconds
1700

1701
  """
1702
  (seconds, microseconds) = timetuple
1703

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

    
1709
  return float(seconds) + (float(microseconds) * 0.000001)
1710

    
1711

    
1712
def GetNodeDaemonPort():
1713
  """Get the node daemon port for this cluster.
1714

1715
  Note that this routine does not read a ganeti-specific file, but
1716
  instead uses C{socket.getservbyname} to allow pre-customization of
1717
  this parameter outside of Ganeti.
1718

1719
  @rtype: int
1720

1721
  """
1722
  try:
1723
    port = socket.getservbyname("ganeti-noded", "tcp")
1724
  except socket.error:
1725
    port = constants.DEFAULT_NODED_PORT
1726

    
1727
  return port
1728

    
1729

    
1730
def SetupLogging(logfile, debug=False, stderr_logging=False, program="",
1731
                 multithreaded=False):
1732
  """Configures the logging module.
1733

1734
  @type logfile: str
1735
  @param logfile: the filename to which we should log
1736
  @type debug: boolean
1737
  @param debug: whether to enable debug messages too or
1738
      only those at C{INFO} and above level
1739
  @type stderr_logging: boolean
1740
  @param stderr_logging: whether we should also log to the standard error
1741
  @type program: str
1742
  @param program: the name under which we should log messages
1743
  @type multithreaded: boolean
1744
  @param multithreaded: if True, will add the thread name to the log file
1745
  @raise EnvironmentError: if we can't open the log file and
1746
      stderr logging is disabled
1747

1748
  """
1749
  fmt = "%(asctime)s: " + program + " pid=%(process)d"
1750
  if multithreaded:
1751
    fmt += "/%(threadName)s"
1752
  if debug:
1753
    fmt += " %(module)s:%(lineno)s"
1754
  fmt += " %(levelname)s %(message)s"
1755
  formatter = logging.Formatter(fmt)
1756

    
1757
  root_logger = logging.getLogger("")
1758
  root_logger.setLevel(logging.NOTSET)
1759

    
1760
  # Remove all previously setup handlers
1761
  for handler in root_logger.handlers:
1762
    handler.close()
1763
    root_logger.removeHandler(handler)
1764

    
1765
  if stderr_logging:
1766
    stderr_handler = logging.StreamHandler()
1767
    stderr_handler.setFormatter(formatter)
1768
    if debug:
1769
      stderr_handler.setLevel(logging.NOTSET)
1770
    else:
1771
      stderr_handler.setLevel(logging.CRITICAL)
1772
    root_logger.addHandler(stderr_handler)
1773

    
1774
  # this can fail, if the logging directories are not setup or we have
1775
  # a permisssion problem; in this case, it's best to log but ignore
1776
  # the error if stderr_logging is True, and if false we re-raise the
1777
  # exception since otherwise we could run but without any logs at all
1778
  try:
1779
    logfile_handler = logging.FileHandler(logfile)
1780
    logfile_handler.setFormatter(formatter)
1781
    if debug:
1782
      logfile_handler.setLevel(logging.DEBUG)
1783
    else:
1784
      logfile_handler.setLevel(logging.INFO)
1785
    root_logger.addHandler(logfile_handler)
1786
  except EnvironmentError:
1787
    if stderr_logging:
1788
      logging.exception("Failed to enable logging to file '%s'", logfile)
1789
    else:
1790
      # we need to re-raise the exception
1791
      raise
1792

    
1793

    
1794
def TailFile(fname, lines=20):
1795
  """Return the last lines from a file.
1796

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

1801
  @param fname: the file name
1802
  @type lines: int
1803
  @param lines: the (maximum) number of lines to return
1804

1805
  """
1806
  fd = open(fname, "r")
1807
  try:
1808
    fd.seek(0, 2)
1809
    pos = fd.tell()
1810
    pos = max(0, pos-4096)
1811
    fd.seek(pos, 0)
1812
    raw_data = fd.read()
1813
  finally:
1814
    fd.close()
1815

    
1816
  rows = raw_data.splitlines()
1817
  return rows[-lines:]
1818

    
1819

    
1820
def SafeEncode(text):
1821
  """Return a 'safe' version of a source string.
1822

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

1832
  @type text: str or unicode
1833
  @param text: input data
1834
  @rtype: str
1835
  @return: a safe version of text
1836

1837
  """
1838
  if isinstance(text, unicode):
1839
    # only if unicode; if str already, we handle it below
1840
    text = text.encode('ascii', 'backslashreplace')
1841
  resu = ""
1842
  for char in text:
1843
    c = ord(char)
1844
    if char  == '\t':
1845
      resu += r'\t'
1846
    elif char == '\n':
1847
      resu += r'\n'
1848
    elif char == '\r':
1849
      resu += r'\'r'
1850
    elif c < 32 or c >= 127: # non-printable
1851
      resu += "\\x%02x" % (c & 0xff)
1852
    else:
1853
      resu += char
1854
  return resu
1855

    
1856

    
1857
def CommaJoin(names):
1858
  """Nicely join a set of identifiers.
1859

1860
  @param names: set, list or tuple
1861
  @return: a string with the formatted results
1862

1863
  """
1864
  return ", ".join(["'%s'" % val for val in names])
1865

    
1866

    
1867
def LockedMethod(fn):
1868
  """Synchronized object access decorator.
1869

1870
  This decorator is intended to protect access to an object using the
1871
  object's own lock which is hardcoded to '_lock'.
1872

1873
  """
1874
  def _LockDebug(*args, **kwargs):
1875
    if debug_locks:
1876
      logging.debug(*args, **kwargs)
1877

    
1878
  def wrapper(self, *args, **kwargs):
1879
    assert hasattr(self, '_lock')
1880
    lock = self._lock
1881
    _LockDebug("Waiting for %s", lock)
1882
    lock.acquire()
1883
    try:
1884
      _LockDebug("Acquired %s", lock)
1885
      result = fn(self, *args, **kwargs)
1886
    finally:
1887
      _LockDebug("Releasing %s", lock)
1888
      lock.release()
1889
      _LockDebug("Released %s", lock)
1890
    return result
1891
  return wrapper
1892

    
1893

    
1894
def LockFile(fd):
1895
  """Locks a file using POSIX locks.
1896

1897
  @type fd: int
1898
  @param fd: the file descriptor we need to lock
1899

1900
  """
1901
  try:
1902
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
1903
  except IOError, err:
1904
    if err.errno == errno.EAGAIN:
1905
      raise errors.LockError("File already locked")
1906
    raise
1907

    
1908

    
1909
class FileLock(object):
1910
  """Utility class for file locks.
1911

1912
  """
1913
  def __init__(self, filename):
1914
    """Constructor for FileLock.
1915

1916
    This will open the file denoted by the I{filename} argument.
1917

1918
    @type filename: str
1919
    @param filename: path to the file to be locked
1920

1921
    """
1922
    self.filename = filename
1923
    self.fd = open(self.filename, "w")
1924

    
1925
  def __del__(self):
1926
    self.Close()
1927

    
1928
  def Close(self):
1929
    """Close the file and release the lock.
1930

1931
    """
1932
    if self.fd:
1933
      self.fd.close()
1934
      self.fd = None
1935

    
1936
  def _flock(self, flag, blocking, timeout, errmsg):
1937
    """Wrapper for fcntl.flock.
1938

1939
    @type flag: int
1940
    @param flag: operation flag
1941
    @type blocking: bool
1942
    @param blocking: whether the operation should be done in blocking mode.
1943
    @type timeout: None or float
1944
    @param timeout: for how long the operation should be retried (implies
1945
                    non-blocking mode).
1946
    @type errmsg: string
1947
    @param errmsg: error message in case operation fails.
1948

1949
    """
1950
    assert self.fd, "Lock was closed"
1951
    assert timeout is None or timeout >= 0, \
1952
      "If specified, timeout must be positive"
1953

    
1954
    if timeout is not None:
1955
      flag |= fcntl.LOCK_NB
1956
      timeout_end = time.time() + timeout
1957

    
1958
    # Blocking doesn't have effect with timeout
1959
    elif not blocking:
1960
      flag |= fcntl.LOCK_NB
1961
      timeout_end = None
1962

    
1963
    retry = True
1964
    while retry:
1965
      try:
1966
        fcntl.flock(self.fd, flag)
1967
        retry = False
1968
      except IOError, err:
1969
        if err.errno in (errno.EAGAIN, ):
1970
          if timeout_end is not None and time.time() < timeout_end:
1971
            # Wait before trying again
1972
            time.sleep(max(0.1, min(1.0, timeout)))
1973
          else:
1974
            raise errors.LockError(errmsg)
1975
        else:
1976
          logging.exception("fcntl.flock failed")
1977
          raise
1978

    
1979
  def Exclusive(self, blocking=False, timeout=None):
1980
    """Locks the file in exclusive mode.
1981

1982
    @type blocking: boolean
1983
    @param blocking: whether to block and wait until we
1984
        can lock the file or return immediately
1985
    @type timeout: int or None
1986
    @param timeout: if not None, the duration to wait for the lock
1987
        (in blocking mode)
1988

1989
    """
1990
    self._flock(fcntl.LOCK_EX, blocking, timeout,
1991
                "Failed to lock %s in exclusive mode" % self.filename)
1992

    
1993
  def Shared(self, blocking=False, timeout=None):
1994
    """Locks the file in shared mode.
1995

1996
    @type blocking: boolean
1997
    @param blocking: whether to block and wait until we
1998
        can lock the file or return immediately
1999
    @type timeout: int or None
2000
    @param timeout: if not None, the duration to wait for the lock
2001
        (in blocking mode)
2002

2003
    """
2004
    self._flock(fcntl.LOCK_SH, blocking, timeout,
2005
                "Failed to lock %s in shared mode" % self.filename)
2006

    
2007
  def Unlock(self, blocking=True, timeout=None):
2008
    """Unlocks the file.
2009

2010
    According to C{flock(2)}, unlocking can also be a nonblocking
2011
    operation::
2012

2013
      To make a non-blocking request, include LOCK_NB with any of the above
2014
      operations.
2015

2016
    @type blocking: boolean
2017
    @param blocking: whether to block and wait until we
2018
        can lock the file or return immediately
2019
    @type timeout: int or None
2020
    @param timeout: if not None, the duration to wait for the lock
2021
        (in blocking mode)
2022

2023
    """
2024
    self._flock(fcntl.LOCK_UN, blocking, timeout,
2025
                "Failed to unlock %s" % self.filename)
2026

    
2027

    
2028
class SignalHandler(object):
2029
  """Generic signal handler class.
2030

2031
  It automatically restores the original handler when deconstructed or
2032
  when L{Reset} is called. You can either pass your own handler
2033
  function in or query the L{called} attribute to detect whether the
2034
  signal was sent.
2035

2036
  @type signum: list
2037
  @ivar signum: the signals we handle
2038
  @type called: boolean
2039
  @ivar called: tracks whether any of the signals have been raised
2040

2041
  """
2042
  def __init__(self, signum):
2043
    """Constructs a new SignalHandler instance.
2044

2045
    @type signum: int or list of ints
2046
    @param signum: Single signal number or set of signal numbers
2047

2048
    """
2049
    if isinstance(signum, (int, long)):
2050
      self.signum = set([signum])
2051
    else:
2052
      self.signum = set(signum)
2053

    
2054
    self.called = False
2055

    
2056
    self._previous = {}
2057
    try:
2058
      for signum in self.signum:
2059
        # Setup handler
2060
        prev_handler = signal.signal(signum, self._HandleSignal)
2061
        try:
2062
          self._previous[signum] = prev_handler
2063
        except:
2064
          # Restore previous handler
2065
          signal.signal(signum, prev_handler)
2066
          raise
2067
    except:
2068
      # Reset all handlers
2069
      self.Reset()
2070
      # Here we have a race condition: a handler may have already been called,
2071
      # but there's not much we can do about it at this point.
2072
      raise
2073

    
2074
  def __del__(self):
2075
    self.Reset()
2076

    
2077
  def Reset(self):
2078
    """Restore previous handler.
2079

2080
    This will reset all the signals to their previous handlers.
2081

2082
    """
2083
    for signum, prev_handler in self._previous.items():
2084
      signal.signal(signum, prev_handler)
2085
      # If successful, remove from dict
2086
      del self._previous[signum]
2087

    
2088
  def Clear(self):
2089
    """Unsets the L{called} flag.
2090

2091
    This function can be used in case a signal may arrive several times.
2092

2093
    """
2094
    self.called = False
2095

    
2096
  def _HandleSignal(self, signum, frame):
2097
    """Actual signal handling function.
2098

2099
    """
2100
    # This is not nice and not absolutely atomic, but it appears to be the only
2101
    # solution in Python -- there are no atomic types.
2102
    self.called = True
2103

    
2104

    
2105
class FieldSet(object):
2106
  """A simple field set.
2107

2108
  Among the features are:
2109
    - checking if a string is among a list of static string or regex objects
2110
    - checking if a whole list of string matches
2111
    - returning the matching groups from a regex match
2112

2113
  Internally, all fields are held as regular expression objects.
2114

2115
  """
2116
  def __init__(self, *items):
2117
    self.items = [re.compile("^%s$" % value) for value in items]
2118

    
2119
  def Extend(self, other_set):
2120
    """Extend the field set with the items from another one"""
2121
    self.items.extend(other_set.items)
2122

    
2123
  def Matches(self, field):
2124
    """Checks if a field matches the current set
2125

2126
    @type field: str
2127
    @param field: the string to match
2128
    @return: either False or a regular expression match object
2129

2130
    """
2131
    for m in itertools.ifilter(None, (val.match(field) for val in self.items)):
2132
      return m
2133
    return False
2134

    
2135
  def NonMatching(self, items):
2136
    """Returns the list of fields not matching the current set
2137

2138
    @type items: list
2139
    @param items: the list of fields to check
2140
    @rtype: list
2141
    @return: list of non-matching fields
2142

2143
    """
2144
    return [val for val in items if not self.Matches(val)]