~lutostag/ubuntu/trusty/maas/1.5.2+packagefix

« back to all changes in this revision

Viewing changes to src/provisioningserver/utils/__init__.py

  • Committer: Package Import Robot
  • Author(s): Andres Rodriguez
  • Date: 2014-03-28 10:43:53 UTC
  • mto: This revision was merged to the branch mainline in revision 57.
  • Revision ID: package-import@ubuntu.com-20140328104353-ekpolg0pm5xnvq2s
Tags: upstream-1.5+bzr2204
ImportĀ upstreamĀ versionĀ 1.5+bzr2204

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright 2012-2014 Canonical Ltd.  This software is licensed under the
 
2
# GNU Affero General Public License version 3 (see the file LICENSE).
 
3
 
 
4
"""Utilities for the provisioning server."""
 
5
 
 
6
from __future__ import (
 
7
    absolute_import,
 
8
    print_function,
 
9
    unicode_literals,
 
10
    )
 
11
 
 
12
str = None
 
13
 
 
14
__metaclass__ = type
 
15
__all__ = [
 
16
    "ActionScript",
 
17
    "atomic_write",
 
18
    "deferred",
 
19
    "filter_dict",
 
20
    "find_ip_via_arp",
 
21
    "import_settings",
 
22
    "incremental_write",
 
23
    "locate_config",
 
24
    "MainScript",
 
25
    "ensure_dir",
 
26
    "parse_key_value_file",
 
27
    "read_text_file",
 
28
    "ShellTemplate",
 
29
    "sudo_write_file",
 
30
    "write_custom_config_section",
 
31
    "write_text_file",
 
32
    ]
 
33
 
 
34
from argparse import ArgumentParser
 
35
import codecs
 
36
from contextlib import contextmanager
 
37
import errno
 
38
from functools import wraps
 
39
import logging
 
40
import os
 
41
from os import fdopen
 
42
from os.path import isdir
 
43
from pipes import quote
 
44
from shutil import rmtree
 
45
import signal
 
46
import string
 
47
import subprocess
 
48
from subprocess import (
 
49
    CalledProcessError,
 
50
    PIPE,
 
51
    Popen,
 
52
    )
 
53
import sys
 
54
import tempfile
 
55
from time import time
 
56
 
 
57
from crochet import run_in_reactor
 
58
from lockfile import FileLock
 
59
from lxml import etree
 
60
import netifaces
 
61
import tempita
 
62
from twisted.internet.defer import maybeDeferred
 
63
from twisted.python.threadable import isInIOThread
 
64
 
 
65
# A table suitable for use with str.translate() to replace each
 
66
# non-printable and non-ASCII character in a byte string with a question
 
67
# mark, mimicking the "replace" strategy when encoding and decoding.
 
68
non_printable_replace_table = b"".join(
 
69
    chr(i) if chr(i) in string.printable else b"?"
 
70
    for i in xrange(0xff + 0x01))
 
71
 
 
72
 
 
73
class ExternalProcessError(CalledProcessError):
 
74
    """Raised when there's a problem calling an external command.
 
75
 
 
76
    Unlike `CalledProcessError`:
 
77
 
 
78
    - `__str__()` returns a string containing the output of the failed
 
79
      external process, if available. All non-printable and non-ASCII
 
80
      characters are filtered out, replaced by question marks.
 
81
 
 
82
    - `__unicode__()` is defined, and tries to return something
 
83
      analagous to `__str__()` but keeping in valid unicode characters
 
84
      from the error message.
 
85
 
 
86
    """
 
87
 
 
88
    @staticmethod
 
89
    def _to_unicode(string):
 
90
        if isinstance(string, bytes):
 
91
            return string.decode("ascii", "replace")
 
92
        else:
 
93
            return unicode(string)
 
94
 
 
95
    @staticmethod
 
96
    def _to_ascii(string, table=non_printable_replace_table):
 
97
        if isinstance(string, unicode):
 
98
            return string.encode("ascii", "replace")
 
99
        else:
 
100
            return bytes(string).translate(table)
 
101
 
 
102
    def __unicode__(self):
 
103
        cmd = u" ".join(quote(self._to_unicode(part)) for part in self.cmd)
 
104
        output = self._to_unicode(self.output)
 
105
        return u"Command `%s` returned non-zero exit status %d:\n%s" % (
 
106
            cmd, self.returncode, output)
 
