| # Copyright 2020 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| import json |
| import logging |
| import os |
| import requests |
| import subprocess |
| import urlparse |
| |
| from autotest_lib.client.bin import utils |
| from autotest_lib.client.common_lib import autotemp |
| from autotest_lib.client.common_lib import error |
| |
| |
| # JSON attributes used in payload properties. Look at nebraska.py for more |
| # information. |
| KEY_PUBLIC_KEY='public_key' |
| KEY_METADATA_SIZE='metadata_size' |
| KEY_SHA256='sha256_hex' |
| |
| |
| class NebraskaWrapper(object): |
| """ |
| A wrapper around nebraska.py |
| |
| This wrapper is used to start a nebraska.py service and allow the |
| update_engine to interact with it. |
| |
| """ |
| |
| def __init__(self, log_dir=None, payload_url=None, **props_to_override): |
| """ |
| Initializes the NebraskaWrapper module. |
| |
| @param log_dir: The directory to write nebraska.log into. |
| @param payload_url: The payload that will be returned in responses for |
| update requests. This can be a single URL string |
| or a list of URLs to return multiple payload URLs |
| (such as a platform payload + DLC payloads) in the |
| responses. |
| @param props_to_override: Dictionary of key/values to use in responses |
| instead of the default values in payload_url's properties file. |
| |
| """ |
| self._nebraska_server = None |
| self._port = None |
| self._log_dir = log_dir |
| |
| # _update_metadata_dir is the directory for storing the json metadata |
| # files associated with the payloads. |
| # _update_payloads_address is the address of the update server where |
| # the payloads are staged. |
| # The _install variables serve the same purpose for payloads intended |
| # for DLC install requests. |
| self._update_metadata_dir = None |
| self._update_payloads_address = None |
| self._install_metadata_dir = None |
| self._install_payloads_address = None |
| |
| # Create a temporary directory for the metadata and download the |
| # metadata files. |
| if payload_url: |
| # Normalize payload_url to be a list. |
| if not isinstance(payload_url, list): |
| payload_url = [payload_url] |
| |
| self._update_metadata_dir = autotemp.tempdir() |
| self._update_payloads_address = ''.join( |
| payload_url[0].rpartition('/')[0:2]) |
| # We can reuse _update_metadata_dir and _update_payloads_address |
| # for the DLC-specific install values for N-N tests, since the |
| # install and update versions will be the same. For the delta |
| # payload case, Nebraska will always use a full payload for |
| # installation and prefer a delta payload for update, so both full |
| # and delta payload metadata files can occupy the same |
| # metadata_dir. The payloads_address can be shared as well, |
| # provided all payloads have the same base URL. |
| self._install_metadata_dir = self._update_metadata_dir |
| self._install_payloads_address = self._update_payloads_address |
| |
| for url in payload_url: |
| self.get_payload_properties_file( |
| url, self._update_metadata_dir.name, |
| **props_to_override) |
| |
| def __enter__(self): |
| """So that NebraskaWrapper can be used as a Context Manager.""" |
| self.start() |
| return self |
| |
| def __exit__(self, *exception_details): |
| """ |
| So that NebraskaWrapper can be used as a Context Manager. |
| |
| @param exception_details: Details of exceptions happened in the |
| ContextManager. |
| |
| """ |
| self.stop() |
| |
| def start(self): |
| """ |
| Starts the Nebraska server. |
| |
| @raise error.TestError: If fails to start the Nebraska server. |
| |
| """ |
| # Any previously-existing files (port, pid and log files) will be |
| # overriden by Nebraska during bring up. |
| runtime_root = '/tmp/nebraska' |
| cmd = ['nebraska.py', '--runtime-root', runtime_root] |
| if self._log_dir: |
| cmd += ['--log-file', os.path.join(self._log_dir, 'nebraska.log')] |
| if self._update_metadata_dir: |
| cmd += ['--update-metadata', self._update_metadata_dir.name] |
| if self._update_payloads_address: |
| cmd += ['--update-payloads-address', self._update_payloads_address] |
| if self._install_metadata_dir: |
| cmd += ['--install-metadata', self._install_metadata_dir.name] |
| if self._install_payloads_address: |
| cmd += ['--install-payloads-address', |
| self._install_payloads_address] |
| |
| logging.info('Starting nebraska.py with command: %s', cmd) |
| |
| try: |
| self._nebraska_server = subprocess.Popen(cmd, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT) |
| |
| # Wait for port file to appear. |
| port_file = os.path.join(runtime_root, 'port') |
| utils.poll_for_condition(lambda: os.path.exists(port_file), |
| timeout=5) |
| |
| with open(port_file, 'r') as f: |
| self._port = int(f.read()) |
| |
| # Send a health_check request to it to make sure its working. |
| requests.get('http://127.0.0.1:%d/health_check' % self._port) |
| |
| except Exception as e: |
| raise error.TestError('Failed to start Nebraska %s' % e) |
| |
| def stop(self): |
| """Stops the Nebraska server.""" |
| if not self._nebraska_server: |
| return |
| try: |
| self._nebraska_server.terminate() |
| stdout, _ = self._nebraska_server.communicate() |
| logging.info('Stopping nebraska.py with stdout %s', stdout) |
| self._nebraska_server.wait() |
| except subprocess.TimeoutExpired: |
| logging.error('Failed to stop Nebraska. Ignoring...') |
| finally: |
| self._nebraska_server = None |
| |
| def get_update_url(self, **kwargs): |
| """ |
| Returns a URL for getting updates from this Nebraska instance. |
| |
| @param kwargs: A set of key/values to form a search query to instruct |
| Nebraska to do a set of activities. See |
| nebraska.py::ResponseProperties for examples key/values. |
| """ |
| |
| query = '&'.join('%s=%s' % (k, v) for k, v in kwargs.items()) |
| url = urlparse.SplitResult(scheme='http', |
| netloc='127.0.0.1:%d' % self._port, |
| path='/update', |
| query=query, |
| fragment='') |
| return urlparse.urlunsplit(url) |
| |
| def get_payload_properties_file(self, payload_url, target_dir, **kwargs): |
| """ |
| Downloads the payload properties file into a directory. |
| |
| @param payload_url: The URL to the update payload file. |
| @param target_dir: The directory to download the file into. |
| @param kwargs: A dictionary of key/values that needs to be overridden on |
| the payload properties file. |
| |
| """ |
| payload_props_url = payload_url + '.json' |
| _, _, file_name = payload_props_url.rpartition('/') |
| try: |
| response = json.loads(requests.get(payload_props_url).text) |
| # Override existing keys if any. |
| for k, v in kwargs.iteritems(): |
| # Don't set default None values. We don't want to override good |
| # values to None. |
| if v is not None: |
| response[k] = v |
| with open(os.path.join(target_dir, file_name), 'w') as fp: |
| json.dump(response, fp) |
| |
| except (requests.exceptions.RequestException, |
| IOError, |
| ValueError) as err: |
| raise error.TestError( |
| 'Failed to get update payload properties: %s with error: %s' % |
| (payload_props_url, err)) |
