# Lint as: python2, python3 # 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. from __future__ import absolute_import from __future__ import division from __future__ import print_function import errno import json import logging import os import requests import subprocess import six import six.moves.urllib.parse 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' # Path to the startup config file. NEBRASKA_DIR = '/usr/local/nebraska' NEBRASKA_CONFIG = os.path.join(NEBRASKA_DIR, 'config.json') NEBRASKA_METADATA_DIR = os.path.join(NEBRASKA_DIR, 'metadata') 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, persist_metadata=False, **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 persist_metadata: True to store the update and install metadata in a location that will survive a reboot. Use this if you plan on starting nebraska at system startup using a conf file. If False, the metadata will be stored in /tmp and will not persist after rebooting the device. @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 # Download the metadata files and save them in a tempdir for general # use, or in a directory that will survive reboot if we want nebraska # to be up after a reboot. If saving to a tempdir, save a reference # to it to ensure its reference count does not go to zero causing the # directory to be deleted. if payload_url: # Normalize payload_url to be a list. if not isinstance(payload_url, list): payload_url = [payload_url] if persist_metadata: self._create_nebraska_dir(metadata=True) self._update_metadata_dir = NEBRASKA_METADATA_DIR else: self._tempdir = autotemp.tempdir() self._update_metadata_dir = self._tempdir.name 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, **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] 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] 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 = six.moves.urllib.parse.SplitResult(scheme='http', netloc='127.0.0.1:%d' % self._port, path='/update', query=query, fragment='') return six.moves.urllib.parse.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 six.iteritems(kwargs): # 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)) def update_config(self, **kwargs): """ Updates the current running nebraska's config. @param kwargs: A dictionary of key/values to update the nebraska's config. See platform/dev/nebraska/nebraska.py for more information. """ requests.post('http://127.0.0.1:%d/update_config' % self._port, json=kwargs) def _create_nebraska_dir(self, metadata=True): """ Creates /usr/local/nebraska for storing the startup conf and persistent metadata files. @param metadata: True to create a subdir for metadata. """ dir_to_make = NEBRASKA_DIR if metadata: dir_to_make = NEBRASKA_METADATA_DIR try: os.makedirs(dir_to_make) except OSError as e: if errno.EEXIST != e.errno: raise error.TestError('Failed to create %s with error: %s', dir_to_make, e) def create_startup_config(self, **kwargs): """ Creates a nebraska startup config file. If this file is present, nebraska will start before update_engine does during system startup. @param kwargs: A dictionary of key/values for nebraska config options. See platform/dev/nebraska/nebraska.py for more info. """ conf = {} if self._update_metadata_dir: conf['update_metadata'] = self._update_metadata_dir if self._update_payloads_address: conf['update_payloads_address'] = self._update_payloads_address if self._install_metadata_dir: conf['install_metadata'] = self._install_metadata_dir if self._install_payloads_address: conf['install_payloads_address'] = self._install_payloads_address for k, v in six.iteritems(kwargs): conf[k] = v self._create_nebraska_dir() with open(NEBRASKA_CONFIG, 'w') as fp: json.dump(conf, fp)