107
 
 
108
    def __str__(self):
 
109
        cmd = b" ".join(quote(self._to_ascii(part)) for part in self.cmd)
 
110
        output = self._to_ascii(self.output)
 
111
        return b"Command `%s` returned non-zero exit status %d:\n%s" % (
 
112
            cmd, self.returncode, output)
 
113
 
 
114
 
 
115
def call_and_check(command, *args, **kwargs):
 
116
    """A wrapper around subprocess.check_call().
 
117
 
 
118
    When an error occurs, raise an ExternalProcessError.
 
119
    """
 
120
    try:
 
121
        return subprocess.check_call(command, *args, **kwargs)
 
122
    except subprocess.CalledProcessError as error:
 
123
        error.__class__ = ExternalProcessError
 
124
        raise
 
125
 
 
126
 
 
127
def call_capture_and_check(command, *args, **kwargs):
 
128
    """A wrapper around subprocess.check_output().
 
129
 
 
130
    When an error occurs, raise an ExternalProcessError.
 
131
    """
 
132
    try:
 
133
        return subprocess.check_output(command, *args, **kwargs)
 
134
    except subprocess.CalledProcessError as error:
 
135
        error.__class__ = ExternalProcessError
 
136
        raise
 
137
 
 
138
 
 
139
def locate_config(*path):
 
140
    """Return the location of a given config file or directory.
 
141
 
 
142
    Defaults to `/etc/maas` (followed by any further path elements you
 
143
    specify), but can be overridden using the `MAAS_CONFIG_DIR` environment
 
144
    variable.  (When running from a branch, this variable will point to the
 
145
    `etc/maas` inside the branch.)
 
146
 
 
147
    The result is absolute and normalized.
 
148
    """
 
149
    # Check for MAAS_CONFIG_DIR.  Count empty string as "not set."
 
150
    env_setting = os.getenv('MAAS_CONFIG_DIR', '')
 
151
    if env_setting == '':
 
152
        # Running from installed package.  Config is in /etc/maas.
 
153
        config_dir = '/etc/maas'
 
154
    else:
 
155
        # Running from branch or other customized setup.  Config is at
 
156
        # $MAAS_CONFIG_DIR/etc/maas.
 
157
        config_dir = env_setting
 
158
 
 
159
    return os.path.abspath(os.path.join(config_dir, *path))
 
160
 
 
161
 
 
162
def find_settings(whence):
 
163
    """Return settings from `whence`, which is assumed to be a module."""
 
164
    # XXX 2012-10-11 JeroenVermeulen, bug=1065456: Put this in a shared
 
165
    # location.  It's currently duplicated from elsewhere.
 
166
    return {
 
167
        name: value
 
168
        for name, value in vars(whence).items()
 
169
        if not name.startswith("_")
 
170
        }
 
171
 
 
172
 
 
173
def import_settings(whence):
 
174
    """Import settings from `whence` into the caller's global scope."""
 
175
    # XXX 2012-10-11 JeroenVermeulen, bug=1065456: Put this in a shared
 
176
    # location.  It's currently duplicated from elsewhere.
 
177
    source = find_settings(whence)
 
178
    target = sys._getframe(1).f_globals
 
179
    target.update(source)
 
180
 
 
181
 
 
182
def deferred(func):
 
183
    """Decorates a function to ensure that it always returns a `Deferred`.
 
184
 
 
185
    This also serves a secondary documentation purpose; functions decorated
 
186
    with this are readily identifiable as asynchronous.
 
187
    """
 
188
    @wraps(func)
 
189
    def wrapper(*args, **kwargs):
 
190
        return maybeDeferred(func, *args, **kwargs)
 
191
    return wrapper
 
192
 
 
193
 
 
194
def asynchronous(func):
 
195
    """Decorates a function to ensure that it always runs in the reactor.
 
196
 
 
197
    If the wrapper is called from the reactor thread, it will call
 
198
    straight through to the wrapped function. It will not be wrapped by
 
199
    `maybeDeferred` for example.
 
200
 
 
201
    If the wrapper is called from another thread, it will return a
 
202
    :class:`crochet.EventualResult`, as if it had been decorated with
 
203
    `crochet.run_in_reactor`.
 
204
 
 
205
    This also serves a secondary documentation purpose; functions decorated
 
206
    with this are readily identifiable as asynchronous.
 
207
    """
 
