# 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))