Add utils.CheckBEParams
[ganeti-local] / lib / utils.py
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 sys
31 import os
32 import sha
33 import time
34 import subprocess
35 import re
36 import socket
37 import tempfile
38 import shutil
39 import errno
40 import pwd
41 import itertools
42 import select
43 import fcntl
44 import resource
45 import logging
46 import signal
47
48 from cStringIO import StringIO
49
50 from ganeti import errors
51 from ganeti import constants
52
53
54 _locksheld = []
55 _re_shell_unquoted = re.compile('^[-.,=:/_+@A-Za-z0-9]+$')
56
57 debug = False
58 debug_locks = False
59
60 #: when set to True, L{RunCmd} is disabled
61 no_fork = False
62
63
64 class RunResult(object):
65   """Holds the result of running external programs.
66
67   @type exit_code: int
68   @ivar exit_code: the exit code of the program, or None (if the program
69       didn't exit())
70   @type signal: int or None
71   @ivar signal: the signal that caused the program to finish, or None
72       (if the program wasn't terminated by a signal)
73   @type stdout: str
74   @ivar stdout: the standard output of the program
75   @type stderr: str
76   @ivar stderr: the standard error of the program
77   @type failed: boolean
78   @ivar failed: True in case the program was
79       terminated by a signal or exited with a non-zero exit code
80   @ivar fail_reason: a string detailing the termination reason
81
82   """
83   __slots__ = ["exit_code", "signal", "stdout", "stderr",
84                "failed", "fail_reason", "cmd"]
85
86
87   def __init__(self, exit_code, signal_, stdout, stderr, cmd):
88     self.cmd = cmd
89     self.exit_code = exit_code
90     self.signal = signal_
91     self.stdout = stdout
92     self.stderr = stderr
93     self.failed = (signal_ is not None or exit_code != 0)
94
95     if self.signal is not None:
96       self.fail_reason = "terminated by signal %s" % self.signal
97     elif self.exit_code is not None:
98       self.fail_reason = "exited with exit code %s" % self.exit_code
99     else:
100       self.fail_reason = "unable to determine termination reason"
101
102     if self.failed:
103       logging.debug("Command '%s' failed (%s); output: %s",
104                     self.cmd, self.fail_reason, self.output)
105
106   def _GetOutput(self):
107     """Returns the combined stdout and stderr for easier usage.
108
109     """
110     return self.stdout + self.stderr
111
112   output = property(_GetOutput, None, None, "Return full output")
113
114
115 def RunCmd(cmd, env=None, output=None, cwd='/'):
116   """Execute a (shell) command.
117
118   The command should not read from its standard input, as it will be
119   closed.
120
121   @type  cmd: string or list
122   @param cmd: Command to run
123   @type env: dict
124   @param env: Additional environment
125   @type output: str
126   @param output: if desired, the output of the command can be
127       saved in a file instead of the RunResult instance; this
128       parameter denotes the file name (if not None)
129   @type cwd: string
130   @param cwd: if specified, will be used as the working
131       directory for the command; the default will be /
132   @rtype: L{RunResult}
133   @return: RunResult instance
134   @raise erors.ProgrammerError: if we call this when forks are disabled
135
136   """
137   if no_fork:
138     raise errors.ProgrammerError("utils.RunCmd() called with fork() disabled")
139
140   if isinstance(cmd, list):
141     cmd = [str(val) for val in cmd]
142     strcmd = " ".join(cmd)
143     shell = False
144   else:
145     strcmd = cmd
146     shell = True
147   logging.debug("RunCmd '%s'", strcmd)
148
149   cmd_env = os.environ.copy()
150   cmd_env["LC_ALL"] = "C"
151   if env is not None:
152     cmd_env.update(env)
153
154   if output is None:
155     out, err, status = _RunCmdPipe(cmd, cmd_env, shell, cwd)
156   else:
157     status = _RunCmdFile(cmd, cmd_env, shell, output, cwd)
158     out = err = ""
159
160   if status >= 0:
161     exitcode = status
162     signal_ = None
163   else:
164     exitcode = None
165     signal_ = -status
166
167   return RunResult(exitcode, signal_, out, err, strcmd)
168
169 def _RunCmdPipe(cmd, env, via_shell, cwd):
170   """Run a command and return its output.
171
172   @type  cmd: string or list
173   @param cmd: Command to run
174   @type env: dict
175   @param env: The environment to use
176   @type via_shell: bool
177   @param via_shell: if we should run via the shell
178   @type cwd: string
179   @param cwd: the working directory for the program
180   @rtype: tuple
181   @return: (out, err, status)
182
183   """
184   poller = select.poll()
185   child = subprocess.Popen(cmd, shell=via_shell,
186                            stderr=subprocess.PIPE,
187                            stdout=subprocess.PIPE,
188                            stdin=subprocess.PIPE,
189                            close_fds=True, env=env,
190                            cwd=cwd)
191
192   child.stdin.close()
193   poller.register(child.stdout, select.POLLIN)
194   poller.register(child.stderr, select.POLLIN)
195   out = StringIO()
196   err = StringIO()
197   fdmap = {
198     child.stdout.fileno(): (out, child.stdout),
199     child.stderr.fileno(): (err, child.stderr),
200     }
201   for fd in fdmap:
202     status = fcntl.fcntl(fd, fcntl.F_GETFL)
203     fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
204
205   while fdmap:
206     for fd, event in poller.poll():
207       if event & select.POLLIN or event & select.POLLPRI:
208         data = fdmap[fd][1].read()
209         # no data from read signifies EOF (the same as POLLHUP)
210         if not data:
211           poller.unregister(fd)
212           del fdmap[fd]
213           continue
214         fdmap[fd][0].write(data)
215       if (event & select.POLLNVAL or event & select.POLLHUP or
216           event & select.POLLERR):
217         poller.unregister(fd)
218         del fdmap[fd]
219
220   out = out.getvalue()
221   err = err.getvalue()
222
223   status = child.wait()
224   return out, err, status
225
226
227 def _RunCmdFile(cmd, env, via_shell, output, cwd):
228   """Run a command and save its output to a file.
229
230   @type  cmd: string or list
231   @param cmd: Command to run
232   @type env: dict
233   @param env: The environment to use
234   @type via_shell: bool
235   @param via_shell: if we should run via the shell
236   @type output: str
237   @param output: the filename in which to save the output
238   @type cwd: string
239   @param cwd: the working directory for the program
240   @rtype: int
241   @return: the exit status
242
243   """
244   fh = open(output, "a")
245   try:
246     child = subprocess.Popen(cmd, shell=via_shell,
247                              stderr=subprocess.STDOUT,
248                              stdout=fh,
249                              stdin=subprocess.PIPE,
250                              close_fds=True, env=env,
251                              cwd=cwd)
252
253     child.stdin.close()
254     status = child.wait()
255   finally:
256     fh.close()
257   return status
258
259
260 def RemoveFile(filename):
261   """Remove a file ignoring some errors.
262
263   Remove a file, ignoring non-existing ones or directories. Other
264   errors are passed.
265
266   @type filename: str
267   @param filename: the file to be removed
268
269   """
270   try:
271     os.unlink(filename)
272   except OSError, err:
273     if err.errno not in (errno.ENOENT, errno.EISDIR):
274       raise
275
276
277 def _FingerprintFile(filename):
278   """Compute the fingerprint of a file.
279
280   If the file does not exist, a None will be returned
281   instead.
282
283   @type filename: str
284   @param filename: the filename to checksum
285   @rtype: str
286   @return: the hex digest of the sha checksum of the contents
287       of the file
288
289   """
290   if not (os.path.exists(filename) and os.path.isfile(filename)):
291     return None
292
293   f = open(filename)
294
295   fp = sha.sha()
296   while True:
297     data = f.read(4096)
298     if not data:
299       break
300
301     fp.update(data)
302
303   return fp.hexdigest()
304
305
306 def FingerprintFiles(files):
307   """Compute fingerprints for a list of files.
308
309   @type files: list
310   @param files: the list of filename to fingerprint
311   @rtype: dict
312   @return: a dictionary filename: fingerprint, holding only
313       existing files
314
315   """
316   ret = {}
317
318   for filename in files:
319     cksum = _FingerprintFile(filename)
320     if cksum:
321       ret[filename] = cksum
322
323   return ret
324
325
326 def CheckDict(target, template, logname=None):
327   """Ensure a dictionary has a required set of keys.
328
329   For the given dictionaries I{target} and I{template}, ensure
330   I{target} has all the keys from I{template}. Missing keys are added
331   with values from template.
332
333   @type target: dict
334   @param target: the dictionary to update
335   @type template: dict
336   @param template: the dictionary holding the default values
337   @type logname: str or None
338   @param logname: if not None, causes the missing keys to be
339       logged with this name
340
341   """
342   missing = []
343   for k in template:
344     if k not in target:
345       missing.append(k)
346       target[k] = template[k]
347
348   if missing and logname:
349     logging.warning('%s missing keys %s', logname, ', '.join(missing))
350
351
352 def IsProcessAlive(pid):
353   """Check if a given pid exists on the system.
354
355   @note: zombie status is not handled, so zombie processes
356       will be returned as alive
357   @type pid: int
358   @param pid: the process ID to check
359   @rtype: boolean
360   @return: True if the process exists
361
362   """
363   if pid <= 0:
364     return False
365
366   try:
367     os.stat("/proc/%d/status" % pid)
368     return True
369   except EnvironmentError, err:
370     if err.errno in (errno.ENOENT, errno.ENOTDIR):
371       return False
372     raise
373
374
375 def ReadPidFile(pidfile):
376   """Read a pid from a file.
377
378   @type  pidfile: string
379   @param pidfile: path to the file containing the pid
380   @rtype: int
381   @return: The process id, if the file exista and contains a valid PID,
382            otherwise 0
383
384   """
385   try:
386     pf = open(pidfile, 'r')
387   except EnvironmentError, err:
388     if err.errno != errno.ENOENT:
389       logging.exception("Can't read pid file?!")
390     return 0
391
392   try:
393     pid = int(pf.read())
394   except ValueError, err:
395     logging.info("Can't parse pid file contents", exc_info=True)
396     return 0
397
398   return pid
399
400
401 def MatchNameComponent(key, name_list):
402   """Try to match a name against a list.
403
404   This function will try to match a name like test1 against a list
405   like C{['test1.example.com', 'test2.example.com', ...]}. Against
406   this list, I{'test1'} as well as I{'test1.example'} will match, but
407   not I{'test1.ex'}. A multiple match will be considered as no match
408   at all (e.g. I{'test1'} against C{['test1.example.com',
409   'test1.example.org']}).
410
411   @type key: str
412   @param key: the name to be searched
413   @type name_list: list
414   @param name_list: the list of strings against which to search the key
415
416   @rtype: None or str
417   @return: None if there is no match I{or} if there are multiple matches,
418       otherwise the element from the list which matches
419
420   """
421   mo = re.compile("^%s(\..*)?$" % re.escape(key))
422   names_filtered = [name for name in name_list if mo.match(name) is not None]
423   if len(names_filtered) != 1:
424     return None
425   return names_filtered[0]
426
427
428 class HostInfo:
429   """Class implementing resolver and hostname functionality
430
431   """
432   def __init__(self, name=None):
433     """Initialize the host name object.
434
435     If the name argument is not passed, it will use this system's
436     name.
437
438     """
439     if name is None:
440       name = self.SysName()
441
442     self.query = name
443     self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
444     self.ip = self.ipaddrs[0]
445
446   def ShortName(self):
447     """Returns the hostname without domain.
448
449     """
450     return self.name.split('.')[0]
451
452   @staticmethod
453   def SysName():
454     """Return the current system's name.
455
456     This is simply a wrapper over C{socket.gethostname()}.
457
458     """
459     return socket.gethostname()
460
461   @staticmethod
462   def LookupHostname(hostname):
463     """Look up hostname
464
465     @type hostname: str
466     @param hostname: hostname to look up
467
468     @rtype: tuple
469     @return: a tuple (name, aliases, ipaddrs) as returned by
470         C{socket.gethostbyname_ex}
471     @raise errors.ResolverError: in case of errors in resolving
472
473     """
474     try:
475       result = socket.gethostbyname_ex(hostname)
476     except socket.gaierror, err:
477       # hostname not found in DNS
478       raise errors.ResolverError(hostname, err.args[0], err.args[1])
479
480     return result
481
482
483 def ListVolumeGroups():
484   """List volume groups and their size
485
486   @rtype: dict
487   @return:
488        Dictionary with keys volume name and values
489        the size of the volume
490
491   """
492   command = "vgs --noheadings --units m --nosuffix -o name,size"
493   result = RunCmd(command)
494   retval = {}
495   if result.failed:
496     return retval
497
498   for line in result.stdout.splitlines():
499     try:
500       name, size = line.split()
501       size = int(float(size))
502     except (IndexError, ValueError), err:
503       logging.error("Invalid output from vgs (%s): %s", err, line)
504       continue
505
506     retval[name] = size
507
508   return retval
509
510
511 def BridgeExists(bridge):
512   """Check whether the given bridge exists in the system
513
514   @type bridge: str
515   @param bridge: the bridge name to check
516   @rtype: boolean
517   @return: True if it does
518
519   """
520   return os.path.isdir("/sys/class/net/%s/bridge" % bridge)
521
522
523 def CheckBEParams(beparams):
524   """Checks whether the user-supplied be-params are valid,
525   and converts them from string format where appropriate.
526
527   @type beparams: dict
528   @param beparams: new params dict
529
530   """
531   if beparams:
532     for item in beparams:
533       if item not in constants.BES_PARAMETERS:
534         raise errors.OpPrereqError("Unknown backend parameter %s" % item)
535       if item in (constants.BE_MEMORY, constants.BE_VCPUS):
536         val = beparams[item]
537         if val != constants.VALUE_DEFAULT:
538           try:
539             val = int(val)
540           except ValueError, err:
541             raise errors.OpPrereqError("Invalid %s size: %s" % (item, str(err)))
542           beparams[item] = val
543       if item in (constants.BE_AUTO_BALANCE):
544         val = beparams[item]
545         if val == constants.VALUE_TRUE:
546           beparams[item] = True
547         elif val == constants.VALUE_FALSE:
548           beparams[item] = False
549         else:
550           raise errors.OpPrereqError("Invalid %s value: %s" % (item, val))
551
552
553 def NiceSort(name_list):
554   """Sort a list of strings based on digit and non-digit groupings.
555
556   Given a list of names C{['a1', 'a10', 'a11', 'a2']} this function
557   will sort the list in the logical order C{['a1', 'a2', 'a10',
558   'a11']}.
559
560   The sort algorithm breaks each name in groups of either only-digits
561   or no-digits. Only the first eight such groups are considered, and
562   after that we just use what's left of the string.
563
564   @type name_list: list
565   @param name_list: the names to be sorted
566   @rtype: list
567   @return: a copy of the name list sorted with our algorithm
568
569   """
570   _SORTER_BASE = "(\D+|\d+)"
571   _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
572                                                   _SORTER_BASE, _SORTER_BASE,
573                                                   _SORTER_BASE, _SORTER_BASE,
574                                                   _SORTER_BASE, _SORTER_BASE)
575   _SORTER_RE = re.compile(_SORTER_FULL)
576   _SORTER_NODIGIT = re.compile("^\D*$")
577   def _TryInt(val):
578     """Attempts to convert a variable to integer."""
579     if val is None or _SORTER_NODIGIT.match(val):
580       return val
581     rval = int(val)
582     return rval
583
584   to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
585              for name in name_list]
586   to_sort.sort()
587   return [tup[1] for tup in to_sort]
588
589
590 def TryConvert(fn, val):
591   """Try to convert a value ignoring errors.
592
593   This function tries to apply function I{fn} to I{val}. If no
594   C{ValueError} or C{TypeError} exceptions are raised, it will return
595   the result, else it will return the original value. Any other
596   exceptions are propagated to the caller.
597
598   @type fn: callable
599   @param fn: function to apply to the value
600   @param val: the value to be converted
601   @return: The converted value if the conversion was successful,
602       otherwise the original value.
603
604   """
605   try:
606     nv = fn(val)
607   except (ValueError, TypeError), err:
608     nv = val
609   return nv
610
611
612 def IsValidIP(ip):
613   """Verifies the syntax of an IPv4 address.
614
615   This function checks if the IPv4 address passes is valid or not based
616   on syntax (not IP range, class calculations, etc.).
617
618   @type ip: str
619   @param ip: the address to be checked
620   @rtype: a regular expression match object
621   @return: a regular epression match object, or None if the
622       address is not valid
623
624   """
625   unit = "(0|[1-9]\d{0,2})"
626   #TODO: convert and return only boolean
627   return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
628
629
630 def IsValidShellParam(word):
631   """Verifies is the given word is safe from the shell's p.o.v.
632
633   This means that we can pass this to a command via the shell and be
634   sure that it doesn't alter the command line and is passed as such to
635   the actual command.
636
637   Note that we are overly restrictive here, in order to be on the safe
638   side.
639
640   @type word: str
641   @param word: the word to check
642   @rtype: boolean
643   @return: True if the word is 'safe'
644
645   """
646   return bool(re.match("^[-a-zA-Z0-9._+/:%@]+$", word))
647
648
649 def BuildShellCmd(template, *args):
650   """Build a safe shell command line from the given arguments.
651
652   This function will check all arguments in the args list so that they
653   are valid shell parameters (i.e. they don't contain shell
654   metacharaters). If everything is ok, it will return the result of
655   template % args.
656
657   @type template: str
658   @param template: the string holding the template for the
659       string formatting
660   @rtype: str
661   @return: the expanded command line
662
663   """
664   for word in args:
665     if not IsValidShellParam(word):
666       raise errors.ProgrammerError("Shell argument '%s' contains"
667                                    " invalid characters" % word)
668   return template % args
669
670
671 def FormatUnit(value, units):
672   """Formats an incoming number of MiB with the appropriate unit.
673
674   @type value: int
675   @param value: integer representing the value in MiB (1048576)
676   @type units: char
677   @param units: the type of formatting we should do:
678       - 'h' for automatic scaling
679       - 'm' for MiBs
680       - 'g' for GiBs
681       - 't' for TiBs
682   @rtype: str
683   @return: the formatted value (with suffix)
684
685   """
686   if units not in ('m', 'g', 't', 'h'):
687     raise errors.ProgrammerError("Invalid unit specified '%s'" % str(units))
688
689   suffix = ''
690
691   if units == 'm' or (units == 'h' and value < 1024):
692     if units == 'h':
693       suffix = 'M'
694     return "%d%s" % (round(value, 0), suffix)
695
696   elif units == 'g' or (units == 'h' and value < (1024 * 1024)):
697     if units == 'h':
698       suffix = 'G'
699     return "%0.1f%s" % (round(float(value) / 1024, 1), suffix)
700
701   else:
702     if units == 'h':
703       suffix = 'T'
704     return "%0.1f%s" % (round(float(value) / 1024 / 1024, 1), suffix)
705
706
707 def ParseUnit(input_string):
708   """Tries to extract number and scale from the given string.
709
710   Input must be in the format C{NUMBER+ [DOT NUMBER+] SPACE*
711   [UNIT]}. If no unit is specified, it defaults to MiB. Return value
712   is always an int in MiB.
713
714   """
715   m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', input_string)
716   if not m:
717     raise errors.UnitParseError("Invalid format")
718
719   value = float(m.groups()[0])
720
721   unit = m.groups()[1]
722   if unit:
723     lcunit = unit.lower()
724   else:
725     lcunit = 'm'
726
727   if lcunit in ('m', 'mb', 'mib'):
728     # Value already in MiB
729     pass
730
731   elif lcunit in ('g', 'gb', 'gib'):
732     value *= 1024
733
734   elif lcunit in ('t', 'tb', 'tib'):
735     value *= 1024 * 1024
736
737   else:
738     raise errors.UnitParseError("Unknown unit: %s" % unit)
739
740   # Make sure we round up
741   if int(value) < value:
742     value += 1
743
744   # Round up to the next multiple of 4
745   value = int(value)
746   if value % 4:
747     value += 4 - value % 4
748
749   return value
750
751
752 def AddAuthorizedKey(file_name, key):
753   """Adds an SSH public key to an authorized_keys file.
754
755   @type file_name: str
756   @param file_name: path to authorized_keys file
757   @type key: str
758   @param key: string containing key
759
760   """
761   key_fields = key.split()
762
763   f = open(file_name, 'a+')
764   try:
765     nl = True
766     for line in f:
767       # Ignore whitespace changes
768       if line.split() == key_fields:
769         break
770       nl = line.endswith('\n')
771     else:
772       if not nl:
773         f.write("\n")
774       f.write(key.rstrip('\r\n'))
775       f.write("\n")
776       f.flush()
777   finally:
778     f.close()
779
780
781 def RemoveAuthorizedKey(file_name, key):
782   """Removes an SSH public key from an authorized_keys file.
783
784   @type file_name: str
785   @param file_name: path to authorized_keys file
786   @type key: str
787   @param key: string containing key
788
789   """
790   key_fields = key.split()
791
792   fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
793   try:
794     out = os.fdopen(fd, 'w')
795     try:
796       f = open(file_name, 'r')
797       try:
798         for line in f:
799           # Ignore whitespace changes while comparing lines
800           if line.split() != key_fields:
801             out.write(line)
802
803         out.flush()
804         os.rename(tmpname, file_name)
805       finally:
806         f.close()
807     finally:
808       out.close()
809   except:
810     RemoveFile(tmpname)
811     raise
812
813
814 def SetEtcHostsEntry(file_name, ip, hostname, aliases):
815   """Sets the name of an IP address and hostname in /etc/hosts.
816
817   @type file_name: str
818   @param file_name: path to the file to modify (usually C{/etc/hosts})
819   @type ip: str
820   @param ip: the IP address
821   @type hostname: str
822   @param hostname: the hostname to be added
823   @type aliases: list
824   @param aliases: the list of aliases to add for the hostname
825
826   """
827   # Ensure aliases are unique
828   aliases = UniqueSequence([hostname] + aliases)[1:]
829
830   fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
831   try:
832     out = os.fdopen(fd, 'w')
833     try:
834       f = open(file_name, 'r')
835       try:
836         written = False
837         for line in f:
838           fields = line.split()
839           if fields and not fields[0].startswith('#') and ip == fields[0]:
840             continue
841           out.write(line)
842
843         out.write("%s\t%s" % (ip, hostname))
844         if aliases:
845           out.write(" %s" % ' '.join(aliases))
846         out.write('\n')
847
848         out.flush()
849         os.fsync(out)
850         os.rename(tmpname, file_name)
851       finally:
852         f.close()
853     finally:
854       out.close()
855   except:
856     RemoveFile(tmpname)
857     raise
858
859
860 def AddHostToEtcHosts(hostname):
861   """Wrapper around SetEtcHostsEntry.
862
863   @type hostname: str
864   @param hostname: a hostname that will be resolved and added to
865       L{constants.ETC_HOSTS}
866
867   """
868   hi = HostInfo(name=hostname)
869   SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])
870
871
872 def RemoveEtcHostsEntry(file_name, hostname):
873   """Removes a hostname from /etc/hosts.
874
875   IP addresses without names are removed from the file.
876
877   @type file_name: str
878   @param file_name: path to the file to modify (usually C{/etc/hosts})
879   @type hostname: str
880   @param hostname: the hostname to be removed
881
882   """
883   fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
884   try:
885     out = os.fdopen(fd, 'w')
886     try:
887       f = open(file_name, 'r')
888       try:
889         for line in f:
890           fields = line.split()
891           if len(fields) > 1 and not fields[0].startswith('#'):
892             names = fields[1:]
893             if hostname in names:
894               while hostname in names:
895                 names.remove(hostname)
896               if names:
897                 out.write("%s %s\n" % (fields[0], ' '.join(names)))
898               continue
899
900           out.write(line)
901
902         out.flush()
903         os.fsync(out)
904         os.rename(tmpname, file_name)
905       finally:
906         f.close()
907     finally:
908       out.close()
909   except:
910     RemoveFile(tmpname)
911     raise
912
913
914 def RemoveHostFromEtcHosts(hostname):
915   """Wrapper around RemoveEtcHostsEntry.
916
917   @type hostname: str
918   @param hostname: hostname that will be resolved and its
919       full and shot name will be removed from
920       L{constants.ETC_HOSTS}
921
922   """
923   hi = HostInfo(name=hostname)
924   RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.name)
925   RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.ShortName())
926
927
928 def CreateBackup(file_name):
929   """Creates a backup of a file.
930
931   @type file_name: str
932   @param file_name: file to be backed up
933   @rtype: str
934   @return: the path to the newly created backup
935   @raise errors.ProgrammerError: for invalid file names
936
937   """
938   if not os.path.isfile(file_name):
939     raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
940                                 file_name)
941
942   prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
943   dir_name = os.path.dirname(file_name)
944
945   fsrc = open(file_name, 'rb')
946   try:
947     (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
948     fdst = os.fdopen(fd, 'wb')
949     try:
950       shutil.copyfileobj(fsrc, fdst)
951     finally:
952       fdst.close()
953   finally:
954     fsrc.close()
955
956   return backup_name
957
958
959 def ShellQuote(value):
960   """Quotes shell argument according to POSIX.
961
962   @type value: str
963   @param value: the argument to be quoted
964   @rtype: str
965   @return: the quoted value
966
967   """
968   if _re_shell_unquoted.match(value):
969     return value
970   else:
971     return "'%s'" % value.replace("'", "'\\''")
972
973
974 def ShellQuoteArgs(args):
975   """Quotes a list of shell arguments.
976
977   @type args: list
978   @param args: list of arguments to be quoted
979   @rtype: str
980   @return: the quoted arguments concatenaned with spaces
981
982   """
983   return ' '.join([ShellQuote(i) for i in args])
984
985
986 def TcpPing(target, port, timeout=10, live_port_needed=False, source=None):
987   """Simple ping implementation using TCP connect(2).
988
989   Check if the given IP is reachable by doing attempting a TCP connect
990   to it.
991
992   @type target: str
993   @param target: the IP or hostname to ping
994   @type port: int
995   @param port: the port to connect to
996   @type timeout: int
997   @param timeout: the timeout on the connection attemp
998   @type live_port_needed: boolean
999   @param live_port_needed: whether a closed port will cause the
1000       function to return failure, as if there was a timeout
1001   @type source: str or None
1002   @param source: if specified, will cause the connect to be made
1003       from this specific source address; failures to bind other
1004       than C{EADDRNOTAVAIL} will be ignored
1005
1006   """
1007   sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
1008
1009   sucess = False
1010
1011   if source is not None:
1012     try:
1013       sock.bind((source, 0))
1014     except socket.error, (errcode, errstring):
1015       if errcode == errno.EADDRNOTAVAIL:
1016         success = False
1017
1018   sock.settimeout(timeout)
1019
1020   try:
1021     sock.connect((target, port))
1022     sock.close()
1023     success = True
1024   except socket.timeout:
1025     success = False
1026   except socket.error, (errcode, errstring):
1027     success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
1028
1029   return success
1030
1031
1032 def OwnIpAddress(address):
1033   """Check if the current host has the the given IP address.
1034
1035   Currently this is done by TCP-pinging the address from the loopback
1036   address.
1037
1038   @type address: string
1039   @param address: the addres to check
1040   @rtype: bool
1041   @return: True if we own the address
1042
1043   """
1044   return TcpPing(address, constants.DEFAULT_NODED_PORT,
1045                  source=constants.LOCALHOST_IP_ADDRESS)
1046
1047
1048 def ListVisibleFiles(path):
1049   """Returns a list of visible files in a directory.
1050
1051   @type path: str
1052   @param path: the directory to enumerate
1053   @rtype: list
1054   @return: the list of all files not starting with a dot
1055
1056   """
1057   files = [i for i in os.listdir(path) if not i.startswith(".")]
1058   files.sort()
1059   return files
1060
1061
1062 def GetHomeDir(user, default=None):
1063   """Try to get the homedir of the given user.
1064
1065   The user can be passed either as a string (denoting the name) or as
1066   an integer (denoting the user id). If the user is not found, the
1067   'default' argument is returned, which defaults to None.
1068
1069   """
1070   try:
1071     if isinstance(user, basestring):
1072       result = pwd.getpwnam(user)
1073     elif isinstance(user, (int, long)):
1074       result = pwd.getpwuid(user)
1075     else:
1076       raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
1077                                    type(user))
1078   except KeyError:
1079     return default
1080   return result.pw_dir
1081
1082
1083 def NewUUID():
1084   """Returns a random UUID.
1085
1086   @note: This is a Linux-specific method as it uses the /proc
1087       filesystem.
1088   @rtype: str
1089
1090   """
1091   f = open("/proc/sys/kernel/random/uuid", "r")
1092   try:
1093     return f.read(128).rstrip("\n")
1094   finally:
1095     f.close()
1096
1097
1098 def GenerateSecret():
1099   """Generates a random secret.
1100
1101   This will generate a pseudo-random secret, and return its sha digest
1102   (so that it can be used where an ASCII string is needed).
1103
1104   @rtype: str
1105   @return: a sha1 hexdigest of a block of 64 random bytes
1106
1107   """
1108   return sha.new(os.urandom(64)).hexdigest()
1109
1110
1111 def ReadFile(file_name, size=None):
1112   """Reads a file.
1113
1114   @type size: None or int
1115   @param size: Read at most size bytes
1116   @rtype: str
1117   @return: the (possibly partial) conent of the file
1118
1119   """
1120   f = open(file_name, "r")
1121   try:
1122     if size is None:
1123       return f.read()
1124     else:
1125       return f.read(size)
1126   finally:
1127     f.close()
1128
1129
1130 def WriteFile(file_name, fn=None, data=None,
1131               mode=None, uid=-1, gid=-1,
1132               atime=None, mtime=None, close=True,
1133               dry_run=False, backup=False,
1134               prewrite=None, postwrite=None):
1135   """(Over)write a file atomically.
1136
1137   The file_name and either fn (a function taking one argument, the
1138   file descriptor, and which should write the data to it) or data (the
1139   contents of the file) must be passed. The other arguments are
1140   optional and allow setting the file mode, owner and group, and the
1141   mtime/atime of the file.
1142
1143   If the function doesn't raise an exception, it has succeeded and the
1144   target file has the new contents. If the file has raised an
1145   exception, an existing target file should be unmodified and the
1146   temporary file should be removed.
1147
1148   @type file_name: str
1149   @param file_name: the target filename
1150   @type fn: callable
1151   @param fn: content writing function, called with
1152       file descriptor as parameter
1153   @type data: sr
1154   @param data: contents of the file
1155   @type mode: int
1156   @param mode: file mode
1157   @type uid: int
1158   @param uid: the owner of the file
1159   @type gid: int
1160   @param gid: the group of the file
1161   @type atime: int
1162   @param atime: a custom access time to be set on the file
1163   @type mtime: int
1164   @param mtime: a custom modification time to be set on the file
1165   @type close: boolean
1166   @param close: whether to close file after writing it
1167   @type prewrite: callable
1168   @param prewrite: function to be called before writing content
1169   @type postwrite: callable
1170   @param postwrite: function to be called after writing content
1171
1172   @rtype: None or int
1173   @return: None if the 'close' parameter evaluates to True,
1174       otherwise the file descriptor
1175
1176   @raise errors.ProgrammerError: if an of the arguments are not valid
1177
1178   """
1179   if not os.path.isabs(file_name):
1180     raise errors.ProgrammerError("Path passed to WriteFile is not"
1181                                  " absolute: '%s'" % file_name)
1182
1183   if [fn, data].count(None) != 1:
1184     raise errors.ProgrammerError("fn or data required")
1185
1186   if [atime, mtime].count(None) == 1:
1187     raise errors.ProgrammerError("Both atime and mtime must be either"
1188                                  " set or None")
1189
1190   if backup and not dry_run and os.path.isfile(file_name):
1191     CreateBackup(file_name)
1192
1193   dir_name, base_name = os.path.split(file_name)
1194   fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
1195   # here we need to make sure we remove the temp file, if any error
1196   # leaves it in place
1197   try:
1198     if uid != -1 or gid != -1:
1199       os.chown(new_name, uid, gid)
1200     if mode:
1201       os.chmod(new_name, mode)
1202     if callable(prewrite):
1203       prewrite(fd)
1204     if data is not None:
1205       os.write(fd, data)
1206     else:
1207       fn(fd)
1208     if callable(postwrite):
1209       postwrite(fd)
1210     os.fsync(fd)
1211     if atime is not None and mtime is not None:
1212       os.utime(new_name, (atime, mtime))
1213     if not dry_run:
1214       os.rename(new_name, file_name)
1215   finally:
1216     if close:
1217       os.close(fd)
1218       result = None
1219     else:
1220       result = fd
1221     RemoveFile(new_name)
1222
1223   return result
1224
1225
1226 def FirstFree(seq, base=0):
1227   """Returns the first non-existing integer from seq.
1228
1229   The seq argument should be a sorted list of positive integers. The
1230   first time the index of an element is smaller than the element
1231   value, the index will be returned.
1232
1233   The base argument is used to start at a different offset,
1234   i.e. C{[3, 4, 6]} with I{offset=3} will return 5.
1235
1236   Example: C{[0, 1, 3]} will return I{2}.
1237
1238   @type seq: sequence
1239   @param seq: the sequence to be analyzed.
1240   @type base: int
1241   @param base: use this value as the base index of the sequence
1242   @rtype: int
1243   @return: the first non-used index in the sequence
1244
1245   """
1246   for idx, elem in enumerate(seq):
1247     assert elem >= base, "Passed element is higher than base offset"
1248     if elem > idx + base:
1249       # idx is not used
1250       return idx + base
1251   return None
1252
1253
1254 def all(seq, pred=bool):
1255   "Returns True if pred(x) is True for every element in the iterable"
1256   for elem in itertools.ifilterfalse(pred, seq):
1257     return False
1258   return True
1259
1260
1261 def any(seq, pred=bool):
1262   "Returns True if pred(x) is True for at least one element in the iterable"
1263   for elem in itertools.ifilter(pred, seq):
1264     return True
1265   return False
1266
1267
1268 def UniqueSequence(seq):
1269   """Returns a list with unique elements.
1270
1271   Element order is preserved.
1272
1273   @type seq: sequence
1274   @param seq: the sequence with the source elementes
1275   @rtype: list
1276   @return: list of unique elements from seq
1277
1278   """
1279   seen = set()
1280   return [i for i in seq if i not in seen and not seen.add(i)]
1281
1282
1283 def IsValidMac(mac):
1284   """Predicate to check if a MAC address is valid.
1285
1286   Checks wether the supplied MAC address is formally correct, only
1287   accepts colon separated format.
1288
1289   @type mac: str
1290   @param mac: the MAC to be validated
1291   @rtype: boolean
1292   @return: True is the MAC seems valid
1293
1294   """
1295   mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
1296   return mac_check.match(mac) is not None
1297
1298
1299 def TestDelay(duration):
1300   """Sleep for a fixed amount of time.
1301
1302   @type duration: float
1303   @param duration: the sleep duration
1304   @rtype: boolean
1305   @return: False for negative value, True otherwise
1306
1307   """
1308   if duration < 0:
1309     return False
1310   time.sleep(duration)
1311   return True
1312
1313
1314 def Daemonize(logfile, noclose_fds=None):
1315   """Daemonize the current process.
1316
1317   This detaches the current process from the controlling terminal and
1318   runs it in the background as a daemon.
1319
1320   @type logfile: str
1321   @param logfile: the logfile to which we should redirect stdout/stderr
1322   @type noclose_fds: list or None
1323   @param noclose_fds: if given, it denotes a list of file descriptor
1324       that should not be closed
1325   @rtype: int
1326   @returns: the value zero
1327
1328   """
1329   UMASK = 077
1330   WORKDIR = "/"
1331   # Default maximum for the number of available file descriptors.
1332   if 'SC_OPEN_MAX' in os.sysconf_names:
1333     try:
1334       MAXFD = os.sysconf('SC_OPEN_MAX')
1335       if MAXFD < 0:
1336         MAXFD = 1024
1337     except OSError:
1338       MAXFD = 1024
1339   else:
1340     MAXFD = 1024
1341
1342   # this might fail
1343   pid = os.fork()
1344   if (pid == 0):  # The first child.
1345     os.setsid()
1346     # this might fail
1347     pid = os.fork() # Fork a second child.
1348     if (pid == 0):  # The second child.
1349       os.chdir(WORKDIR)
1350       os.umask(UMASK)
1351     else:
1352       # exit() or _exit()?  See below.
1353       os._exit(0) # Exit parent (the first child) of the second child.
1354   else:
1355     os._exit(0) # Exit parent of the first child.
1356   maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1357   if (maxfd == resource.RLIM_INFINITY):
1358     maxfd = MAXFD
1359
1360   # Iterate through and close all file descriptors.
1361   for fd in range(0, maxfd):
1362     if noclose_fds and fd in noclose_fds:
1363       continue
1364     try:
1365       os.close(fd)
1366     except OSError: # ERROR, fd wasn't open to begin with (ignored)
1367       pass
1368   os.open(logfile, os.O_RDWR|os.O_CREAT|os.O_APPEND, 0600)
1369   # Duplicate standard input to standard output and standard error.
1370   os.dup2(0, 1)     # standard output (1)
1371   os.dup2(0, 2)     # standard error (2)
1372   return 0
1373
1374
1375 def DaemonPidFileName(name):
1376   """Compute a ganeti pid file absolute path
1377
1378   @type name: str
1379   @param name: the daemon name
1380   @rtype: str
1381   @return: the full path to the pidfile corresponding to the given
1382       daemon name
1383
1384   """
1385   return os.path.join(constants.RUN_GANETI_DIR, "%s.pid" % name)
1386
1387
1388 def WritePidFile(name):
1389   """Write the current process pidfile.
1390
1391   The file will be written to L{constants.RUN_GANETI_DIR}I{/name.pid}
1392
1393   @type name: str
1394   @param name: the daemon name to use
1395   @raise errors.GenericError: if the pid file already exists and
1396       points to a live process
1397
1398   """
1399   pid = os.getpid()
1400   pidfilename = DaemonPidFileName(name)
1401   if IsProcessAlive(ReadPidFile(pidfilename)):
1402     raise errors.GenericError("%s contains a live process" % pidfilename)
1403
1404   WriteFile(pidfilename, data="%d\n" % pid)
1405
1406
1407 def RemovePidFile(name):
1408   """Remove the current process pidfile.
1409
1410   Any errors are ignored.
1411
1412   @type name: str
1413   @param name: the daemon name used to derive the pidfile name
1414
1415   """
1416   pid = os.getpid()
1417   pidfilename = DaemonPidFileName(name)
1418   # TODO: we could check here that the file contains our pid
1419   try:
1420     RemoveFile(pidfilename)
1421   except:
1422     pass
1423
1424
1425 def KillProcess(pid, signal_=signal.SIGTERM, timeout=30,
1426                 waitpid=False):
1427   """Kill a process given by its pid.
1428
1429   @type pid: int
1430   @param pid: The PID to terminate.
1431   @type signal_: int
1432   @param signal_: The signal to send, by default SIGTERM
1433   @type timeout: int
1434   @param timeout: The timeout after which, if the process is still alive,
1435                   a SIGKILL will be sent. If not positive, no such checking
1436                   will be done
1437   @type waitpid: boolean
1438   @param waitpid: If true, we should waitpid on this process after
1439       sending signals, since it's our own child and otherwise it
1440       would remain as zombie
1441
1442   """
1443   def _helper(pid, signal_, wait):
1444     """Simple helper to encapsulate the kill/waitpid sequence"""
1445     os.kill(pid, signal_)
1446     if wait:
1447       try:
1448         os.waitpid(pid, os.WNOHANG)
1449       except OSError:
1450         pass
1451
1452   if pid <= 0:
1453     # kill with pid=0 == suicide
1454     raise errors.ProgrammerError("Invalid pid given '%s'" % pid)
1455
1456   if not IsProcessAlive(pid):
1457     return
1458   _helper(pid, signal_, waitpid)
1459   if timeout <= 0:
1460     return
1461   end = time.time() + timeout
1462   while time.time() < end and IsProcessAlive(pid):
1463     time.sleep(0.1)
1464   if IsProcessAlive(pid):
1465     _helper(pid, signal.SIGKILL, waitpid)
1466
1467
1468 def FindFile(name, search_path, test=os.path.exists):
1469   """Look for a filesystem object in a given path.
1470
1471   This is an abstract method to search for filesystem object (files,
1472   dirs) under a given search path.
1473
1474   @type name: str
1475   @param name: the name to look for
1476   @type search_path: str
1477   @param search_path: location to start at
1478   @type test: callable
1479   @param test: a function taking one argument that should return True
1480       if the a given object is valid; the default value is
1481       os.path.exists, causing only existing files to be returned
1482   @rtype: str or None
1483   @return: full path to the object if found, None otherwise
1484
1485   """
1486   for dir_name in search_path:
1487     item_name = os.path.sep.join([dir_name, name])
1488     if test(item_name):
1489       return item_name
1490   return None
1491
1492
1493 def CheckVolumeGroupSize(vglist, vgname, minsize):
1494   """Checks if the volume group list is valid.
1495
1496   The function will check if a given volume group is in the list of
1497   volume groups and has a minimum size.
1498
1499   @type vglist: dict
1500   @param vglist: dictionary of volume group names and their size
1501   @type vgname: str
1502   @param vgname: the volume group we should check
1503   @type minsize: int
1504   @param minsize: the minimum size we accept
1505   @rtype: None or str
1506   @return: None for success, otherwise the error message
1507
1508   """
1509   vgsize = vglist.get(vgname, None)
1510   if vgsize is None:
1511     return "volume group '%s' missing" % vgname
1512   elif vgsize < minsize:
1513     return ("volume group '%s' too small (%s MiB required, %d MiB found)" %
1514             (vgname, minsize, vgsize))
1515   return None
1516
1517
1518 def SplitTime(value):
1519   """Splits time as floating point number into a tuple.
1520
1521   @param value: Time in seconds
1522   @type value: int or float
1523   @return: Tuple containing (seconds, microseconds)
1524
1525   """
1526   (seconds, microseconds) = divmod(int(value * 1000000), 1000000)
1527
1528   assert 0 <= seconds, \
1529     "Seconds must be larger than or equal to 0, but are %s" % seconds
1530   assert 0 <= microseconds <= 999999, \
1531     "Microseconds must be 0-999999, but are %s" % microseconds
1532
1533   return (int(seconds), int(microseconds))
1534
1535
1536 def MergeTime(timetuple):
1537   """Merges a tuple into time as a floating point number.
1538
1539   @param timetuple: Time as tuple, (seconds, microseconds)
1540   @type timetuple: tuple
1541   @return: Time as a floating point number expressed in seconds
1542
1543   """
1544   (seconds, microseconds) = timetuple
1545
1546   assert 0 <= seconds, \
1547     "Seconds must be larger than or equal to 0, but are %s" % seconds
1548   assert 0 <= microseconds <= 999999, \
1549     "Microseconds must be 0-999999, but are %s" % microseconds
1550
1551   return float(seconds) + (float(microseconds) * 0.000001)
1552
1553
1554 def GetNodeDaemonPort():
1555   """Get the node daemon port for this cluster.
1556
1557   Note that this routine does not read a ganeti-specific file, but
1558   instead uses C{socket.getservbyname} to allow pre-customization of
1559   this parameter outside of Ganeti.
1560
1561   @rtype: int
1562
1563   """
1564   try:
1565     port = socket.getservbyname("ganeti-noded", "tcp")
1566   except socket.error:
1567     port = constants.DEFAULT_NODED_PORT
1568
1569   return port
1570
1571
1572 def SetupLogging(logfile, debug=False, stderr_logging=False, program=""):
1573   """Configures the logging module.
1574
1575   @type logfile: str
1576   @param logfile: the filename to which we should log
1577   @type debug: boolean
1578   @param debug: whether to enable debug messages too or
1579       only those at C{INFO} and above level
1580   @type stderr_logging: boolean
1581   @param stderr_logging: whether we should also log to the standard error
1582   @type program: str
1583   @param program: the name under which we should log messages
1584   @raise EnvironmentError: if we can't open the log file and
1585       stderr logging is disabled
1586
1587   """
1588   fmt = "%(asctime)s: " + program + " "
1589   if debug:
1590     fmt += ("pid=%(process)d/%(threadName)s %(levelname)s"
1591            " %(module)s:%(lineno)s %(message)s")
1592   else:
1593     fmt += "pid=%(process)d %(levelname)s %(message)s"
1594   formatter = logging.Formatter(fmt)
1595
1596   root_logger = logging.getLogger("")
1597   root_logger.setLevel(logging.NOTSET)
1598
1599   # Remove all previously setup handlers
1600   for handler in root_logger.handlers:
1601     root_logger.removeHandler(handler)
1602
1603   if stderr_logging:
1604     stderr_handler = logging.StreamHandler()
1605     stderr_handler.setFormatter(formatter)
1606     if debug:
1607       stderr_handler.setLevel(logging.NOTSET)
1608     else:
1609       stderr_handler.setLevel(logging.CRITICAL)
1610     root_logger.addHandler(stderr_handler)
1611
1612   # this can fail, if the logging directories are not setup or we have
1613   # a permisssion problem; in this case, it's best to log but ignore
1614   # the error if stderr_logging is True, and if false we re-raise the
1615   # exception since otherwise we could run but without any logs at all
1616   try:
1617     logfile_handler = logging.FileHandler(logfile)
1618     logfile_handler.setFormatter(formatter)
1619     if debug:
1620       logfile_handler.setLevel(logging.DEBUG)
1621     else:
1622       logfile_handler.setLevel(logging.INFO)
1623     root_logger.addHandler(logfile_handler)
1624   except EnvironmentError, err:
1625     if stderr_logging:
1626       logging.exception("Failed to enable logging to file '%s'", logfile)
1627     else:
1628       # we need to re-raise the exception
1629       raise
1630
1631
1632 def LockedMethod(fn):
1633   """Synchronized object access decorator.
1634
1635   This decorator is intended to protect access to an object using the
1636   object's own lock which is hardcoded to '_lock'.
1637
1638   """
1639   def _LockDebug(*args, **kwargs):
1640     if debug_locks:
1641       logging.debug(*args, **kwargs)
1642
1643   def wrapper(self, *args, **kwargs):
1644     assert hasattr(self, '_lock')
1645     lock = self._lock
1646     _LockDebug("Waiting for %s", lock)
1647     lock.acquire()
1648     try:
1649       _LockDebug("Acquired %s", lock)
1650       result = fn(self, *args, **kwargs)
1651     finally:
1652       _LockDebug("Releasing %s", lock)
1653       lock.release()
1654       _LockDebug("Released %s", lock)
1655     return result
1656   return wrapper
1657
1658
1659 def LockFile(fd):
1660   """Locks a file using POSIX locks.
1661
1662   @type fd: int
1663   @param fd: the file descriptor we need to lock
1664
1665   """
1666   try:
1667     fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
1668   except IOError, err:
1669     if err.errno == errno.EAGAIN:
1670       raise errors.LockError("File already locked")
1671     raise
1672
1673
1674 class FileLock(object):
1675   """Utility class for file locks.
1676
1677   """
1678   def __init__(self, filename):
1679     """Constructor for FileLock.
1680
1681     This will open the file denoted by the I{filename} argument.
1682
1683     @type filename: str
1684     @param filename: path to the file to be locked
1685
1686     """
1687     self.filename = filename
1688     self.fd = open(self.filename, "w")
1689
1690   def __del__(self):
1691     self.Close()
1692
1693   def Close(self):
1694     """Close the file and release the lock.
1695
1696     """
1697     if self.fd:
1698       self.fd.close()
1699       self.fd = None
1700
1701   def _flock(self, flag, blocking, timeout, errmsg):
1702     """Wrapper for fcntl.flock.
1703
1704     @type flag: int
1705     @param flag: operation flag
1706     @type blocking: bool
1707     @param blocking: whether the operation should be done in blocking mode.
1708     @type timeout: None or float
1709     @param timeout: for how long the operation should be retried (implies
1710                     non-blocking mode).
1711     @type errmsg: string
1712     @param errmsg: error message in case operation fails.
1713
1714     """
1715     assert self.fd, "Lock was closed"
1716     assert timeout is None or timeout >= 0, \
1717       "If specified, timeout must be positive"
1718
1719     if timeout is not None:
1720       flag |= fcntl.LOCK_NB
1721       timeout_end = time.time() + timeout
1722
1723     # Blocking doesn't have effect with timeout
1724     elif not blocking:
1725       flag |= fcntl.LOCK_NB
1726       timeout_end = None
1727
1728     retry = True
1729     while retry:
1730       try:
1731         fcntl.flock(self.fd, flag)
1732         retry = False
1733       except IOError, err:
1734         if err.errno in (errno.EAGAIN, ):
1735           if timeout_end is not None and time.time() < timeout_end:
1736             # Wait before trying again
1737             time.sleep(max(0.1, min(1.0, timeout)))
1738           else:
1739             raise errors.LockError(errmsg)
1740         else:
1741           logging.exception("fcntl.flock failed")
1742           raise
1743
1744   def Exclusive(self, blocking=False, timeout=None):
1745     """Locks the file in exclusive mode.
1746
1747     @type blocking: boolean
1748     @param blocking: whether to block and wait until we
1749         can lock the file or return immediately
1750     @type timeout: int or None
1751     @param timeout: if not None, the duration to wait for the lock
1752         (in blocking mode)
1753
1754     """
1755     self._flock(fcntl.LOCK_EX, blocking, timeout,
1756                 "Failed to lock %s in exclusive mode" % self.filename)
1757
1758   def Shared(self, blocking=False, timeout=None):
1759     """Locks the file in shared mode.
1760
1761     @type blocking: boolean
1762     @param blocking: whether to block and wait until we
1763         can lock the file or return immediately
1764     @type timeout: int or None
1765     @param timeout: if not None, the duration to wait for the lock
1766         (in blocking mode)
1767
1768     """
1769     self._flock(fcntl.LOCK_SH, blocking, timeout,
1770                 "Failed to lock %s in shared mode" % self.filename)
1771
1772   def Unlock(self, blocking=True, timeout=None):
1773     """Unlocks the file.
1774
1775     According to C{flock(2)}, unlocking can also be a nonblocking
1776     operation::
1777
1778       To make a non-blocking request, include LOCK_NB with any of the above
1779       operations.
1780
1781     @type blocking: boolean
1782     @param blocking: whether to block and wait until we
1783         can lock the file or return immediately
1784     @type timeout: int or None
1785     @param timeout: if not None, the duration to wait for the lock
1786         (in blocking mode)
1787
1788     """
1789     self._flock(fcntl.LOCK_UN, blocking, timeout,
1790                 "Failed to unlock %s" % self.filename)
1791
1792
1793 class SignalHandler(object):
1794   """Generic signal handler class.
1795
1796   It automatically restores the original handler when deconstructed or
1797   when L{Reset} is called. You can either pass your own handler
1798   function in or query the L{called} attribute to detect whether the
1799   signal was sent.
1800
1801   @type signum: list
1802   @ivar signum: the signals we handle
1803   @type called: boolean
1804   @ivar called: tracks whether any of the signals have been raised
1805
1806   """
1807   def __init__(self, signum):
1808     """Constructs a new SignalHandler instance.
1809
1810     @type signum: int or list of ints
1811     @param signum: Single signal number or set of signal numbers
1812
1813     """
1814     if isinstance(signum, (int, long)):
1815       self.signum = set([signum])
1816     else:
1817       self.signum = set(signum)
1818
1819     self.called = False
1820
1821     self._previous = {}
1822     try:
1823       for signum in self.signum:
1824         # Setup handler
1825         prev_handler = signal.signal(signum, self._HandleSignal)
1826         try:
1827           self._previous[signum] = prev_handler
1828         except:
1829           # Restore previous handler
1830           signal.signal(signum, prev_handler)
1831           raise
1832     except:
1833       # Reset all handlers
1834       self.Reset()
1835       # Here we have a race condition: a handler may have already been called,
1836       # but there's not much we can do about it at this point.
1837       raise
1838
1839   def __del__(self):
1840     self.Reset()
1841
1842   def Reset(self):
1843     """Restore previous handler.
1844
1845     This will reset all the signals to their previous handlers.
1846
1847     """
1848     for signum, prev_handler in self._previous.items():
1849       signal.signal(signum, prev_handler)
1850       # If successful, remove from dict
1851       del self._previous[signum]
1852
1853   def Clear(self):
1854     """Unsets the L{called} flag.
1855
1856     This function can be used in case a signal may arrive several times.
1857
1858     """
1859     self.called = False
1860
1861   def _HandleSignal(self, signum, frame):
1862     """Actual signal handling function.
1863
1864     """
1865     # This is not nice and not absolutely atomic, but it appears to be the only
1866     # solution in Python -- there are no atomic types.
1867     self.called = True
1868
1869
1870 class FieldSet(object):
1871   """A simple field set.
1872
1873   Among the features are:
1874     - checking if a string is among a list of static string or regex objects
1875     - checking if a whole list of string matches
1876     - returning the matching groups from a regex match
1877
1878   Internally, all fields are held as regular expression objects.
1879
1880   """
1881   def __init__(self, *items):
1882     self.items = [re.compile("^%s$" % value) for value in items]
1883
1884   def Extend(self, other_set):
1885     """Extend the field set with the items from another one"""
1886     self.items.extend(other_set.items)
1887
1888   def Matches(self, field):
1889     """Checks if a field matches the current set
1890
1891     @type field: str
1892     @param field: the string to match
1893     @return: either False or a regular expression match object
1894
1895     """
1896     for m in itertools.ifilter(None, (val.match(field) for val in self.items)):
1897       return m
1898     return False
1899
1900   def NonMatching(self, items):
1901     """Returns the list of fields not matching the current set
1902
1903     @type items: list
1904     @param items: the list of fields to check
1905     @rtype: list
1906     @return: list of non-matching fields
1907
1908     """
1909     return [val for val in items if not self.Matches(val)]