208
    func_in_reactor = run_in_reactor(func)
 
209
 
 
210
    @wraps(func)
 
211
    def wrapper(*args, **kwargs):
 
212
        if isInIOThread():
 
213
            return func(*args, **kwargs)
 
214
        else:
 
215
            return func_in_reactor(*args, **kwargs)
 
216
    return wrapper
 
217
 
 
218
 
 
219
def synchronous(func):
 
220
    """Decorator to ensure that `func` never runs in the reactor thread.
 
221
 
 
222
    If the wrapped function is called from the reactor thread, this will
 
223
    raise a :class:`AssertionError`, implying that this is a programming
 
224
    error. Calls from outside the reactor will proceed unaffected.
 
225
 
 
226
    There is an asymettry with the `asynchronous` decorator. The reason
 
227
    is that it is essential to be aware when `deferToThread()` is being
 
228
    used, so that in-reactor code knows to synchronise with it, to add a
 
229
    callback to the :class:`Deferred` that it returns, for example. The
 
230
    expectation with `asynchronous` is that the return value is always
 
231
    important, and will be appropriate to the environment in which it is
 
232
    utilised.
 
233
 
 
234
    This also serves a secondary documentation purpose; functions decorated
 
235
    with this are readily identifiable as synchronous, or blocking.
 
236
 
 
237
    :raises AssertionError: When called inside the reactor thread.
 
238
    """
 
239
    @wraps(func)
 
240
    def wrapper(*args, **kwargs):
 
241
        if isInIOThread():
 
242
            raise AssertionError(
 
243
                "Function %s(...) must not be called in the "
 
244
                "reactor thread." % func.__name__)
 
245
        else:
 
246
            return func(*args, **kwargs)
 
247
    return wrapper
 
248
 
 
249
 
 
250
def filter_dict(dictionary, desired_keys):
 
251
    """Return a version of `dictionary` restricted to `desired_keys`.
 
252
 
 
253
    This is like a set union, except the values from `dictionary` come along.
 
254
    (Actually `desired_keys` can be a `dict`, but its values will be ignored).
 
255
    """
 
256
    return {
 
257
        key: value
 
258
        for key, value in dictionary.items()
 
259
        if key in desired_keys
 
260
    }
 
261
 
 
262
 
 
263
def _write_temp_file(content, filename):
 
264
    """Write the given `content` in a temporary file next to `filename`."""
 
265
    # Write the file to a temporary place (next to the target destination,
 
266
    # to ensure that it is on the same filesystem).
 
267
    directory = os.path.dirname(filename)
 
268
    prefix = ".%s." % os.path.basename(filename)
 
269
    suffix = ".tmp"
 
270
    try:
 
271
        temp_fd, temp_file = tempfile.mkstemp(
 
272
            dir=directory, suffix=suffix, prefix=prefix)
 
273
    except OSError, error:
 
274
        if error.filename is None:
 
275
            error.filename = os.path.join(
 
276
                directory, prefix + "XXXXXX" + suffix)
 
277
        raise
 
278
    else:
 
279
        with os.fdopen(temp_fd, "wb") as f:
 
280
            f.write(content)
 
281
            # Finish writing this file to the filesystem, and then, tell the
 
282
            # filesystem to push it down onto persistent storage.  This
 
283
            # prevents a nasty hazard in aggressively optimized filesystems
 
284
            # where you replace an old but consistent file with a new one that
 
285
            # is still in cache, and lose power before the new file can be made
 
286
            # fully persistent.
 
287
            # This was a particular problem with ext4 at one point; it may
 
288
            # still be.
 
289
            f.flush()
 
290
            os.fsync(f)
 
291
        return temp_file
 
292
 
 
293
 
 
294
def atomic_write(content, filename, overwrite=True, mode=0600):
 
295
    """Write `content` into the file `filename` in an atomic fashion.
 
296
 
 
297
    This requires write permissions to the directory that `filename` is in.
 
298
    It creates a temporary file in the same directory (so that it will be
 
299
    on the same filesystem as the destination) and then renames it to
 
300
    replace the original, if any.  Such a rename is atomic in POSIX.
 
301
 
 
302
    :param overwrite: Overwrite `filename` if it already exists?  Default
 
303
        is True.
 
304
    :param mode: Access permissions for the file, if written.
 
305
    """
 
