Statistics
| Branch: | Tag: | Revision:

root / lib / utils.py @ 3cd62121

History | View | Annotate | Download (30.7 kB)

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

24 a8083063 Iustin Pop
"""
25 a8083063 Iustin Pop
26 a8083063 Iustin Pop
27 a8083063 Iustin Pop
import sys
28 a8083063 Iustin Pop
import os
29 a8083063 Iustin Pop
import sha
30 a8083063 Iustin Pop
import time
31 113b55aa Iustin Pop
import subprocess
32 a8083063 Iustin Pop
import re
33 a8083063 Iustin Pop
import socket
34 a8083063 Iustin Pop
import tempfile
35 a8083063 Iustin Pop
import shutil
36 4ca1b175 Alexander Schreiber
import errno
37 2f8b60b3 Iustin Pop
import pwd
38 78feb6fb Guido Trotter
import itertools
39 9c233417 Iustin Pop
import select
40 9c233417 Iustin Pop
import fcntl
41 8f765069 Iustin Pop
import resource
42 bb698c1f Iustin Pop
import logging
43 de499029 Michael Hanselmann
import signal
44 9c233417 Iustin Pop
45 9c233417 Iustin Pop
from cStringIO import StringIO
46 a8083063 Iustin Pop
47 a8083063 Iustin Pop
from ganeti import errors
48 3aecd2c7 Iustin Pop
from ganeti import constants
49 a8083063 Iustin Pop
50 16abfbc2 Alexander Schreiber
51 a8083063 Iustin Pop
_locksheld = []
52 a8083063 Iustin Pop
_re_shell_unquoted = re.compile('^[-.,=:/_+@A-Za-z0-9]+$')
53 a8083063 Iustin Pop
54 f362096f Iustin Pop
debug = False
55 b74159ee Iustin Pop
no_fork = False
56 f362096f Iustin Pop
57 7c0d6283 Michael Hanselmann
58 a8083063 Iustin Pop
class RunResult(object):
59 a8083063 Iustin Pop
  """Simple class for holding the result of running external programs.
60 a8083063 Iustin Pop

61 a8083063 Iustin Pop
  Instance variables:
62 a8083063 Iustin Pop
    exit_code: the exit code of the program, or None (if the program
63 a8083063 Iustin Pop
               didn't exit())
64 a8083063 Iustin Pop
    signal: numeric signal that caused the program to finish, or None
65 a8083063 Iustin Pop
            (if the program wasn't terminated by a signal)
66 a8083063 Iustin Pop
    stdout: the standard output of the program
67 a8083063 Iustin Pop
    stderr: the standard error of the program
68 a8083063 Iustin Pop
    failed: a Boolean value which is True in case the program was
69 a8083063 Iustin Pop
            terminated by a signal or exited with a non-zero exit code
70 a8083063 Iustin Pop
    fail_reason: a string detailing the termination reason
71 a8083063 Iustin Pop

72 a8083063 Iustin Pop
  """
73 a8083063 Iustin Pop
  __slots__ = ["exit_code", "signal", "stdout", "stderr",
74 a8083063 Iustin Pop
               "failed", "fail_reason", "cmd"]
75 a8083063 Iustin Pop
76 a8083063 Iustin Pop
77 a8083063 Iustin Pop
  def __init__(self, exit_code, signal, stdout, stderr, cmd):
78 a8083063 Iustin Pop
    self.cmd = cmd
79 a8083063 Iustin Pop
    self.exit_code = exit_code
80 a8083063 Iustin Pop
    self.signal = signal
81 a8083063 Iustin Pop
    self.stdout = stdout
82 a8083063 Iustin Pop
    self.stderr = stderr
83 a8083063 Iustin Pop
    self.failed = (signal is not None or exit_code != 0)
84 a8083063 Iustin Pop
85 a8083063 Iustin Pop
    if self.signal is not None:
86 a8083063 Iustin Pop
      self.fail_reason = "terminated by signal %s" % self.signal
87 a8083063 Iustin Pop
    elif self.exit_code is not None:
88 a8083063 Iustin Pop
      self.fail_reason = "exited with exit code %s" % self.exit_code
89 a8083063 Iustin Pop
    else:
90 a8083063 Iustin Pop
      self.fail_reason = "unable to determine termination reason"
91 a8083063 Iustin Pop
92 bb698c1f Iustin Pop
    if self.failed:
93 bb698c1f Iustin Pop
      logging.debug("Command '%s' failed (%s); output: %s",
94 bb698c1f Iustin Pop
                    self.cmd, self.fail_reason, self.output)
95 f362096f Iustin Pop
96 a8083063 Iustin Pop
  def _GetOutput(self):
97 a8083063 Iustin Pop
    """Returns the combined stdout and stderr for easier usage.
98 a8083063 Iustin Pop

99 a8083063 Iustin Pop
    """
100 a8083063 Iustin Pop
    return self.stdout + self.stderr
101 a8083063 Iustin Pop
102 a8083063 Iustin Pop
  output = property(_GetOutput, None, None, "Return full output")
103 a8083063 Iustin Pop
104 a8083063 Iustin Pop
105 a8083063 Iustin Pop
def RunCmd(cmd):
106 a8083063 Iustin Pop
  """Execute a (shell) command.
107 a8083063 Iustin Pop

108 a8083063 Iustin Pop
  The command should not read from its standard input, as it will be
109 a8083063 Iustin Pop
  closed.
110 a8083063 Iustin Pop

111 a8083063 Iustin Pop
  Args:
112 a8083063 Iustin Pop
    cmd: command to run. (str)
113 a8083063 Iustin Pop

114 a8083063 Iustin Pop
  Returns: `RunResult` instance
115 a8083063 Iustin Pop

116 a8083063 Iustin Pop
  """
117 b74159ee Iustin Pop
  if no_fork:
118 b74159ee Iustin Pop
    raise errors.ProgrammerError("utils.RunCmd() called with fork() disabled")
119 b74159ee Iustin Pop
120 a8083063 Iustin Pop
  if isinstance(cmd, list):
121 a8083063 Iustin Pop
    cmd = [str(val) for val in cmd]
122 113b55aa Iustin Pop
    strcmd = " ".join(cmd)
123 113b55aa Iustin Pop
    shell = False
124 113b55aa Iustin Pop
  else:
125 113b55aa Iustin Pop
    strcmd = cmd
126 113b55aa Iustin Pop
    shell = True
127 bb698c1f Iustin Pop
  logging.debug("RunCmd '%s'", strcmd)
128 23f41a3e Michael Hanselmann
  env = os.environ.copy()
129 23f41a3e Michael Hanselmann
  env["LC_ALL"] = "C"
130 9c233417 Iustin Pop
  poller = select.poll()
131 113b55aa Iustin Pop
  child = subprocess.Popen(cmd, shell=shell,
132 113b55aa Iustin Pop
                           stderr=subprocess.PIPE,
133 113b55aa Iustin Pop
                           stdout=subprocess.PIPE,
134 113b55aa Iustin Pop
                           stdin=subprocess.PIPE,
135 23f41a3e Michael Hanselmann
                           close_fds=True, env=env)
136 113b55aa Iustin Pop
137 113b55aa Iustin Pop
  child.stdin.close()
138 9c233417 Iustin Pop
  poller.register(child.stdout, select.POLLIN)
139 9c233417 Iustin Pop
  poller.register(child.stderr, select.POLLIN)
140 9c233417 Iustin Pop
  out = StringIO()
141 9c233417 Iustin Pop
  err = StringIO()
142 9c233417 Iustin Pop
  fdmap = {
143 9c233417 Iustin Pop
    child.stdout.fileno(): (out, child.stdout),
144 9c233417 Iustin Pop
    child.stderr.fileno(): (err, child.stderr),
145 9c233417 Iustin Pop
    }
146 9c233417 Iustin Pop
  for fd in fdmap:
147 9c233417 Iustin Pop
    status = fcntl.fcntl(fd, fcntl.F_GETFL)
148 9c233417 Iustin Pop
    fcntl.fcntl(fd, fcntl.F_SETFL, status | os.O_NONBLOCK)
149 9c233417 Iustin Pop
150 9c233417 Iustin Pop
  while fdmap:
151 9c233417 Iustin Pop
    for fd, event in poller.poll():
152 9c233417 Iustin Pop
      if event & select.POLLIN or event & select.POLLPRI:
153 9c233417 Iustin Pop
        data = fdmap[fd][1].read()
154 9c233417 Iustin Pop
        # no data from read signifies EOF (the same as POLLHUP)
155 9c233417 Iustin Pop
        if not data:
156 9c233417 Iustin Pop
          poller.unregister(fd)
157 9c233417 Iustin Pop
          del fdmap[fd]
158 9c233417 Iustin Pop
          continue
159 9c233417 Iustin Pop
        fdmap[fd][0].write(data)
160 9c233417 Iustin Pop
      if (event & select.POLLNVAL or event & select.POLLHUP or
161 9c233417 Iustin Pop
          event & select.POLLERR):
162 9c233417 Iustin Pop
        poller.unregister(fd)
163 9c233417 Iustin Pop
        del fdmap[fd]
164 9c233417 Iustin Pop
165 9c233417 Iustin Pop
  out = out.getvalue()
166 9c233417 Iustin Pop
  err = err.getvalue()
167 a8083063 Iustin Pop
168 a8083063 Iustin Pop
  status = child.wait()
169 113b55aa Iustin Pop
  if status >= 0:
170 113b55aa Iustin Pop
    exitcode = status
171 a8083063 Iustin Pop
    signal = None
172 a8083063 Iustin Pop
  else:
173 a8083063 Iustin Pop
    exitcode = None
174 113b55aa Iustin Pop
    signal = -status
175 a8083063 Iustin Pop
176 a8083063 Iustin Pop
  return RunResult(exitcode, signal, out, err, strcmd)
177 a8083063 Iustin Pop
178 a8083063 Iustin Pop
179 a8083063 Iustin Pop
def RemoveFile(filename):
180 a8083063 Iustin Pop
  """Remove a file ignoring some errors.
181 a8083063 Iustin Pop

182 a8083063 Iustin Pop
  Remove a file, ignoring non-existing ones or directories. Other
183 a8083063 Iustin Pop
  errors are passed.
184 a8083063 Iustin Pop

185 a8083063 Iustin Pop
  """
186 a8083063 Iustin Pop
  try:
187 a8083063 Iustin Pop
    os.unlink(filename)
188 a8083063 Iustin Pop
  except OSError, err:
189 4ca1b175 Alexander Schreiber
    if err.errno not in (errno.ENOENT, errno.EISDIR):
190 a8083063 Iustin Pop
      raise
191 a8083063 Iustin Pop
192 a8083063 Iustin Pop
193 a8083063 Iustin Pop
def _FingerprintFile(filename):
194 a8083063 Iustin Pop
  """Compute the fingerprint of a file.
195 a8083063 Iustin Pop

196 a8083063 Iustin Pop
  If the file does not exist, a None will be returned
197 a8083063 Iustin Pop
  instead.
198 a8083063 Iustin Pop

199 a8083063 Iustin Pop
  Args:
200 a8083063 Iustin Pop
    filename - Filename (str)
201 a8083063 Iustin Pop

202 a8083063 Iustin Pop
  """
203 a8083063 Iustin Pop
  if not (os.path.exists(filename) and os.path.isfile(filename)):
204 a8083063 Iustin Pop
    return None
205 a8083063 Iustin Pop
206 a8083063 Iustin Pop
  f = open(filename)
207 a8083063 Iustin Pop
208 a8083063 Iustin Pop
  fp = sha.sha()
209 a8083063 Iustin Pop
  while True:
210 a8083063 Iustin Pop
    data = f.read(4096)
211 a8083063 Iustin Pop
    if not data:
212 a8083063 Iustin Pop
      break
213 a8083063 Iustin Pop
214 a8083063 Iustin Pop
    fp.update(data)
215 a8083063 Iustin Pop
216 a8083063 Iustin Pop
  return fp.hexdigest()
217 a8083063 Iustin Pop
218 a8083063 Iustin Pop
219 a8083063 Iustin Pop
def FingerprintFiles(files):
220 a8083063 Iustin Pop
  """Compute fingerprints for a list of files.
221 a8083063 Iustin Pop

222 a8083063 Iustin Pop
  Args:
223 a8083063 Iustin Pop
    files - array of filenames.  ( [str, ...] )
224 a8083063 Iustin Pop

225 a8083063 Iustin Pop
  Return value:
226 a8083063 Iustin Pop
    dictionary of filename: fingerprint for the files that exist
227 a8083063 Iustin Pop

228 a8083063 Iustin Pop
  """
229 a8083063 Iustin Pop
  ret = {}
230 a8083063 Iustin Pop
231 a8083063 Iustin Pop
  for filename in files:
232 a8083063 Iustin Pop
    cksum = _FingerprintFile(filename)
233 a8083063 Iustin Pop
    if cksum:
234 a8083063 Iustin Pop
      ret[filename] = cksum
235 a8083063 Iustin Pop
236 a8083063 Iustin Pop
  return ret
237 a8083063 Iustin Pop
238 a8083063 Iustin Pop
239 a8083063 Iustin Pop
def CheckDict(target, template, logname=None):
240 a8083063 Iustin Pop
  """Ensure a dictionary has a required set of keys.
241 a8083063 Iustin Pop

242 a8083063 Iustin Pop
  For the given dictionaries `target` and `template`, ensure target
243 a8083063 Iustin Pop
  has all the keys from template. Missing keys are added with values
244 a8083063 Iustin Pop
  from template.
245 a8083063 Iustin Pop

246 a8083063 Iustin Pop
  Args:
247 a8083063 Iustin Pop
    target   - the dictionary to check
248 a8083063 Iustin Pop
    template - template dictionary
249 a8083063 Iustin Pop
    logname  - a caller-chosen string to identify the debug log
250 a8083063 Iustin Pop
               entry; if None, no logging will be done
251 a8083063 Iustin Pop

252 a8083063 Iustin Pop
  Returns value:
253 a8083063 Iustin Pop
    None
254 a8083063 Iustin Pop

255 a8083063 Iustin Pop
  """
256 a8083063 Iustin Pop
  missing = []
257 a8083063 Iustin Pop
  for k in template:
258 a8083063 Iustin Pop
    if k not in target:
259 a8083063 Iustin Pop
      missing.append(k)
260 a8083063 Iustin Pop
      target[k] = template[k]
261 a8083063 Iustin Pop
262 a8083063 Iustin Pop
  if missing and logname:
263 bb698c1f Iustin Pop
    logging.warning('%s missing keys %s', logname, ', '.join(missing))
264 a8083063 Iustin Pop
265 a8083063 Iustin Pop
266 a8083063 Iustin Pop
def IsProcessAlive(pid):
267 a8083063 Iustin Pop
  """Check if a given pid exists on the system.
268 a8083063 Iustin Pop

269 a8083063 Iustin Pop
  Returns: true or false, depending on if the pid exists or not
270 a8083063 Iustin Pop

271 a8083063 Iustin Pop
  Remarks: zombie processes treated as not alive
272 a8083063 Iustin Pop

273 a8083063 Iustin Pop
  """
274 a8083063 Iustin Pop
  try:
275 a8083063 Iustin Pop
    f = open("/proc/%d/status" % pid)
276 a8083063 Iustin Pop
  except IOError, err:
277 4ca1b175 Alexander Schreiber
    if err.errno in (errno.ENOENT, errno.ENOTDIR):
278 a8083063 Iustin Pop
      return False
279 a8083063 Iustin Pop
280 a8083063 Iustin Pop
  alive = True
281 a8083063 Iustin Pop
  try:
282 a8083063 Iustin Pop
    data = f.readlines()
283 a8083063 Iustin Pop
    if len(data) > 1:
284 a8083063 Iustin Pop
      state = data[1].split()
285 a8083063 Iustin Pop
      if len(state) > 1 and state[1] == "Z":
286 a8083063 Iustin Pop
        alive = False
287 a8083063 Iustin Pop
  finally:
288 a8083063 Iustin Pop
    f.close()
289 a8083063 Iustin Pop
290 a8083063 Iustin Pop
  return alive
291 a8083063 Iustin Pop
292 a8083063 Iustin Pop
293 fee80e90 Guido Trotter
def IsPidFileAlive(pidfile):
294 fee80e90 Guido Trotter
  """Check whether the given pidfile points to a live process.
295 fee80e90 Guido Trotter

296 fee80e90 Guido Trotter
    @param pidfile: Path to a file containing the pid to be checked
297 fee80e90 Guido Trotter
    @type  pidfile: string (filename)
298 fee80e90 Guido Trotter

299 fee80e90 Guido Trotter
  """
300 fee80e90 Guido Trotter
  try:
301 fee80e90 Guido Trotter
    pf = open(pidfile, 'r')
302 fee80e90 Guido Trotter
  except EnvironmentError, open_err:
303 fee80e90 Guido Trotter
    if open_err.errno == errno.ENOENT:
304 fee80e90 Guido Trotter
      return False
305 fee80e90 Guido Trotter
    else:
306 fee80e90 Guido Trotter
      raise errors.GenericError("Cannot open file %s. Error: %s" %
307 fee80e90 Guido Trotter
                                (pidfile, str(open_err)))
308 fee80e90 Guido Trotter
309 fee80e90 Guido Trotter
  try:
310 fee80e90 Guido Trotter
    pid = int(pf.read())
311 fee80e90 Guido Trotter
  except ValueError:
312 fee80e90 Guido Trotter
    raise errors.GenericError("Invalid pid string in %s" %
313 fee80e90 Guido Trotter
                              (pidfile,))
314 fee80e90 Guido Trotter
315 fee80e90 Guido Trotter
  return IsProcessAlive(pid)
316 fee80e90 Guido Trotter
317 fee80e90 Guido Trotter
318 a8083063 Iustin Pop
def MatchNameComponent(key, name_list):
319 a8083063 Iustin Pop
  """Try to match a name against a list.
320 a8083063 Iustin Pop

321 a8083063 Iustin Pop
  This function will try to match a name like test1 against a list
322 a8083063 Iustin Pop
  like ['test1.example.com', 'test2.example.com', ...]. Against this
323 a8083063 Iustin Pop
  list, 'test1' as well as 'test1.example' will match, but not
324 a8083063 Iustin Pop
  'test1.ex'. A multiple match will be considered as no match at all
325 a8083063 Iustin Pop
  (e.g. 'test1' against ['test1.example.com', 'test1.example.org']).
326 a8083063 Iustin Pop

327 a8083063 Iustin Pop
  Args:
328 a8083063 Iustin Pop
    key: the name to be searched
329 a8083063 Iustin Pop
    name_list: the list of strings against which to search the key
330 a8083063 Iustin Pop

331 a8083063 Iustin Pop
  Returns:
332 a8083063 Iustin Pop
    None if there is no match *or* if there are multiple matches
333 a8083063 Iustin Pop
    otherwise the element from the list which matches
334 a8083063 Iustin Pop

335 a8083063 Iustin Pop
  """
336 a8083063 Iustin Pop
  mo = re.compile("^%s(\..*)?$" % re.escape(key))
337 a8083063 Iustin Pop
  names_filtered = [name for name in name_list if mo.match(name) is not None]
338 a8083063 Iustin Pop
  if len(names_filtered) != 1:
339 a8083063 Iustin Pop
    return None
340 a8083063 Iustin Pop
  return names_filtered[0]
341 a8083063 Iustin Pop
342 a8083063 Iustin Pop
343 bcf043c9 Iustin Pop
class HostInfo:
344 89e1fc26 Iustin Pop
  """Class implementing resolver and hostname functionality
345 bcf043c9 Iustin Pop

346 bcf043c9 Iustin Pop
  """
347 89e1fc26 Iustin Pop
  def __init__(self, name=None):
348 bcf043c9 Iustin Pop
    """Initialize the host name object.
349 bcf043c9 Iustin Pop

350 89e1fc26 Iustin Pop
    If the name argument is not passed, it will use this system's
351 89e1fc26 Iustin Pop
    name.
352 bcf043c9 Iustin Pop

353 bcf043c9 Iustin Pop
    """
354 89e1fc26 Iustin Pop
    if name is None:
355 89e1fc26 Iustin Pop
      name = self.SysName()
356 89e1fc26 Iustin Pop
357 89e1fc26 Iustin Pop
    self.query = name
358 89e1fc26 Iustin Pop
    self.name, self.aliases, self.ipaddrs = self.LookupHostname(name)
359 bcf043c9 Iustin Pop
    self.ip = self.ipaddrs[0]
360 bcf043c9 Iustin Pop
361 c8a0948f Michael Hanselmann
  def ShortName(self):
362 c8a0948f Michael Hanselmann
    """Returns the hostname without domain.
363 c8a0948f Michael Hanselmann

364 c8a0948f Michael Hanselmann
    """
365 c8a0948f Michael Hanselmann
    return self.name.split('.')[0]
366 c8a0948f Michael Hanselmann
367 89e1fc26 Iustin Pop
  @staticmethod
368 89e1fc26 Iustin Pop
  def SysName():
369 89e1fc26 Iustin Pop
    """Return the current system's name.
370 bcf043c9 Iustin Pop

371 89e1fc26 Iustin Pop
    This is simply a wrapper over socket.gethostname()
372 a8083063 Iustin Pop

373 89e1fc26 Iustin Pop
    """
374 89e1fc26 Iustin Pop
    return socket.gethostname()
375 a8083063 Iustin Pop
376 89e1fc26 Iustin Pop
  @staticmethod
377 89e1fc26 Iustin Pop
  def LookupHostname(hostname):
378 89e1fc26 Iustin Pop
    """Look up hostname
379 a8083063 Iustin Pop

380 89e1fc26 Iustin Pop
    Args:
381 89e1fc26 Iustin Pop
      hostname: hostname to look up
382 89e1fc26 Iustin Pop

383 89e1fc26 Iustin Pop
    Returns:
384 89e1fc26 Iustin Pop
      a tuple (name, aliases, ipaddrs) as returned by socket.gethostbyname_ex
385 89e1fc26 Iustin Pop
      in case of errors in resolving, we raise a ResolverError
386 89e1fc26 Iustin Pop

387 89e1fc26 Iustin Pop
    """
388 89e1fc26 Iustin Pop
    try:
389 89e1fc26 Iustin Pop
      result = socket.gethostbyname_ex(hostname)
390 89e1fc26 Iustin Pop
    except socket.gaierror, err:
391 89e1fc26 Iustin Pop
      # hostname not found in DNS
392 89e1fc26 Iustin Pop
      raise errors.ResolverError(hostname, err.args[0], err.args[1])
393 a8083063 Iustin Pop
394 89e1fc26 Iustin Pop
    return result
395 a8083063 Iustin Pop
396 a8083063 Iustin Pop
397 a8083063 Iustin Pop
def ListVolumeGroups():
398 a8083063 Iustin Pop
  """List volume groups and their size
399 a8083063 Iustin Pop

400 a8083063 Iustin Pop
  Returns:
401 a8083063 Iustin Pop
     Dictionary with keys volume name and values the size of the volume
402 a8083063 Iustin Pop

403 a8083063 Iustin Pop
  """
404 a8083063 Iustin Pop
  command = "vgs --noheadings --units m --nosuffix -o name,size"
405 a8083063 Iustin Pop
  result = RunCmd(command)
406 a8083063 Iustin Pop
  retval = {}
407 a8083063 Iustin Pop
  if result.failed:
408 a8083063 Iustin Pop
    return retval
409 a8083063 Iustin Pop
410 a8083063 Iustin Pop
  for line in result.stdout.splitlines():
411 a8083063 Iustin Pop
    try:
412 a8083063 Iustin Pop
      name, size = line.split()
413 a8083063 Iustin Pop
      size = int(float(size))
414 a8083063 Iustin Pop
    except (IndexError, ValueError), err:
415 bb698c1f Iustin Pop
      logging.error("Invalid output from vgs (%s): %s", err, line)
416 a8083063 Iustin Pop
      continue
417 a8083063 Iustin Pop
418 a8083063 Iustin Pop
    retval[name] = size
419 a8083063 Iustin Pop
420 a8083063 Iustin Pop
  return retval
421 a8083063 Iustin Pop
422 a8083063 Iustin Pop
423 a8083063 Iustin Pop
def BridgeExists(bridge):
424 a8083063 Iustin Pop
  """Check whether the given bridge exists in the system
425 a8083063 Iustin Pop

426 a8083063 Iustin Pop
  Returns:
427 a8083063 Iustin Pop
     True if it does, false otherwise.
428 a8083063 Iustin Pop

429 a8083063 Iustin Pop
  """
430 a8083063 Iustin Pop
  return os.path.isdir("/sys/class/net/%s/bridge" % bridge)
431 a8083063 Iustin Pop
432 a8083063 Iustin Pop
433 a8083063 Iustin Pop
def NiceSort(name_list):
434 a8083063 Iustin Pop
  """Sort a list of strings based on digit and non-digit groupings.
435 a8083063 Iustin Pop

436 a8083063 Iustin Pop
  Given a list of names ['a1', 'a10', 'a11', 'a2'] this function will
437 a8083063 Iustin Pop
  sort the list in the logical order ['a1', 'a2', 'a10', 'a11'].
438 a8083063 Iustin Pop

439 a8083063 Iustin Pop
  The sort algorithm breaks each name in groups of either only-digits
440 a8083063 Iustin Pop
  or no-digits. Only the first eight such groups are considered, and
441 a8083063 Iustin Pop
  after that we just use what's left of the string.
442 a8083063 Iustin Pop

443 a8083063 Iustin Pop
  Return value
444 a8083063 Iustin Pop
    - a copy of the list sorted according to our algorithm
445 a8083063 Iustin Pop

446 a8083063 Iustin Pop
  """
447 a8083063 Iustin Pop
  _SORTER_BASE = "(\D+|\d+)"
448 a8083063 Iustin Pop
  _SORTER_FULL = "^%s%s?%s?%s?%s?%s?%s?%s?.*$" % (_SORTER_BASE, _SORTER_BASE,
449 a8083063 Iustin Pop
                                                  _SORTER_BASE, _SORTER_BASE,
450 a8083063 Iustin Pop
                                                  _SORTER_BASE, _SORTER_BASE,
451 a8083063 Iustin Pop
                                                  _SORTER_BASE, _SORTER_BASE)
452 a8083063 Iustin Pop
  _SORTER_RE = re.compile(_SORTER_FULL)
453 a8083063 Iustin Pop
  _SORTER_NODIGIT = re.compile("^\D*$")
454 a8083063 Iustin Pop
  def _TryInt(val):
455 a8083063 Iustin Pop
    """Attempts to convert a variable to integer."""
456 a8083063 Iustin Pop
    if val is None or _SORTER_NODIGIT.match(val):
457 a8083063 Iustin Pop
      return val
458 a8083063 Iustin Pop
    rval = int(val)
459 a8083063 Iustin Pop
    return rval
460 a8083063 Iustin Pop
461 a8083063 Iustin Pop
  to_sort = [([_TryInt(grp) for grp in _SORTER_RE.match(name).groups()], name)
462 a8083063 Iustin Pop
             for name in name_list]
463 a8083063 Iustin Pop
  to_sort.sort()
464 a8083063 Iustin Pop
  return [tup[1] for tup in to_sort]
465 a8083063 Iustin Pop
466 a8083063 Iustin Pop
467 a8083063 Iustin Pop
def TryConvert(fn, val):
468 a8083063 Iustin Pop
  """Try to convert a value ignoring errors.
469 a8083063 Iustin Pop

470 a8083063 Iustin Pop
  This function tries to apply function `fn` to `val`. If no
471 a8083063 Iustin Pop
  ValueError or TypeError exceptions are raised, it will return the
472 a8083063 Iustin Pop
  result, else it will return the original value. Any other exceptions
473 a8083063 Iustin Pop
  are propagated to the caller.
474 a8083063 Iustin Pop

475 a8083063 Iustin Pop
  """
476 a8083063 Iustin Pop
  try:
477 a8083063 Iustin Pop
    nv = fn(val)
478 a8083063 Iustin Pop
  except (ValueError, TypeError), err:
479 a8083063 Iustin Pop
    nv = val
480 a8083063 Iustin Pop
  return nv
481 a8083063 Iustin Pop
482 a8083063 Iustin Pop
483 a8083063 Iustin Pop
def IsValidIP(ip):
484 a8083063 Iustin Pop
  """Verifies the syntax of an IP address.
485 a8083063 Iustin Pop

486 a8083063 Iustin Pop
  This function checks if the ip address passes is valid or not based
487 a8083063 Iustin Pop
  on syntax (not ip range, class calculations or anything).
488 a8083063 Iustin Pop

489 a8083063 Iustin Pop
  """
490 a8083063 Iustin Pop
  unit = "(0|[1-9]\d{0,2})"
491 a8083063 Iustin Pop
  return re.match("^%s\.%s\.%s\.%s$" % (unit, unit, unit, unit), ip)
492 a8083063 Iustin Pop
493 a8083063 Iustin Pop
494 a8083063 Iustin Pop
def IsValidShellParam(word):
495 a8083063 Iustin Pop
  """Verifies is the given word is safe from the shell's p.o.v.
496 a8083063 Iustin Pop

497 a8083063 Iustin Pop
  This means that we can pass this to a command via the shell and be
498 a8083063 Iustin Pop
  sure that it doesn't alter the command line and is passed as such to
499 a8083063 Iustin Pop
  the actual command.
500 a8083063 Iustin Pop

501 a8083063 Iustin Pop
  Note that we are overly restrictive here, in order to be on the safe
502 a8083063 Iustin Pop
  side.
503 a8083063 Iustin Pop

504 a8083063 Iustin Pop
  """
505 a8083063 Iustin Pop
  return bool(re.match("^[-a-zA-Z0-9._+/:%@]+$", word))
506 a8083063 Iustin Pop
507 a8083063 Iustin Pop
508 a8083063 Iustin Pop
def BuildShellCmd(template, *args):
509 a8083063 Iustin Pop
  """Build a safe shell command line from the given arguments.
510 a8083063 Iustin Pop

511 a8083063 Iustin Pop
  This function will check all arguments in the args list so that they
512 a8083063 Iustin Pop
  are valid shell parameters (i.e. they don't contain shell
513 a8083063 Iustin Pop
  metacharaters). If everything is ok, it will return the result of
514 a8083063 Iustin Pop
  template % args.
515 a8083063 Iustin Pop

516 a8083063 Iustin Pop
  """
517 a8083063 Iustin Pop
  for word in args:
518 a8083063 Iustin Pop
    if not IsValidShellParam(word):
519 3ecf6786 Iustin Pop
      raise errors.ProgrammerError("Shell argument '%s' contains"
520 3ecf6786 Iustin Pop
                                   " invalid characters" % word)
521 a8083063 Iustin Pop
  return template % args
522 a8083063 Iustin Pop
523 a8083063 Iustin Pop
524 a8083063 Iustin Pop
def FormatUnit(value):
525 a8083063 Iustin Pop
  """Formats an incoming number of MiB with the appropriate unit.
526 a8083063 Iustin Pop

527 a8083063 Iustin Pop
  Value needs to be passed as a numeric type. Return value is always a string.
528 a8083063 Iustin Pop

529 a8083063 Iustin Pop
  """
530 a8083063 Iustin Pop
  if value < 1024:
531 a8083063 Iustin Pop
    return "%dM" % round(value, 0)
532 a8083063 Iustin Pop
533 a8083063 Iustin Pop
  elif value < (1024 * 1024):
534 a8083063 Iustin Pop
    return "%0.1fG" % round(float(value) / 1024, 1)
535 a8083063 Iustin Pop
536 a8083063 Iustin Pop
  else:
537 a8083063 Iustin Pop
    return "%0.1fT" % round(float(value) / 1024 / 1024, 1)
538 a8083063 Iustin Pop
539 a8083063 Iustin Pop
540 a8083063 Iustin Pop
def ParseUnit(input_string):
541 a8083063 Iustin Pop
  """Tries to extract number and scale from the given string.
542 a8083063 Iustin Pop

543 a8083063 Iustin Pop
  Input must be in the format NUMBER+ [DOT NUMBER+] SPACE* [UNIT]. If no unit
544 a8083063 Iustin Pop
  is specified, it defaults to MiB. Return value is always an int in MiB.
545 a8083063 Iustin Pop

546 a8083063 Iustin Pop
  """
547 a8083063 Iustin Pop
  m = re.match('^([.\d]+)\s*([a-zA-Z]+)?$', input_string)
548 a8083063 Iustin Pop
  if not m:
549 3ecf6786 Iustin Pop
    raise errors.UnitParseError("Invalid format")
550 a8083063 Iustin Pop
551 a8083063 Iustin Pop
  value = float(m.groups()[0])
552 a8083063 Iustin Pop
553 a8083063 Iustin Pop
  unit = m.groups()[1]
554 a8083063 Iustin Pop
  if unit:
555 a8083063 Iustin Pop
    lcunit = unit.lower()
556 a8083063 Iustin Pop
  else:
557 a8083063 Iustin Pop
    lcunit = 'm'
558 a8083063 Iustin Pop
559 a8083063 Iustin Pop
  if lcunit in ('m', 'mb', 'mib'):
560 a8083063 Iustin Pop
    # Value already in MiB
561 a8083063 Iustin Pop
    pass
562 a8083063 Iustin Pop
563 a8083063 Iustin Pop
  elif lcunit in ('g', 'gb', 'gib'):
564 a8083063 Iustin Pop
    value *= 1024
565 a8083063 Iustin Pop
566 a8083063 Iustin Pop
  elif lcunit in ('t', 'tb', 'tib'):
567 a8083063 Iustin Pop
    value *= 1024 * 1024
568 a8083063 Iustin Pop
569 a8083063 Iustin Pop
  else:
570 3ecf6786 Iustin Pop
    raise errors.UnitParseError("Unknown unit: %s" % unit)
571 a8083063 Iustin Pop
572 a8083063 Iustin Pop
  # Make sure we round up
573 a8083063 Iustin Pop
  if int(value) < value:
574 a8083063 Iustin Pop
    value += 1
575 a8083063 Iustin Pop
576 a8083063 Iustin Pop
  # Round up to the next multiple of 4
577 a8083063 Iustin Pop
  value = int(value)
578 a8083063 Iustin Pop
  if value % 4:
579 a8083063 Iustin Pop
    value += 4 - value % 4
580 a8083063 Iustin Pop
581 a8083063 Iustin Pop
  return value
582 a8083063 Iustin Pop
583 a8083063 Iustin Pop
584 a8083063 Iustin Pop
def AddAuthorizedKey(file_name, key):
585 a8083063 Iustin Pop
  """Adds an SSH public key to an authorized_keys file.
586 a8083063 Iustin Pop

587 a8083063 Iustin Pop
  Args:
588 a8083063 Iustin Pop
    file_name: Path to authorized_keys file
589 a8083063 Iustin Pop
    key: String containing key
590 a8083063 Iustin Pop
  """
591 a8083063 Iustin Pop
  key_fields = key.split()
592 a8083063 Iustin Pop
593 a8083063 Iustin Pop
  f = open(file_name, 'a+')
594 a8083063 Iustin Pop
  try:
595 a8083063 Iustin Pop
    nl = True
596 a8083063 Iustin Pop
    for line in f:
597 a8083063 Iustin Pop
      # Ignore whitespace changes
598 a8083063 Iustin Pop
      if line.split() == key_fields:
599 a8083063 Iustin Pop
        break
600 a8083063 Iustin Pop
      nl = line.endswith('\n')
601 a8083063 Iustin Pop
    else:
602 a8083063 Iustin Pop
      if not nl:
603 a8083063 Iustin Pop
        f.write("\n")
604 a8083063 Iustin Pop
      f.write(key.rstrip('\r\n'))
605 a8083063 Iustin Pop
      f.write("\n")
606 a8083063 Iustin Pop
      f.flush()
607 a8083063 Iustin Pop
  finally:
608 a8083063 Iustin Pop
    f.close()
609 a8083063 Iustin Pop
610 a8083063 Iustin Pop
611 a8083063 Iustin Pop
def RemoveAuthorizedKey(file_name, key):
612 a8083063 Iustin Pop
  """Removes an SSH public key from an authorized_keys file.
613 a8083063 Iustin Pop

614 a8083063 Iustin Pop
  Args:
615 a8083063 Iustin Pop
    file_name: Path to authorized_keys file
616 a8083063 Iustin Pop
    key: String containing key
617 a8083063 Iustin Pop
  """
618 a8083063 Iustin Pop
  key_fields = key.split()
619 a8083063 Iustin Pop
620 a8083063 Iustin Pop
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
621 a8083063 Iustin Pop
  try:
622 59f82e3f Michael Hanselmann
    out = os.fdopen(fd, 'w')
623 a8083063 Iustin Pop
    try:
624 59f82e3f Michael Hanselmann
      f = open(file_name, 'r')
625 59f82e3f Michael Hanselmann
      try:
626 59f82e3f Michael Hanselmann
        for line in f:
627 59f82e3f Michael Hanselmann
          # Ignore whitespace changes while comparing lines
628 59f82e3f Michael Hanselmann
          if line.split() != key_fields:
629 59f82e3f Michael Hanselmann
            out.write(line)
630 899d2a81 Michael Hanselmann
631 899d2a81 Michael Hanselmann
        out.flush()
632 899d2a81 Michael Hanselmann
        os.rename(tmpname, file_name)
633 899d2a81 Michael Hanselmann
      finally:
634 899d2a81 Michael Hanselmann
        f.close()
635 899d2a81 Michael Hanselmann
    finally:
636 899d2a81 Michael Hanselmann
      out.close()
637 899d2a81 Michael Hanselmann
  except:
638 899d2a81 Michael Hanselmann
    RemoveFile(tmpname)
639 899d2a81 Michael Hanselmann
    raise
640 899d2a81 Michael Hanselmann
641 899d2a81 Michael Hanselmann
642 9440aeab Michael Hanselmann
def SetEtcHostsEntry(file_name, ip, hostname, aliases):
643 9440aeab Michael Hanselmann
  """Sets the name of an IP address and hostname in /etc/hosts.
644 899d2a81 Michael Hanselmann

645 899d2a81 Michael Hanselmann
  """
646 7fbb1f65 Michael Hanselmann
  # Ensure aliases are unique
647 7fbb1f65 Michael Hanselmann
  aliases = UniqueSequence([hostname] + aliases)[1:]
648 7fbb1f65 Michael Hanselmann
649 9440aeab Michael Hanselmann
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
650 899d2a81 Michael Hanselmann
  try:
651 9440aeab Michael Hanselmann
    out = os.fdopen(fd, 'w')
652 9440aeab Michael Hanselmann
    try:
653 9440aeab Michael Hanselmann
      f = open(file_name, 'r')
654 9440aeab Michael Hanselmann
      try:
655 9440aeab Michael Hanselmann
        written = False
656 9440aeab Michael Hanselmann
        for line in f:
657 9440aeab Michael Hanselmann
          fields = line.split()
658 7e3dbb94 Iustin Pop
          if fields and not fields[0].startswith('#') and ip == fields[0]:
659 9440aeab Michael Hanselmann
            continue
660 9440aeab Michael Hanselmann
          out.write(line)
661 9440aeab Michael Hanselmann
662 7e3dbb94 Iustin Pop
        out.write("%s\t%s" % (ip, hostname))
663 9440aeab Michael Hanselmann
        if aliases:
664 9440aeab Michael Hanselmann
          out.write(" %s" % ' '.join(aliases))
665 9440aeab Michael Hanselmann
        out.write('\n')
666 9440aeab Michael Hanselmann
667 9440aeab Michael Hanselmann
        out.flush()
668 2e3e75b7 Michael Hanselmann
        os.fsync(out)
669 9440aeab Michael Hanselmann
        os.rename(tmpname, file_name)
670 9440aeab Michael Hanselmann
      finally:
671 9440aeab Michael Hanselmann
        f.close()
672 9440aeab Michael Hanselmann
    finally:
673 9440aeab Michael Hanselmann
      out.close()
674 9440aeab Michael Hanselmann
  except:
675 9440aeab Michael Hanselmann
    RemoveFile(tmpname)
676 9440aeab Michael Hanselmann
    raise
677 899d2a81 Michael Hanselmann
678 899d2a81 Michael Hanselmann
679 d9c02ca6 Michael Hanselmann
def AddHostToEtcHosts(hostname):
680 d9c02ca6 Michael Hanselmann
  """Wrapper around SetEtcHostsEntry.
681 d9c02ca6 Michael Hanselmann

682 d9c02ca6 Michael Hanselmann
  """
683 d9c02ca6 Michael Hanselmann
  hi = HostInfo(name=hostname)
684 d9c02ca6 Michael Hanselmann
  SetEtcHostsEntry(constants.ETC_HOSTS, hi.ip, hi.name, [hi.ShortName()])
685 d9c02ca6 Michael Hanselmann
686 d9c02ca6 Michael Hanselmann
687 899d2a81 Michael Hanselmann
def RemoveEtcHostsEntry(file_name, hostname):
688 3e1cdf9f Michael Hanselmann
  """Removes a hostname from /etc/hosts.
689 899d2a81 Michael Hanselmann

690 9440aeab Michael Hanselmann
  IP addresses without names are removed from the file.
691 899d2a81 Michael Hanselmann
  """
692 899d2a81 Michael Hanselmann
  fd, tmpname = tempfile.mkstemp(dir=os.path.dirname(file_name))
693 899d2a81 Michael Hanselmann
  try:
694 899d2a81 Michael Hanselmann
    out = os.fdopen(fd, 'w')
695 899d2a81 Michael Hanselmann
    try:
696 899d2a81 Michael Hanselmann
      f = open(file_name, 'r')
697 899d2a81 Michael Hanselmann
      try:
698 899d2a81 Michael Hanselmann
        for line in f:
699 899d2a81 Michael Hanselmann
          fields = line.split()
700 899d2a81 Michael Hanselmann
          if len(fields) > 1 and not fields[0].startswith('#'):
701 899d2a81 Michael Hanselmann
            names = fields[1:]
702 899d2a81 Michael Hanselmann
            if hostname in names:
703 899d2a81 Michael Hanselmann
              while hostname in names:
704 899d2a81 Michael Hanselmann
                names.remove(hostname)
705 899d2a81 Michael Hanselmann
              if names:
706 9440aeab Michael Hanselmann
                out.write("%s %s\n" % (fields[0], ' '.join(names)))
707 899d2a81 Michael Hanselmann
              continue
708 899d2a81 Michael Hanselmann
709 899d2a81 Michael Hanselmann
          out.write(line)
710 59f82e3f Michael Hanselmann
711 59f82e3f Michael Hanselmann
        out.flush()
712 2e3e75b7 Michael Hanselmann
        os.fsync(out)
713 59f82e3f Michael Hanselmann
        os.rename(tmpname, file_name)
714 59f82e3f Michael Hanselmann
      finally:
715 59f82e3f Michael Hanselmann
        f.close()
716 a8083063 Iustin Pop
    finally:
717 59f82e3f Michael Hanselmann
      out.close()
718 59f82e3f Michael Hanselmann
  except:
719 59f82e3f Michael Hanselmann
    RemoveFile(tmpname)
720 59f82e3f Michael Hanselmann
    raise
721 a8083063 Iustin Pop
722 a8083063 Iustin Pop
723 d9c02ca6 Michael Hanselmann
def RemoveHostFromEtcHosts(hostname):
724 d9c02ca6 Michael Hanselmann
  """Wrapper around RemoveEtcHostsEntry.
725 d9c02ca6 Michael Hanselmann

726 d9c02ca6 Michael Hanselmann
  """
727 d9c02ca6 Michael Hanselmann
  hi = HostInfo(name=hostname)
728 d9c02ca6 Michael Hanselmann
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.name)
729 d9c02ca6 Michael Hanselmann
  RemoveEtcHostsEntry(constants.ETC_HOSTS, hi.ShortName())
730 d9c02ca6 Michael Hanselmann
731 d9c02ca6 Michael Hanselmann
732 a8083063 Iustin Pop
def CreateBackup(file_name):
733 a8083063 Iustin Pop
  """Creates a backup of a file.
734 a8083063 Iustin Pop

735 a8083063 Iustin Pop
  Returns: the path to the newly created backup file.
736 a8083063 Iustin Pop

737 a8083063 Iustin Pop
  """
738 a8083063 Iustin Pop
  if not os.path.isfile(file_name):
739 3ecf6786 Iustin Pop
    raise errors.ProgrammerError("Can't make a backup of a non-file '%s'" %
740 3ecf6786 Iustin Pop
                                file_name)
741 a8083063 Iustin Pop
742 081b1e69 Michael Hanselmann
  prefix = '%s.backup-%d.' % (os.path.basename(file_name), int(time.time()))
743 65fe4693 Iustin Pop
  dir_name = os.path.dirname(file_name)
744 081b1e69 Michael Hanselmann
745 081b1e69 Michael Hanselmann
  fsrc = open(file_name, 'rb')
746 081b1e69 Michael Hanselmann
  try:
747 65fe4693 Iustin Pop
    (fd, backup_name) = tempfile.mkstemp(prefix=prefix, dir=dir_name)
748 081b1e69 Michael Hanselmann
    fdst = os.fdopen(fd, 'wb')
749 081b1e69 Michael Hanselmann
    try:
750 081b1e69 Michael Hanselmann
      shutil.copyfileobj(fsrc, fdst)
751 081b1e69 Michael Hanselmann
    finally:
752 081b1e69 Michael Hanselmann
      fdst.close()
753 081b1e69 Michael Hanselmann
  finally:
754 081b1e69 Michael Hanselmann
    fsrc.close()
755 081b1e69 Michael Hanselmann
756 a8083063 Iustin Pop
  return backup_name
757 a8083063 Iustin Pop
758 a8083063 Iustin Pop
759 a8083063 Iustin Pop
def ShellQuote(value):
760 a8083063 Iustin Pop
  """Quotes shell argument according to POSIX.
761 3ecf6786 Iustin Pop

762 a8083063 Iustin Pop
  """
763 a8083063 Iustin Pop
  if _re_shell_unquoted.match(value):
764 a8083063 Iustin Pop
    return value
765 a8083063 Iustin Pop
  else:
766 a8083063 Iustin Pop
    return "'%s'" % value.replace("'", "'\\''")
767 a8083063 Iustin Pop
768 a8083063 Iustin Pop
769 a8083063 Iustin Pop
def ShellQuoteArgs(args):
770 a8083063 Iustin Pop
  """Quotes all given shell arguments and concatenates using spaces.
771 a8083063 Iustin Pop

772 a8083063 Iustin Pop
  """
773 a8083063 Iustin Pop
  return ' '.join([ShellQuote(i) for i in args])
774 88d14415 Michael Hanselmann
775 88d14415 Michael Hanselmann
776 b15d625f Iustin Pop
def TcpPing(target, port, timeout=10, live_port_needed=False, source=None):
777 2c30e9d7 Alexander Schreiber
  """Simple ping implementation using TCP connect(2).
778 2c30e9d7 Alexander Schreiber

779 b15d625f Iustin Pop
  Try to do a TCP connect(2) from an optional source IP to the
780 b15d625f Iustin Pop
  specified target IP and the specified target port. If the optional
781 b15d625f Iustin Pop
  parameter live_port_needed is set to true, requires the remote end
782 b15d625f Iustin Pop
  to accept the connection. The timeout is specified in seconds and
783 b15d625f Iustin Pop
  defaults to 10 seconds. If the source optional argument is not
784 b15d625f Iustin Pop
  passed, the source address selection is left to the kernel,
785 b15d625f Iustin Pop
  otherwise we try to connect using the passed address (failures to
786 b15d625f Iustin Pop
  bind other than EADDRNOTAVAIL will be ignored).
787 2c30e9d7 Alexander Schreiber

788 2c30e9d7 Alexander Schreiber
  """
789 2c30e9d7 Alexander Schreiber
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
790 2c30e9d7 Alexander Schreiber
791 2c30e9d7 Alexander Schreiber
  sucess = False
792 2c30e9d7 Alexander Schreiber
793 b15d625f Iustin Pop
  if source is not None:
794 b15d625f Iustin Pop
    try:
795 b15d625f Iustin Pop
      sock.bind((source, 0))
796 b15d625f Iustin Pop
    except socket.error, (errcode, errstring):
797 b15d625f Iustin Pop
      if errcode == errno.EADDRNOTAVAIL:
798 b15d625f Iustin Pop
        success = False
799 2c30e9d7 Alexander Schreiber
800 2c30e9d7 Alexander Schreiber
  sock.settimeout(timeout)
801 2c30e9d7 Alexander Schreiber
802 2c30e9d7 Alexander Schreiber
  try:
803 2c30e9d7 Alexander Schreiber
    sock.connect((target, port))
804 2c30e9d7 Alexander Schreiber
    sock.close()
805 2c30e9d7 Alexander Schreiber
    success = True
806 2c30e9d7 Alexander Schreiber
  except socket.timeout:
807 2c30e9d7 Alexander Schreiber
    success = False
808 2c30e9d7 Alexander Schreiber
  except socket.error, (errcode, errstring):
809 4ca1b175 Alexander Schreiber
    success = (not live_port_needed) and (errcode == errno.ECONNREFUSED)
810 2c30e9d7 Alexander Schreiber
811 2c30e9d7 Alexander Schreiber
  return success
812 eedbda4b Michael Hanselmann
813 eedbda4b Michael Hanselmann
814 eedbda4b Michael Hanselmann
def ListVisibleFiles(path):
815 eedbda4b Michael Hanselmann
  """Returns a list of all visible files in a directory.
816 eedbda4b Michael Hanselmann

817 eedbda4b Michael Hanselmann
  """
818 f3299a07 Michael Hanselmann
  files = [i for i in os.listdir(path) if not i.startswith(".")]
819 f3299a07 Michael Hanselmann
  files.sort()
820 f3299a07 Michael Hanselmann
  return files
821 2f8b60b3 Iustin Pop
822 2f8b60b3 Iustin Pop
823 257f4c0a Iustin Pop
def GetHomeDir(user, default=None):
824 257f4c0a Iustin Pop
  """Try to get the homedir of the given user.
825 257f4c0a Iustin Pop

826 257f4c0a Iustin Pop
  The user can be passed either as a string (denoting the name) or as
827 257f4c0a Iustin Pop
  an integer (denoting the user id). If the user is not found, the
828 257f4c0a Iustin Pop
  'default' argument is returned, which defaults to None.
829 2f8b60b3 Iustin Pop

830 2f8b60b3 Iustin Pop
  """
831 2f8b60b3 Iustin Pop
  try:
832 257f4c0a Iustin Pop
    if isinstance(user, basestring):
833 257f4c0a Iustin Pop
      result = pwd.getpwnam(user)
834 257f4c0a Iustin Pop
    elif isinstance(user, (int, long)):
835 257f4c0a Iustin Pop
      result = pwd.getpwuid(user)
836 257f4c0a Iustin Pop
    else:
837 257f4c0a Iustin Pop
      raise errors.ProgrammerError("Invalid type passed to GetHomeDir (%s)" %
838 257f4c0a Iustin Pop
                                   type(user))
839 2f8b60b3 Iustin Pop
  except KeyError:
840 2f8b60b3 Iustin Pop
    return default
841 2f8b60b3 Iustin Pop
  return result.pw_dir
842 59072e7e Michael Hanselmann
843 59072e7e Michael Hanselmann
844 24818e8f Michael Hanselmann
def NewUUID():
845 59072e7e Michael Hanselmann
  """Returns a random UUID.
846 59072e7e Michael Hanselmann

847 59072e7e Michael Hanselmann
  """
848 59072e7e Michael Hanselmann
  f = open("/proc/sys/kernel/random/uuid", "r")
849 59072e7e Michael Hanselmann
  try:
850 59072e7e Michael Hanselmann
    return f.read(128).rstrip("\n")
851 59072e7e Michael Hanselmann
  finally:
852 59072e7e Michael Hanselmann
    f.close()
853 087b34fe Iustin Pop
854 087b34fe Iustin Pop
855 087b34fe Iustin Pop
def WriteFile(file_name, fn=None, data=None,
856 087b34fe Iustin Pop
              mode=None, uid=-1, gid=-1,
857 71714516 Michael Hanselmann
              atime=None, mtime=None, close=True,
858 04a8d789 Michael Hanselmann
              dry_run=False, backup=False,
859 71714516 Michael Hanselmann
              prewrite=None, postwrite=None):
860 087b34fe Iustin Pop
  """(Over)write a file atomically.
861 087b34fe Iustin Pop

862 087b34fe Iustin Pop
  The file_name and either fn (a function taking one argument, the
863 087b34fe Iustin Pop
  file descriptor, and which should write the data to it) or data (the
864 087b34fe Iustin Pop
  contents of the file) must be passed. The other arguments are
865 087b34fe Iustin Pop
  optional and allow setting the file mode, owner and group, and the
866 087b34fe Iustin Pop
  mtime/atime of the file.
867 087b34fe Iustin Pop

868 087b34fe Iustin Pop
  If the function doesn't raise an exception, it has succeeded and the
869 087b34fe Iustin Pop
  target file has the new contents. If the file has raised an
870 087b34fe Iustin Pop
  exception, an existing target file should be unmodified and the
871 087b34fe Iustin Pop
  temporary file should be removed.
872 087b34fe Iustin Pop

873 71714516 Michael Hanselmann
  Args:
874 71714516 Michael Hanselmann
    file_name: New filename
875 71714516 Michael Hanselmann
    fn: Content writing function, called with file descriptor as parameter
876 71714516 Michael Hanselmann
    data: Content as string
877 71714516 Michael Hanselmann
    mode: File mode
878 71714516 Michael Hanselmann
    uid: Owner
879 71714516 Michael Hanselmann
    gid: Group
880 71714516 Michael Hanselmann
    atime: Access time
881 71714516 Michael Hanselmann
    mtime: Modification time
882 71714516 Michael Hanselmann
    close: Whether to close file after writing it
883 71714516 Michael Hanselmann
    prewrite: Function object called before writing content
884 71714516 Michael Hanselmann
    postwrite: Function object called after writing content
885 71714516 Michael Hanselmann

886 71714516 Michael Hanselmann
  Returns:
887 71714516 Michael Hanselmann
    None if "close" parameter evaluates to True, otherwise file descriptor.
888 71714516 Michael Hanselmann

889 087b34fe Iustin Pop
  """
890 04a8d789 Michael Hanselmann
  if not os.path.isabs(file_name):
891 087b34fe Iustin Pop
    raise errors.ProgrammerError("Path passed to WriteFile is not"
892 087b34fe Iustin Pop
                                 " absolute: '%s'" % file_name)
893 087b34fe Iustin Pop
894 087b34fe Iustin Pop
  if [fn, data].count(None) != 1:
895 087b34fe Iustin Pop
    raise errors.ProgrammerError("fn or data required")
896 087b34fe Iustin Pop
897 087b34fe Iustin Pop
  if [atime, mtime].count(None) == 1:
898 087b34fe Iustin Pop
    raise errors.ProgrammerError("Both atime and mtime must be either"
899 087b34fe Iustin Pop
                                 " set or None")
900 087b34fe Iustin Pop
901 70f4497c Michael Hanselmann
  if backup and not dry_run and os.path.isfile(file_name):
902 70f4497c Michael Hanselmann
    CreateBackup(file_name)
903 087b34fe Iustin Pop
904 087b34fe Iustin Pop
  dir_name, base_name = os.path.split(file_name)
905 087b34fe Iustin Pop
  fd, new_name = tempfile.mkstemp('.new', base_name, dir_name)
906 087b34fe Iustin Pop
  # here we need to make sure we remove the temp file, if any error
907 087b34fe Iustin Pop
  # leaves it in place
908 087b34fe Iustin Pop
  try:
909 087b34fe Iustin Pop
    if uid != -1 or gid != -1:
910 087b34fe Iustin Pop
      os.chown(new_name, uid, gid)
911 087b34fe Iustin Pop
    if mode:
912 087b34fe Iustin Pop
      os.chmod(new_name, mode)
913 71714516 Michael Hanselmann
    if callable(prewrite):
914 71714516 Michael Hanselmann
      prewrite(fd)
915 087b34fe Iustin Pop
    if data is not None:
916 087b34fe Iustin Pop
      os.write(fd, data)
917 087b34fe Iustin Pop
    else:
918 087b34fe Iustin Pop
      fn(fd)
919 71714516 Michael Hanselmann
    if callable(postwrite):
920 71714516 Michael Hanselmann
      postwrite(fd)
921 087b34fe Iustin Pop
    os.fsync(fd)
922 087b34fe Iustin Pop
    if atime is not None and mtime is not None:
923 087b34fe Iustin Pop
      os.utime(new_name, (atime, mtime))
924 70f4497c Michael Hanselmann
    if not dry_run:
925 70f4497c Michael Hanselmann
      os.rename(new_name, file_name)
926 087b34fe Iustin Pop
  finally:
927 71714516 Michael Hanselmann
    if close:
928 71714516 Michael Hanselmann
      os.close(fd)
929 71714516 Michael Hanselmann
      result = None
930 71714516 Michael Hanselmann
    else:
931 71714516 Michael Hanselmann
      result = fd
932 087b34fe Iustin Pop
    RemoveFile(new_name)
933 78feb6fb Guido Trotter
934 71714516 Michael Hanselmann
  return result
935 71714516 Michael Hanselmann
936 78feb6fb Guido Trotter
937 7b4126b7 Iustin Pop
def FirstFree(seq, base=0):
938 7b4126b7 Iustin Pop
  """Returns the first non-existing integer from seq.
939 7b4126b7 Iustin Pop

940 7b4126b7 Iustin Pop
  The seq argument should be a sorted list of positive integers. The
941 7b4126b7 Iustin Pop
  first time the index of an element is smaller than the element
942 7b4126b7 Iustin Pop
  value, the index will be returned.
943 7b4126b7 Iustin Pop

944 7b4126b7 Iustin Pop
  The base argument is used to start at a different offset,
945 7b4126b7 Iustin Pop
  i.e. [3, 4, 6] with offset=3 will return 5.
946 7b4126b7 Iustin Pop

947 7b4126b7 Iustin Pop
  Example: [0, 1, 3] will return 2.
948 7b4126b7 Iustin Pop

949 7b4126b7 Iustin Pop
  """
950 7b4126b7 Iustin Pop
  for idx, elem in enumerate(seq):
951 7b4126b7 Iustin Pop
    assert elem >= base, "Passed element is higher than base offset"
952 7b4126b7 Iustin Pop
    if elem > idx + base:
953 7b4126b7 Iustin Pop
      # idx is not used
954 7b4126b7 Iustin Pop
      return idx + base
955 7b4126b7 Iustin Pop
  return None
956 7b4126b7 Iustin Pop
957 7b4126b7 Iustin Pop
958 78feb6fb Guido Trotter
def all(seq, pred=bool):
959 78feb6fb Guido Trotter
  "Returns True if pred(x) is True for every element in the iterable"
960 78feb6fb Guido Trotter
  for elem in itertools.ifilterfalse(pred, seq):
961 78feb6fb Guido Trotter
    return False
962 78feb6fb Guido Trotter
  return True
963 78feb6fb Guido Trotter
964 78feb6fb Guido Trotter
965 78feb6fb Guido Trotter
def any(seq, pred=bool):
966 78feb6fb Guido Trotter
  "Returns True if pred(x) is True for at least one element in the iterable"
967 78feb6fb Guido Trotter
  for elem in itertools.ifilter(pred, seq):
968 78feb6fb Guido Trotter
    return True
969 78feb6fb Guido Trotter
  return False
970 f7414041 Michael Hanselmann
971 f7414041 Michael Hanselmann
972 f7414041 Michael Hanselmann
def UniqueSequence(seq):
973 f7414041 Michael Hanselmann
  """Returns a list with unique elements.
974 f7414041 Michael Hanselmann

975 f7414041 Michael Hanselmann
  Element order is preserved.
976 f7414041 Michael Hanselmann
  """
977 f7414041 Michael Hanselmann
  seen = set()
978 f7414041 Michael Hanselmann
  return [i for i in seq if i not in seen and not seen.add(i)]
979 1862d460 Alexander Schreiber
980 1862d460 Alexander Schreiber
981 1862d460 Alexander Schreiber
def IsValidMac(mac):
982 1862d460 Alexander Schreiber
  """Predicate to check if a MAC address is valid.
983 1862d460 Alexander Schreiber

984 1862d460 Alexander Schreiber
  Checks wether the supplied MAC address is formally correct, only
985 1862d460 Alexander Schreiber
  accepts colon separated format.
986 1862d460 Alexander Schreiber
  """
987 1862d460 Alexander Schreiber
  mac_check = re.compile("^([0-9a-f]{2}(:|$)){6}$")
988 1862d460 Alexander Schreiber
  return mac_check.match(mac) is not None
989 06009e27 Iustin Pop
990 06009e27 Iustin Pop
991 06009e27 Iustin Pop
def TestDelay(duration):
992 06009e27 Iustin Pop
  """Sleep for a fixed amount of time.
993 06009e27 Iustin Pop

994 06009e27 Iustin Pop
  """
995 06009e27 Iustin Pop
  if duration < 0:
996 06009e27 Iustin Pop
    return False
997 06009e27 Iustin Pop
  time.sleep(duration)
998 06009e27 Iustin Pop
  return True
999 8f765069 Iustin Pop
1000 8f765069 Iustin Pop
1001 8ff612c2 Iustin Pop
def Daemonize(logfile, noclose_fds=None):
1002 8f765069 Iustin Pop
  """Daemonize the current process.
1003 8f765069 Iustin Pop

1004 8f765069 Iustin Pop
  This detaches the current process from the controlling terminal and
1005 8f765069 Iustin Pop
  runs it in the background as a daemon.
1006 8f765069 Iustin Pop

1007 8f765069 Iustin Pop
  """
1008 8f765069 Iustin Pop
  UMASK = 077
1009 8f765069 Iustin Pop
  WORKDIR = "/"
1010 8f765069 Iustin Pop
  # Default maximum for the number of available file descriptors.
1011 8f765069 Iustin Pop
  if 'SC_OPEN_MAX' in os.sysconf_names:
1012 8f765069 Iustin Pop
    try:
1013 8f765069 Iustin Pop
      MAXFD = os.sysconf('SC_OPEN_MAX')
1014 8f765069 Iustin Pop
      if MAXFD < 0:
1015 8f765069 Iustin Pop
        MAXFD = 1024
1016 8f765069 Iustin Pop
    except OSError:
1017 8f765069 Iustin Pop
      MAXFD = 1024
1018 8f765069 Iustin Pop
  else:
1019 8f765069 Iustin Pop
    MAXFD = 1024
1020 8f765069 Iustin Pop
1021 8f765069 Iustin Pop
  # this might fail
1022 8f765069 Iustin Pop
  pid = os.fork()
1023 8f765069 Iustin Pop
  if (pid == 0):  # The first child.
1024 8f765069 Iustin Pop
    os.setsid()
1025 8f765069 Iustin Pop
    # this might fail
1026 8f765069 Iustin Pop
    pid = os.fork() # Fork a second child.
1027 8f765069 Iustin Pop
    if (pid == 0):  # The second child.
1028 8f765069 Iustin Pop
      os.chdir(WORKDIR)
1029 8f765069 Iustin Pop
      os.umask(UMASK)
1030 8f765069 Iustin Pop
    else:
1031 8f765069 Iustin Pop
      # exit() or _exit()?  See below.
1032 8f765069 Iustin Pop
      os._exit(0) # Exit parent (the first child) of the second child.
1033 8f765069 Iustin Pop
  else:
1034 8f765069 Iustin Pop
    os._exit(0) # Exit parent of the first child.
1035 8f765069 Iustin Pop
  maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
1036 8f765069 Iustin Pop
  if (maxfd == resource.RLIM_INFINITY):
1037 8f765069 Iustin Pop
    maxfd = MAXFD
1038 8f765069 Iustin Pop
1039 8f765069 Iustin Pop
  # Iterate through and close all file descriptors.
1040 8f765069 Iustin Pop
  for fd in range(0, maxfd):
1041 8ff612c2 Iustin Pop
    if noclose_fds and fd in noclose_fds:
1042 8ff612c2 Iustin Pop
      continue
1043 8f765069 Iustin Pop
    try:
1044 8f765069 Iustin Pop
      os.close(fd)
1045 8f765069 Iustin Pop
    except OSError: # ERROR, fd wasn't open to begin with (ignored)
1046 8f765069 Iustin Pop
      pass
1047 8f765069 Iustin Pop
  os.open(logfile, os.O_RDWR|os.O_CREAT|os.O_APPEND, 0600)
1048 8f765069 Iustin Pop
  # Duplicate standard input to standard output and standard error.
1049 8f765069 Iustin Pop
  os.dup2(0, 1)     # standard output (1)
1050 8f765069 Iustin Pop
  os.dup2(0, 2)     # standard error (2)
1051 8f765069 Iustin Pop
  return 0
1052 57c177af Iustin Pop
1053 57c177af Iustin Pop
1054 b330ac0b Guido Trotter
def _DaemonPidFileName(name):
1055 b330ac0b Guido Trotter
  """Compute a ganeti pid file absolute path, given the daemon name.
1056 b330ac0b Guido Trotter

1057 b330ac0b Guido Trotter
  """
1058 b330ac0b Guido Trotter
  return os.path.join(constants.RUN_GANETI_DIR, "%s.pid" % name)
1059 b330ac0b Guido Trotter
1060 b330ac0b Guido Trotter
1061 b330ac0b Guido Trotter
def WritePidFile(name):
1062 b330ac0b Guido Trotter
  """Write the current process pidfile.
1063 b330ac0b Guido Trotter

1064 b330ac0b Guido Trotter
  The file will be written to constants.RUN_GANETI_DIR/name.pid
1065 b330ac0b Guido Trotter

1066 b330ac0b Guido Trotter
  """
1067 b330ac0b Guido Trotter
  pid = os.getpid()
1068 b330ac0b Guido Trotter
  pidfilename = _DaemonPidFileName(name)
1069 b330ac0b Guido Trotter
  if IsPidFileAlive(pidfilename):
1070 533bb4b1 Michael Hanselmann
    raise errors.GenericError("%s contains a live process" % pidfilename)
1071 b330ac0b Guido Trotter
1072 b330ac0b Guido Trotter
  WriteFile(pidfilename, data="%d\n" % pid)
1073 b330ac0b Guido Trotter
1074 b330ac0b Guido Trotter
1075 b330ac0b Guido Trotter
def RemovePidFile(name):
1076 b330ac0b Guido Trotter
  """Remove the current process pidfile.
1077 b330ac0b Guido Trotter

1078 b330ac0b Guido Trotter
  Any errors are ignored.
1079 b330ac0b Guido Trotter

1080 b330ac0b Guido Trotter
  """
1081 b330ac0b Guido Trotter
  pid = os.getpid()
1082 b330ac0b Guido Trotter
  pidfilename = _DaemonPidFileName(name)
1083 b330ac0b Guido Trotter
  # TODO: we could check here that the file contains our pid
1084 b330ac0b Guido Trotter
  try:
1085 b330ac0b Guido Trotter
    RemoveFile(pidfilename)
1086 b330ac0b Guido Trotter
  except:
1087 b330ac0b Guido Trotter
    pass
1088 b330ac0b Guido Trotter
1089 b330ac0b Guido Trotter
1090 57c177af Iustin Pop
def FindFile(name, search_path, test=os.path.exists):
1091 57c177af Iustin Pop
  """Look for a filesystem object in a given path.
1092 57c177af Iustin Pop

1093 57c177af Iustin Pop
  This is an abstract method to search for filesystem object (files,
1094 57c177af Iustin Pop
  dirs) under a given search path.
1095 57c177af Iustin Pop

1096 57c177af Iustin Pop
  Args:
1097 57c177af Iustin Pop
    - name: the name to look for
1098 57c177af Iustin Pop
    - search_path: list of directory names
1099 57c177af Iustin Pop
    - test: the test which the full path must satisfy
1100 57c177af Iustin Pop
      (defaults to os.path.exists)
1101 57c177af Iustin Pop

1102 57c177af Iustin Pop
  Returns:
1103 57c177af Iustin Pop
    - full path to the item if found
1104 57c177af Iustin Pop
    - None otherwise
1105 57c177af Iustin Pop

1106 57c177af Iustin Pop
  """
1107 57c177af Iustin Pop
  for dir_name in search_path:
1108 57c177af Iustin Pop
    item_name = os.path.sep.join([dir_name, name])
1109 57c177af Iustin Pop
    if test(item_name):
1110 57c177af Iustin Pop
      return item_name
1111 57c177af Iustin Pop
  return None
1112 8d1a2a64 Michael Hanselmann
1113 8d1a2a64 Michael Hanselmann
1114 8d1a2a64 Michael Hanselmann
def CheckVolumeGroupSize(vglist, vgname, minsize):
1115 8d1a2a64 Michael Hanselmann
  """Checks if the volume group list is valid.
1116 8d1a2a64 Michael Hanselmann

1117 8d1a2a64 Michael Hanselmann
  A non-None return value means there's an error, and the return value
1118 8d1a2a64 Michael Hanselmann
  is the error message.
1119 8d1a2a64 Michael Hanselmann

1120 8d1a2a64 Michael Hanselmann
  """
1121 8d1a2a64 Michael Hanselmann
  vgsize = vglist.get(vgname, None)
1122 8d1a2a64 Michael Hanselmann
  if vgsize is None:
1123 8d1a2a64 Michael Hanselmann
    return "volume group '%s' missing" % vgname
1124 8d1a2a64 Michael Hanselmann
  elif vgsize < minsize:
1125 8d1a2a64 Michael Hanselmann
    return ("volume group '%s' too small (%s MiB required, %d MiB found)" %
1126 8d1a2a64 Michael Hanselmann
            (vgname, minsize, vgsize))
1127 8d1a2a64 Michael Hanselmann
  return None
1128 7996a135 Iustin Pop
1129 7996a135 Iustin Pop
1130 7996a135 Iustin Pop
def LockedMethod(fn):
1131 7996a135 Iustin Pop
  """Synchronized object access decorator.
1132 7996a135 Iustin Pop

1133 7996a135 Iustin Pop
  This decorator is intended to protect access to an object using the
1134 7996a135 Iustin Pop
  object's own lock which is hardcoded to '_lock'.
1135 7996a135 Iustin Pop

1136 7996a135 Iustin Pop
  """
1137 7996a135 Iustin Pop
  def wrapper(self, *args, **kwargs):
1138 7996a135 Iustin Pop
    assert hasattr(self, '_lock')
1139 7996a135 Iustin Pop
    lock = self._lock
1140 7996a135 Iustin Pop
    lock.acquire()
1141 7996a135 Iustin Pop
    try:
1142 7996a135 Iustin Pop
      result = fn(self, *args, **kwargs)
1143 7996a135 Iustin Pop
    finally:
1144 7996a135 Iustin Pop
      lock.release()
1145 7996a135 Iustin Pop
    return result
1146 7996a135 Iustin Pop
  return wrapper
1147 eb0f0ce0 Michael Hanselmann
1148 eb0f0ce0 Michael Hanselmann
1149 eb0f0ce0 Michael Hanselmann
def LockFile(fd):
1150 eb0f0ce0 Michael Hanselmann
  """Locks a file using POSIX locks.
1151 eb0f0ce0 Michael Hanselmann

1152 eb0f0ce0 Michael Hanselmann
  """
1153 eb0f0ce0 Michael Hanselmann
  try:
1154 eb0f0ce0 Michael Hanselmann
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
1155 eb0f0ce0 Michael Hanselmann
  except IOError, err:
1156 eb0f0ce0 Michael Hanselmann
    if err.errno == errno.EAGAIN:
1157 eb0f0ce0 Michael Hanselmann
      raise errors.LockError("File already locked")
1158 eb0f0ce0 Michael Hanselmann
    raise
1159 de499029 Michael Hanselmann
1160 de499029 Michael Hanselmann
1161 de499029 Michael Hanselmann
class SignalHandler(object):
1162 de499029 Michael Hanselmann
  """Generic signal handler class.
1163 de499029 Michael Hanselmann

1164 de499029 Michael Hanselmann
  It automatically restores the original handler when deconstructed or when
1165 de499029 Michael Hanselmann
  Reset() is called. You can either pass your own handler function in or query
1166 de499029 Michael Hanselmann
  the "called" attribute to detect whether the signal was sent.
1167 de499029 Michael Hanselmann

1168 de499029 Michael Hanselmann
  """
1169 de499029 Michael Hanselmann
  def __init__(self, signum):
1170 de499029 Michael Hanselmann
    """Constructs a new SignalHandler instance.
1171 de499029 Michael Hanselmann

1172 de499029 Michael Hanselmann
    @param signum: Single signal number or set of signal numbers
1173 de499029 Michael Hanselmann

1174 de499029 Michael Hanselmann
    """
1175 de499029 Michael Hanselmann
    if isinstance(signum, (int, long)):
1176 de499029 Michael Hanselmann
      self.signum = set([signum])
1177 de499029 Michael Hanselmann
    else:
1178 de499029 Michael Hanselmann
      self.signum = set(signum)
1179 de499029 Michael Hanselmann
1180 de499029 Michael Hanselmann
    self.called = False
1181 de499029 Michael Hanselmann
1182 de499029 Michael Hanselmann
    self._previous = {}
1183 de499029 Michael Hanselmann
    try:
1184 de499029 Michael Hanselmann
      for signum in self.signum:
1185 de499029 Michael Hanselmann
        # Setup handler
1186 de499029 Michael Hanselmann
        prev_handler = signal.signal(signum, self._HandleSignal)
1187 de499029 Michael Hanselmann
        try:
1188 de499029 Michael Hanselmann
          self._previous[signum] = prev_handler
1189 de499029 Michael Hanselmann
        except:
1190 de499029 Michael Hanselmann
          # Restore previous handler
1191 de499029 Michael Hanselmann
          signal.signal(signum, prev_handler)
1192 de499029 Michael Hanselmann
          raise
1193 de499029 Michael Hanselmann
    except:
1194 de499029 Michael Hanselmann
      # Reset all handlers
1195 de499029 Michael Hanselmann
      self.Reset()
1196 de499029 Michael Hanselmann
      # Here we have a race condition: a handler may have already been called,
1197 de499029 Michael Hanselmann
      # but there's not much we can do about it at this point.
1198 de499029 Michael Hanselmann
      raise
1199 de499029 Michael Hanselmann
1200 de499029 Michael Hanselmann
  def __del__(self):
1201 de499029 Michael Hanselmann
    self.Reset()
1202 de499029 Michael Hanselmann
1203 de499029 Michael Hanselmann
  def Reset(self):
1204 de499029 Michael Hanselmann
    """Restore previous handler.
1205 de499029 Michael Hanselmann

1206 de499029 Michael Hanselmann
    """
1207 de499029 Michael Hanselmann
    for signum, prev_handler in self._previous.items():
1208 de499029 Michael Hanselmann
      signal.signal(signum, prev_handler)
1209 de499029 Michael Hanselmann
      # If successful, remove from dict
1210 de499029 Michael Hanselmann
      del self._previous[signum]
1211 de499029 Michael Hanselmann
1212 de499029 Michael Hanselmann
  def Clear(self):
1213 de499029 Michael Hanselmann
    """Unsets "called" flag.
1214 de499029 Michael Hanselmann

1215 de499029 Michael Hanselmann
    This function can be used in case a signal may arrive several times.
1216 de499029 Michael Hanselmann

1217 de499029 Michael Hanselmann
    """
1218 de499029 Michael Hanselmann
    self.called = False
1219 de499029 Michael Hanselmann
1220 de499029 Michael Hanselmann
  def _HandleSignal(self, signum, frame):
1221 de499029 Michael Hanselmann
    """Actual signal handling function.
1222 de499029 Michael Hanselmann

1223 de499029 Michael Hanselmann
    """
1224 de499029 Michael Hanselmann
    # This is not nice and not absolutely atomic, but it appears to be the only
1225 de499029 Michael Hanselmann
    # solution in Python -- there are no atomic types.
1226 de499029 Michael Hanselmann
    self.called = True