1# Copyright 2020 The Chromium OS 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
5import json
6import logging
7import os
8import requests
9import subprocess
10import urlparse
11
12from autotest_lib.client.bin import utils
13from autotest_lib.client.common_lib import autotemp
14from autotest_lib.client.common_lib import error
15
16
17# JSON attributes used in payload properties. Look at nebraska.py for more
18# information.
19KEY_PUBLIC_KEY='public_key'
20KEY_METADATA_SIZE='metadata_size'
21KEY_SHA256='sha256_hex'
22
23
24class NebraskaWrapper(object):
25    """
26    A wrapper around nebraska.py
27
28    This wrapper is used to start a nebraska.py service and allow the
29    update_engine to interact with it.
30
31    """
32
33    def __init__(self, log_dir=None, payload_url=None, **props_to_override):
34        """
35        Initializes the NebraskaWrapper module.
36
37        @param log_dir: The directory to write nebraska.log into.
38        @param payload_url: The payload that will be returned in responses for
39                            update requests. This can be a single URL string
40                            or a list of URLs to return multiple payload URLs
41                            (such as a platform payload + DLC payloads) in the
42                            responses.
43        @param props_to_override: Dictionary of key/values to use in responses
44                instead of the default values in payload_url's properties file.
45
46        """
47        self._nebraska_server = None
48        self._port = None
49        self._log_dir = log_dir
50
51        # _update_metadata_dir is the directory for storing the json metadata
52        # files associated with the payloads.
53        # _update_payloads_address is the address of the update server where
54        # the payloads are staged.
55        # The _install variables serve the same purpose for payloads intended
56        # for DLC install requests.
57        self._update_metadata_dir = None
58        self._update_payloads_address = None
59        self._install_metadata_dir = None
60        self._install_payloads_address = None
61
62        # Create a temporary directory for the metadata and download the
63        # metadata files.
64        if payload_url:
65            # Normalize payload_url to be a list.
66            if not isinstance(payload_url, list):
67                payload_url = [payload_url]
68
69            self._update_metadata_dir = autotemp.tempdir()
70            self._update_payloads_address = ''.join(
71                payload_url[0].rpartition('/')[0:2])
72            # We can reuse _update_metadata_dir and _update_payloads_address
73            # for the DLC-specific install values for N-N tests, since the
74            # install and update versions will be the same. For the delta
75            # payload case, Nebraska will always use a full payload for
76            # installation and prefer a delta payload for update, so both full
77            # and delta payload metadata files can occupy the same
78            # metadata_dir. The payloads_address can be shared as well,
79            # provided all payloads have the same base URL.
80            self._install_metadata_dir = self._update_metadata_dir
81            self._install_payloads_address = self._update_payloads_address
82
83            for url in payload_url:
84                self.get_payload_properties_file(
85                    url, self._update_metadata_dir.name,
86                    **props_to_override)
87
88    def __enter__(self):
89        """So that NebraskaWrapper can be used as a Context Manager."""
90        self.start()
91        return self
92
93    def __exit__(self, *exception_details):
94        """
95        So that NebraskaWrapper can be used as a Context Manager.
96
97        @param exception_details: Details of exceptions happened in the
98                ContextManager.
99
100        """
101        self.stop()
102
103    def start(self):
104        """
105        Starts the Nebraska server.
106
107        @raise error.TestError: If fails to start the Nebraska server.
108
109        """
110        # Any previously-existing files (port, pid and log files) will be
111        # overriden by Nebraska during bring up.
112        runtime_root = '/tmp/nebraska'
113        cmd = ['nebraska.py', '--runtime-root', runtime_root]
114        if self._log_dir:
115            cmd += ['--log-file', os.path.join(self._log_dir, 'nebraska.log')]
116        if self._update_metadata_dir:
117            cmd += ['--update-metadata', self._update_metadata_dir.name]
118        if self._update_payloads_address:
119            cmd += ['--update-payloads-address', self._update_payloads_address]
120        if self._install_metadata_dir:
121            cmd += ['--install-metadata', self._install_metadata_dir.name]
122        if self._install_payloads_address:
123            cmd += ['--install-payloads-address',
124                    self._install_payloads_address]
125
126        logging.info('Starting nebraska.py with command: %s', cmd)
127
128        try:
129            self._nebraska_server = subprocess.Popen(cmd,
130                                                     stdout=subprocess.PIPE,
131                                                     stderr=subprocess.STDOUT)
132
133            # Wait for port file to appear.
134            port_file = os.path.join(runtime_root, 'port')
135            utils.poll_for_condition(lambda: os.path.exists(port_file),
136                                     timeout=5)
137
138            with open(port_file, 'r') as f:
139                self._port = int(f.read())
140
141            # Send a health_check request to it to make sure its working.
142            requests.get('http://127.0.0.1:%d/health_check' % self._port)
143
144        except Exception as e:
145            raise error.TestError('Failed to start Nebraska %s' % e)
146
147    def stop(self):
148        """Stops the Nebraska server."""
149        if not self._nebraska_server:
150            return
151        try:
152            self._nebraska_server.terminate()
153            stdout, _ = self._nebraska_server.communicate()
154            logging.info('Stopping nebraska.py with stdout %s', stdout)
155            self._nebraska_server.wait()
156        except subprocess.TimeoutExpired:
157            logging.error('Failed to stop Nebraska. Ignoring...')
158        finally:
159            self._nebraska_server = None
160
161    def get_update_url(self, **kwargs):
162        """
163        Returns a URL for getting updates from this Nebraska instance.
164
165        @param kwargs: A set of key/values to form a search query to instruct
166                Nebraska to do a set of activities. See
167                nebraska.py::ResponseProperties for examples key/values.
168        """
169
170        query = '&'.join('%s=%s' % (k, v) for k, v in kwargs.items())
171        url = urlparse.SplitResult(scheme='http',
172                                   netloc='127.0.0.1:%d' % self._port,
173                                   path='/update',
174                                   query=query,
175                                   fragment='')
176        return urlparse.urlunsplit(url)
177
178    def get_payload_properties_file(self, payload_url, target_dir, **kwargs):
179        """
180        Downloads the payload properties file into a directory.
181
182        @param payload_url: The URL to the update payload file.
183        @param target_dir: The directory to download the file into.
184        @param kwargs: A dictionary of key/values that needs to be overridden on
185                the payload properties file.
186
187        """
188        payload_props_url = payload_url + '.json'
189        _, _, file_name = payload_props_url.rpartition('/')
190        try:
191            response = json.loads(requests.get(payload_props_url).text)
192            # Override existing keys if any.
193            for k, v in kwargs.iteritems():
194                # Don't set default None values. We don't want to override good
195                # values to None.
196                if v is not None:
197                    response[k] = v
198            with open(os.path.join(target_dir, file_name), 'w') as fp:
199                json.dump(response, fp)
200
201        except (requests.exceptions.RequestException,
202                IOError,
203                ValueError) as err:
204            raise error.TestError(
205                'Failed to get update payload properties: %s with error: %s' %
206                (payload_props_url, err))
207