306
    temp_file = _write_temp_file(content, filename)
 
307
    os.chmod(temp_file, mode)
 
308
    try:
 
309
        if overwrite:
 
310
            os.rename(temp_file, filename)
 
311
        else:
 
312
            lock = FileLock(filename)
 
313
            lock.acquire()
 
314
            try:
 
315
                if not os.path.isfile(filename):
 
316
                    os.rename(temp_file, filename)
 
317
            finally:
 
318
                lock.release()
 
319
    finally:
 
320
        if os.path.isfile(temp_file):
 
321
            os.remove(temp_file)
 
322
 
 
323
 
 
324
def incremental_write(content, filename, mode=0600):
 
325
    """Write the given `content` into the file `filename` and
 
326
    increment the modification time by 1 sec.
 
327
 
 
328
    :param mode: Access permissions for the file.
 
329
    """
 
330
    old_mtime = get_mtime(filename)
 
331
    atomic_write(content, filename, mode=mode)
 
332
    new_mtime = pick_new_mtime(old_mtime)
 
333
    os.utime(filename, (new_mtime, new_mtime))
 
334
 
 
335
 
 
336
def get_mtime(filename):
 
337
    """Return a file's modification time, or None if it does not exist."""
 
338
    try:
 
339
        return os.stat(filename).st_mtime
 
340
    except OSError as e:
 
341
        if e.errno == errno.ENOENT:
 
342
            # File does not exist.  Be helpful, return None.
 
343
            return None
 
344
        else:
 
345
            # Other failure.  The caller will want to know.
 
346
            raise
 
347
 
 
348
 
 
349
def pick_new_mtime(old_mtime=None, starting_age=1000):
 
350
    """Choose a new modification time for a file that needs it updated.
 
351
 
 
352
    This function is used to manage the modification time of files
 
353
    for which we need to see an increment in the modification time
 
354
    each time the file is modified.  This is the case for DNS zone
 
355
    files which only get properly reloaded if BIND sees that the
 
356
    modification time is > to the time it has in its database.
 
357
 
 
358
    Modification time can have a resolution as low as one second in
 
359
    some relevant environments (we have observed this with ext3).
 
360
    To produce mtime changes regardless, we set a file's modification
 
361
    time in the past when it is first written, and
 
362
    increment it by 1 second on each subsequent write.
 
363
 
 
364
    However we also want to be careful not to set the modification time
 
365
    in the future, mostly because BIND does not deal with that very
 
366
    well.
 
367
 
 
368
    :param old_mtime: File's previous modification time, as a number
 
369
        with a unity of one second, or None if it did not previously
 
370
        exist.
 
371
    :param starting_age: If the file did not exist previously, set its
 
372
        modification time this many seconds in the past.
 
373
    """
 
374
    now = time()
 
375
    if old_mtime is None:
 
376
        # File is new.  Set modification time in the past to have room for
 
377
        # sub-second modifications.
 
378
        return now - starting_age
 
379
    elif old_mtime + 1 <= now:
 
380
        # There is room to increment the file's mtime by one second
 
381
        # without ending up in the future.
 
382
        return old_mtime + 1
 
383
    else:
 
384
        # We can't increase the file's modification time.  Give up and
 
385
        # return the previous modification time.
 
386
        return old_mtime
 
387
 
 
388
 
 
389
def split_lines(input, separator):
 
390
    """Split each item from `input` into a key/value pair."""
 
391
    return (line.split(separator, 1) for line in input if line.strip() != '')
 
392
 
 
393
 
 
394
def strip_pairs(input):
 
395
    """Strip whitespace of each key/value pair in input."""
 
396
    return ((key.strip(), value.strip()) for (key, value) in input)
 
397
 
 
398
 
 
399
def parse_key_value_file(file_name, separator=":"):
 
400
    """Parse a text file into a dict of key/value pairs.
 
401
 
 
402
    Use this for simple key:value or key=value files. There are no
 
403
    sections, as required for python's ConfigParse. Whitespace and empty
 
404
    lines are ignored.
 
405
 
 
406
    :param file_name: Name of file to parse.
 
407
    :param separator: The text that separates each key from its value.
 
408
    """
 
409
    with open(file_name, 'rb') as input:
 
410
        return dict(strip_pairs(split_lines(input, separator)))
 
