295 lines
12 KiB
Python
295 lines
12 KiB
Python
# 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)
|