1# Copyright 2017 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Kludges to support legacy Autotest code.
6
7Autotest imports should be done by calling monkeypatch() first and then
8calling load().  monkeypatch() should only be called once from a
9script's main function.
10
11chromite imports should be done with chromite_load(), and any third
12party packages should be imported with deps_load().  The reason for this
13is to present a clear API for these unsafe imports, making it easier to
14identify which imports are currently unsafe.  Eventually, everything
15should be moved to virtualenv, but that will not be in the near future.
16"""
17
18from __future__ import absolute_import
19from __future__ import division
20from __future__ import print_function
21
22import ast
23import contextlib
24import imp
25import importlib
26import logging
27import os
28import site
29import subprocess
30import sys
31
32import autotest_lib
33
34AUTOTEST_DIR = autotest_lib.__path__[0]
35_SITEPKG_DIR = os.path.join(AUTOTEST_DIR, 'site-packages')
36_SYSTEM_PYTHON = '/usr/bin/python2.7'
37
38_setup_done = False
39
40logger = logging.getLogger(__name__)
41
42
43def monkeypatch():
44    """Monkeypatch everything needed to import Autotest.
45
46    This should be called before any calls to load().  Only the main
47    function in scripts should call this function.
48    """
49    with _global_setup():
50        _monkeypatch_body()
51
52
53@contextlib.contextmanager
54def _global_setup():
55    """Context manager for checking and setting global _setup_done variable."""
56    global _setup_done
57    assert not _setup_done
58    try:
59        yield
60    except Exception:  # pragma: no cover
61        # We cannot recover from this since we leave the interpreter in
62        # an unknown state.
63        logger.exception('Uncaught exception escaped Autotest setup')
64        sys.exit(1)
65    else:
66        _setup_done = True
67
68
69def _monkeypatch_body():
70    """The body of monkeypatch() running within _global_setup() context."""
71    # Add Autotest's site-packages.
72    site.addsitedir(_SITEPKG_DIR)
73
74    # Dummy out common imports as they may cause problems.
75    sys.meta_path.insert(0, _CommonRemovingFinder())
76
77    # Add chromite's third-party to the import path (chromite does this
78    # on import).
79    try:
80        importlib.import_module('chromite')
81    except ImportError:
82        # Moblab does not run build_externals; dependencies like
83        # chromite are installed system-wide.
84        logger.info("""\
85Could not find chromite; adding system packages and retrying \
86(This should only happen on Moblab)""")
87        for d in _system_site_packages():
88            site.addsitedir(d)
89        importlib.import_module('chromite')
90
91    # Set up Django environment variables.
92    importlib.import_module('autotest_lib.frontend.setup_django_environment')
93
94    # Make Django app paths absolute.
95    settings = importlib.import_module('autotest_lib.frontend.settings')
96    settings.INSTALLED_APPS = (
97            'autotest_lib.frontend.afe',
98            'autotest_lib.frontend.tko',
99            'django.contrib.admin',
100            'django.contrib.auth',
101            'django.contrib.contenttypes',
102            'django.contrib.sessions',
103            'django.contrib.sites',
104    )
105
106    # drone_utility uses this.
107    common = importlib.import_module('autotest_lib.scheduler.common')
108    common.autotest_dir = AUTOTEST_DIR
109
110
111def _system_site_packages():
112    """Get list of system site-package directories.
113
114    This is needed for Moblab because dependencies are installed
115    system-wide instead of using build_externals.py.
116    """
117    output = subprocess.check_output([
118        _SYSTEM_PYTHON, '-c',
119        'import site; print repr(site.getsitepackages())'])
120    return ast.literal_eval(output)
121
122
123class _CommonRemovingFinder(object):
124    """Python import finder that neuters Autotest's common.py
125
126    The common module is replaced with an empty module everywhere it is
127    imported.  common.py should have only been imported for side
128    effects, so nothing should actually use the imported module.
129
130    See also https://www.python.org/dev/peps/pep-0302/
131    """
132
133    def find_module(self, fullname, path=None):
134        """Find module."""
135        del path  # unused
136        if not self._is_autotest_common(fullname):
137            return None
138        logger.debug('Dummying out %s import', fullname)
139        return self
140
141    def _is_autotest_common(self, fullname):
142        return (fullname.partition('.')[0] == 'autotest_lib'
143                and fullname.rpartition('.')[-1] == 'common')
144
145    def load_module(self, fullname):
146        """Load module."""
147        if fullname in sys.modules:  # pragma: no cover
148            return sys.modules[fullname]
149        mod = imp.new_module(fullname)
150        mod.__file__ = '<removed>'
151        mod.__loader__ = self
152        mod.__package__ = fullname.rpartition('.')[0]
153        sys.modules[fullname] = mod
154        return mod
155
156
157def load(name):
158    """Import module from autotest.
159
160    This enforces that monkeypatch() is called first.
161
162    @param name: name of module as string, e.g., 'frontend.afe.models'
163    """
164    return _load('autotest_lib.%s' % name)
165
166
167def chromite_load(name):
168    """Import module from chromite.lib.
169
170    This enforces that monkeypatch() is called first.
171
172    @param name: name of module as string, e.g., 'metrics'
173    """
174    return _load('chromite.lib.%s' % name)
175
176
177def deps_load(name):
178    """Import module from chromite.lib.
179
180    This enforces that monkeypatch() is called first.
181
182    @param name: name of module as string, e.g., 'metrics'
183    """
184    assert not name.startswith('autotest_lib')
185    assert not name.startswith('chromite.lib')
186    return _load(name)
187
188
189def _load(name):
190    """Import a module.
191
192    This enforces that monkeypatch() is called first.
193
194    @param name: name of module as string
195    """
196    if not _setup_done:
197        raise ImportError('cannot load chromite modules before monkeypatching')
198    return importlib.import_module(name)
199