411
 
 
412
 
 
413
# Header and footer comments for MAAS custom config sections, as managed
 
414
# by write_custom_config_section.
 
415
maas_custom_config_markers = (
 
416
    "## Begin MAAS settings.  Do not edit; MAAS will overwrite this section.",
 
417
    "## End MAAS settings.",
 
418
    )
 
419
 
 
420
 
 
421
def find_list_item(item, in_list, starting_at=0):
 
422
    """Return index of `item` in `in_list`, or None if not found."""
 
423
    try:
 
424
        return in_list.index(item, starting_at)
 
425
    except ValueError:
 
426
        return None
 
427
 
 
428
 
 
429
def write_custom_config_section(original_text, custom_section):
 
430
    """Insert or replace a custom section in a configuration file's text.
 
431
 
 
432
    This allows you to rewrite configuration files that are not owned by
 
433
    MAAS, but where MAAS will have one section for its own settings.  It
 
434
    doesn't read or write any files; this is a pure text operation.
 
435
 
 
436
    Appends `custom_section` to the end of `original_text` if there was no
 
437
    custom MAAS section yet.  Otherwise, replaces the existing custom MAAS
 
438
    section with `custom_section`.  Returns the new text.
 
439
 
 
440
    Assumes that the configuration file's format accepts lines starting with
 
441
    hash marks (#) as comments.  The custom section will be bracketed by
 
442
    special marker comments that make it clear that MAAS wrote the section
 
443
    and it should not be edited by hand.
 
444
 
 
445
    :param original_text: The config file's current text.
 
446
    :type original_text: unicode
 
447
    :param custom_section: Custom config section to insert.
 
448
    :type custom_section: unicode
 
449
    :return: New config file text.
 
450
    :rtype: unicode
 
451
    """
 
452
    header, footer = maas_custom_config_markers
 
453
    lines = original_text.splitlines()
 
454
    header_index = find_list_item(header, lines)
 
455
    if header_index is not None:
 
456
        footer_index = find_list_item(footer, lines, header_index)
 
457
        if footer_index is None:
 
458
            # There's a header but no footer.  Pretend we didn't see the
 
459
            # header; just append a new custom section at the end.  Any
 
460
            # subsequent rewrite will replace the part starting at the
 
461
            # header and ending at the header we will add here.  At that
 
462
            # point there will be no trace of the strange situation
 
463
            # left.
 
464
            header_index = None
 
465
 
 
466
    if header_index is None:
 
467
        # There was no MAAS custom section in this file.  Append it at
 
468
        # the end.
 
469
        lines += [
 
470
            header,
 
471
            custom_section,
 
472
            footer,
 
473
            ]
 
474
    else:
 
475
        # There is a MAAS custom section in the file.  Replace it.
 
476
        lines = (
 
477
            lines[:(header_index + 1)] +
 
478
            [custom_section] +
 
479
            lines[footer_index:])
 
480
 
 
481
    return '\n'.join(lines) + '\n'
 
482
 
 
483
 
 
484
def sudo_write_file(filename, contents, encoding='utf-8', mode=0644):
 
485
    """Write (or overwrite) file as root.  USE WITH EXTREME CARE.
 
486
 
 
487
    Runs an atomic update using non-interactive `sudo`.  This will fail if
 
488
    it needs to prompt for a password.
 
489
    """
 
490
    raw_contents = contents.encode(encoding)
 
491
    command = [
 
492
        'sudo', '-n', 'maas-provision', 'atomic-write',
 
493
        '--filename', filename,
 
494
        '--mode', oct(mode),
 
495
        ]
 
496
    proc = Popen(command, stdin=PIPE)
 
497
    stdout, stderr = proc.communicate(raw_contents)
 
498
    if proc.returncode != 0:
 
499
        raise ExternalProcessError(proc.returncode, command, stderr)
 
500
 
 
501
 
 
502
class Safe:
 
503
    """An object that is safe to render as-is."""
 
504
 
 
505
    __slots__ = ("value",)
 
506
 
 
507
    def __init__(self, value):
 
508
        self.value = value
 
509
 
 
510
    def __repr__(self):
 
511
        return "<%s %r>" % (
 
512
            self.__class__.__name__, self.value)
 
513
 
 
514
 
 
515
class ShellTemplate(tempita.Template):
 
516
    """A Tempita template specialised for writing shell scripts.
 
517
 
 
518
    By default, substitutions will be escaped using `pipes.quote`, unless
 
519
    they're marked as safe. This can be done using Tempita's filter syntax::
 
520
 
 
521
      {{foobar|safe}}
 
522
 
 
523
    or as a plain Python expression::
 
524
 
 
525
      {{safe(foobar)}}
 
526
 
 
527
    """
 
528
 
 
529
    default_namespace = dict(
 
530
        tempita.Template.default_namespace,
 
531
        safe=Safe)
 
532
 
 
533
    def _repr(self, value, pos):
 
534
        """Shell-quote the value by default."""
 
535
        rep = super(ShellTemplate, self)._repr
 
536
        if isinstance(value, Safe):
 
537
            return rep(value.value, pos)
 
538
        else:
 
539
            return quote(rep(value, pos))
 
540
 
 
541
 
 
542
class ActionScript:
 
543
    """A command-line script that follows a command+verb pattern."""
 
544
 
 
545
    def __init__(self, description):
 
546
        super(ActionScript, self).__init__()
 
547
        # See http://docs.python.org/release/2.7/library/argparse.html.
 
548
        self.parser = ArgumentParser(description=description)
 
549
        self.subparsers = self.parser.add_subparsers(title="actions")
 
550
 
 
551
    @staticmethod
 
552
    def setup():
 
553
        # Ensure stdout and stderr are line-bufferred.
 
554
        sys.stdout = fdopen(sys.stdout.fileno(), "ab", 1)
 
555
        sys.stderr = fdopen(sys.stderr.fileno(), "ab", 1)
 
556
        # Run the SIGINT handler on SIGTERM; `svc -d` sends SIGTERM.
 
557
        signal.signal(signal.SIGTERM, signal.default_int_handler)
 
558
 
 
559
    def register(self, name, handler, *args, **kwargs):
 
560
        """Register an action for the given name.
 
561
 
 
562
        :param name: The name of the action.
 
563
        :param handler: An object, a module for example, that has `run` and
 
564
            `add_arguments` callables. The docstring of the `run` callable is
 
565
            used as the help text for the newly registered action.
 
566
        :param args: Additional positional arguments for the subparser_.
 
567
        :param kwargs: Additional named arguments for the subparser_.
 
568
 
 
569
        .. _subparser:
 
570
          http://docs.python.org/
 
571
            release/2.7/library/argparse.html#sub-commands
 
572
        """
 
573
        parser = self.subparsers.add_parser(
 
574
            name, *args, help=handler.run.__doc__, **kwargs)
 
575
        parser.set_defaults(handler=handler)
 
576
        handler.add_arguments(parser)
 
577
        return parser
 
578
 
 
579
    def execute(self, argv=None):
 
580
        """Execute this action.
 
581
 
 
582
        This is intended for in-process invocation of an action, though it may
 
583
        still raise L{SystemExit}. The L{__call__} method is intended for when
 
584
        this object is executed as a script proper.
 
585
        """
 
586
        args = self.parser.parse_args(argv)
 
587
        args.handler.run(args)
 
588
 
 
589
    def __call__(self, argv=None):
 
590
        try:
 
591
            self.setup()
 
592
            self.execute(argv)
 
593
        except CalledProcessError as error:
 
594
            # Print error.cmd and error.output too?
 
595
            raise SystemExit(error.returncode)
 
596
        except KeyboardInterrupt:
 
597
            raise SystemExit(1)
 
598
        else:
 
599
            raise SystemExit(0)
 
600
 
 
601
 
 
602
class MainScript(ActionScript):
 
603
    """An `ActionScript` that always accepts a `--config-file` option.
 
604
 
 
605
    The `--config-file` option defaults to the value of
 
606
    `MAAS_PROVISIONING_SETTINGS` in the process's environment, or absent
 
607
    that, `$MAAS_CONFIG_DIR/pserv.yaml` (normally /etc/maas/pserv.yaml for
 
608
    packaged installations, or when running from branch, the equivalent
 
609
    inside that branch).
 
610
    """
 
611
 
 
612
    def __init__(self, description):
 
613
        # Avoid circular imports.
 
614
        from provisioningserver.config import Config
 
615
 
 
616
        super(MainScript, self).__init__(description)
 
617
        self.parser.add_argument(
 
618
            "-c", "--config-file", metavar="FILENAME",
 
619
            help="Configuration file to load [%(default)s].",
 
620
            default=Config.DEFAULT_FILENAME)
 
621
 
 
622
 
 
623
class AtomicWriteScript:
 
624
    """Wrap the atomic_write function turning it into an ActionScript.
 
625
 
 
626
    To use:
 
627
    >>> main = MainScript(atomic_write.__doc__)
 
628
    >>> main.register("myscriptname", AtomicWriteScript)
 
629
    >>> main()
 
630
    """
 
631
 
 
632
    @staticmethod
 
633
    def add_arguments(parser):
 
634
        """Initialise options for writing files atomically.
 
635
 
 
636
        :param parser: An instance of :class:`ArgumentParser`.
 
637
        """
 
638
        parser.add_argument(
 
639
            "--no-overwrite", action="store_true", required=False,
 
640
            default=False, help="Don't overwrite file if it exists")
 
641
        parser.add_argument(
 
642
            "--filename", action="store", required=True, help=(
 
643
                "The name of the file in which to store contents of stdin"))
 
644
        parser.add_argument(
 
645
            "--mode", action="store", required=False, default=None, help=(
 
646
                "They permissions to set on the file. If not set "
 
647
                "will be r/w only to owner"))
 
648
 
 
649
    @staticmethod
 
650
    def run(args):
 
651
        """Take content from stdin and write it atomically to a file."""
 
652
        content = sys.stdin.read()
 
653
        if args.mode is not None:
 
654
            mode = int(args.mode, 8)
 
655
        else:
 
656
            mode = 0600
 
657
        atomic_write(
 
658
            content, args.filename, overwrite=not args.no_overwrite,
 
659
            mode=mode)
 
660
 
 
661
 
 
662
def get_all_interface_addresses():
 
663
    """For each network interface, yield its IPv4 address."""
 
664
    for interface in netifaces.interfaces():
 
665
        addresses = netifaces.ifaddresses(interface)
 
666
        if netifaces.AF_INET in addresses:
 
667
            for inet_address in addresses[netifaces.AF_INET]:
 
668
                if "addr" in inet_address:
 
669
                    yield inet_address["addr"]
 
670
 
 
671
 
 
672
def ensure_dir(path):
 
673
    """Do the equivalent of `mkdir -p`, creating `path` if it didn't exist."""
 
674
    try:
 
675
        os.makedirs(path)
 
676
    except OSError as e:
 
677
        if e.errno != errno.EEXIST:
 
678
            raise
 
679
        if not isdir(path):
 
680
            # Path exists, but isn't a directory.
 
681
            raise
 
682
        # Otherwise, the error is that the directory already existed.
 
683
        # Which is actually success.
 
684
 
 
685
 
 
686
@contextmanager
 
687
def tempdir(suffix=b'', prefix=b'maas-', location=None):
 
688
    """Context manager: temporary directory.
 
689
 
 
690
    Creates a temporary directory (yielding its path, as `unicode`), and
 
691
    cleans it up again when exiting the context.
 
692
 
 
693
    The directory will be readable, writable, and searchable only to the
 
694
    system user who creates it.
 
695
 
 
696
    >>> with tempdir() as playground:
 
697
    ...     my_file = os.path.join(playground, "my-file")
 
698
    ...     with open(my_file, 'wb') as handle:
 
699
    ...         handle.write(b"Hello.\n")
 
700
    ...     files = os.listdir(playground)
 
701
    >>> files
 
702
    [u'my-file']
 
703
    >>> os.path.isdir(playground)
 
704
    False
 
705
    """
 
706
    path = tempfile.mkdtemp(suffix, prefix, location)
 
707
    if isinstance(path, bytes):
 
708
        path = path.decode(sys.getfilesystemencoding())
 
709
    assert isinstance(path, unicode)
 
710
    try:
 
711
        yield path
 
712
    finally:
 
713
        rmtree(path, ignore_errors=True)
 
714
 
 
715
 
 
716
def read_text_file(path, encoding='utf-8'):
 
717
    """Read and decode the text file at the given path."""
 
718
    with codecs.open(path, encoding=encoding) as infile:
 
719
        return infile.read()
 
720
 
 
721
 
 
722
def write_text_file(path, text, encoding='utf-8'):
 
723
    """Write the given unicode text to the given file path.
 
724
 
 
725
    If the file existed, it will be overwritten.
 
726
    """
 
727
    with codecs.open(path, 'w', encoding) as outfile:
 
728
        outfile.write(text)
 
729
 
 
730
 
 
731
def is_compiled_xpath(xpath):
 
732
    """Is `xpath` a compiled expression?"""
 
733
    return isinstance(xpath, etree.XPath)
 
734
 
 
735
 
 
736
def is_compiled_doc(doc):
 
737
    """Is `doc` a compiled XPath document evaluator?"""
 
738
    return isinstance(doc, etree.XPathDocumentEvaluator)
 
739
 
 
740
 
 
741
def match_xpath(xpath, doc):
 
742
    """Return a match of expression `xpath` against document `doc`.
 
743
 
 
744
    :type xpath: Either `unicode` or `etree.XPath`
 
745
    :type doc: Either `etree._ElementTree` or `etree.XPathDocumentEvaluator`
 
746
 
 
747
    :rtype: bool
 
748
    """
 
749
    is_xpath_compiled = is_compiled_xpath(xpath)
 
750
    is_doc_compiled = is_compiled_doc(doc)
 
751
 
 
752
    if is_xpath_compiled and is_doc_compiled:
 
753
        return doc(xpath.path)
 
754
    elif is_xpath_compiled:
 
755
        return xpath(doc)
 
756
    elif is_doc_compiled:
 
757
        return doc(xpath)
 
758
    else:
 
759
        return doc.xpath(xpath)
 
760
 
 
761
 
 
762
def try_match_xpath(xpath, doc, logger=logging):
 
763
    """See if the XPath expression matches the given XML document.
 
764
 
 
765
    Invalid XPath expressions are logged, and are returned as a
 
766
    non-match.
 
767
 
 
768
    :type xpath: Either `unicode` or `etree.XPath`
 
769
    :type doc: Either `etree._ElementTree` or `etree.XPathDocumentEvaluator`
 
770
 
 
771
    :rtype: bool
 
772
    """
 
773
    try:
 
774
        # Evaluating an XPath expression against a document with LXML
 
775
        # can return a list or a string, and perhaps other types.
 
776
        # Casting the return value into a boolean context appears to
 
777
        # be the most reliable way of detecting a match.
 
778
        return bool(match_xpath(xpath, doc))
 
779
    except etree.XPathEvalError:
 
780
        # Get a plaintext version of `xpath`.
 
781
        expr = xpath.path if is_compiled_xpath(xpath) else xpath
 
782
        logger.exception("Invalid expression: %s", expr)
 
783
        return False
 
784
 
 
785
 
 
786
def classify(func, subjects):
 
787
    """Classify `subjects` according to `func`.
 
788
 
 
789
    Splits `subjects` into two lists: one for those which `func`
 
790
    returns a truth-like value, and one for the others.
 
791
 
 
792
    :param subjects: An iterable of `(ident, subject)` tuples, where
 
793
        `subject` is an argument that can be passed to `func` for
 
794
        classification.
 
795
    :param func: A function that takes a single argument.
 
796
 
 
797
    :return: A ``(matched, other)`` tuple, where ``matched`` and
 
798
        ``other`` are `list`s of `ident` values; `subject` values are
 
799
        not returned.
 
800
    """
 
801
    matched, other = [], []
 
802
    for ident, subject in subjects:
 
803
        bucket = matched if func(subject) else other
 
804
        bucket.append(ident)
 
805
    return matched, other
 
806
 
 
807
 
 
808
def find_ip_via_arp(mac):
 
809
    """Find the IP address for `mac` by reading the output of arp -n.
 
810
 
 
811
    Returns `None` if the MAC is not found.
 
812
 
 
813
    We do this because we aren't necessarily the only DHCP server on the
 
814
    network, so we can't check our own leases file and be guaranteed to find an
 
815
    IP that matches.
 
816
 
 
817
    :param mac: The mac address, e.g. '1c:6f:65:d5:56:98'.
 
818
    """
 
819
 
 
820
    output = call_capture_and_check(['arp', '-n']).split('\n')
 
821
 
 
822
    for line in output:
 
823
        columns = line.split()
 
824
        if len(columns) == 5 and columns[2] == mac:
 
825
            return columns[0]
 
826
    return None