1# Lint as: python2, python3 2# Copyright 2020 The Chromium OS Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6from __future__ import absolute_import 7from __future__ import division 8from __future__ import print_function 9 10import errno 11import json 12import logging 13import os 14import requests 15import subprocess 16import six 17import six.moves.urllib.parse 18 19from autotest_lib.client.bin import utils 20from autotest_lib.client.common_lib import autotemp 21from autotest_lib.client.common_lib import error 22 23 24# JSON attributes used in payload properties. Look at nebraska.py for more 25# information. 26KEY_PUBLIC_KEY='public_key' 27KEY_METADATA_SIZE='metadata_size' 28KEY_SHA256='sha256_hex' 29 30# Path to the startup config file. 31NEBRASKA_DIR = '/usr/local/nebraska' 32NEBRASKA_CONFIG = os.path.join(NEBRASKA_DIR, 'config.json') 33NEBRASKA_METADATA_DIR = os.path.join(NEBRASKA_DIR, 'metadata') 34 35 36class NebraskaWrapper(object): 37 """ 38 A wrapper around nebraska.py 39 40 This wrapper is used to start a nebraska.py service and allow the 41 update_engine to interact with it. 42 43 """ 44 45 def __init__(self, 46 log_dir=None, 47 payload_url=None, 48 persist_metadata=False, 49 **props_to_override): 50 """ 51 Initializes the NebraskaWrapper module. 52 53 @param log_dir: The directory to write nebraska.log into. 54 @param payload_url: The payload that will be returned in responses for 55 update requests. This can be a single URL string 56 or a list of URLs to return multiple payload URLs 57 (such as a platform payload + DLC payloads) in the 58 responses. 59 @param persist_metadata: True to store the update and install metadata 60 in a location that will survive a reboot. Use 61 this if you plan on starting nebraska at 62 system startup using a conf file. If False, 63 the metadata will be stored in /tmp and will 64 not persist after rebooting the device. 65 @param props_to_override: Dictionary of key/values to use in responses 66 instead of the default values in payload_url's properties file. 67 68 """ 69 self._nebraska_server = None 70 self._port = None 71 self._log_dir = log_dir 72 73 # _update_metadata_dir is the directory for storing the json metadata 74 # files associated with the payloads. 75 # _update_payloads_address is the address of the update server where 76 # the payloads are staged. 77 # The _install variables serve the same purpose for payloads intended 78 # for DLC install requests. 79 self._update_metadata_dir = None 80 self._update_payloads_address = None 81 self._install_metadata_dir = None 82 self._install_payloads_address = None 83 84 # Download the metadata files and save them in a tempdir for general 85 # use, or in a directory that will survive reboot if we want nebraska 86 # to be up after a reboot. If saving to a tempdir, save a reference 87 # to it to ensure its reference count does not go to zero causing the 88 # directory to be deleted. 89 if payload_url: 90 # Normalize payload_url to be a list. 91 if not isinstance(payload_url, list): 92 payload_url = [payload_url] 93 94 if persist_metadata: 95 self._create_nebraska_dir(metadata=True) 96 self._update_metadata_dir = NEBRASKA_METADATA_DIR 97 else: 98 self._tempdir = autotemp.tempdir() 99 self._update_metadata_dir = self._tempdir.name 100 101 self._update_payloads_address = ''.join( 102 payload_url[0].rpartition('/')[0:2]) 103 # We can reuse _update_metadata_dir and _update_payloads_address 104 # for the DLC-specific install values for N-N tests, since the 105 # install and update versions will be the same. For the delta 106 # payload case, Nebraska will always use a full payload for 107 # installation and prefer a delta payload for update, so both full 108 # and delta payload metadata files can occupy the same 109 # metadata_dir. The payloads_address can be shared as well, 110 # provided all payloads have the same base URL. 111 self._install_metadata_dir = self._update_metadata_dir 112 self._install_payloads_address = self._update_payloads_address 113 114 for url in payload_url: 115 self.get_payload_properties_file(url, 116 self._update_metadata_dir, 117 **props_to_override) 118 119 def __enter__(self): 120 """So that NebraskaWrapper can be used as a Context Manager.""" 121 self.start() 122 return self 123 124 def __exit__(self, *exception_details): 125 """ 126 So that NebraskaWrapper can be used as a Context Manager. 127 128 @param exception_details: Details of exceptions happened in the 129 ContextManager. 130 131 """ 132 self.stop() 133 134 def start(self): 135 """ 136 Starts the Nebraska server. 137 138 @raise error.TestError: If fails to start the Nebraska server. 139 140 """ 141 # Any previously-existing files (port, pid and log files) will be 142 # overriden by Nebraska during bring up. 143 runtime_root = '/tmp/nebraska' 144 cmd = ['nebraska.py', '--runtime-root', runtime_root] 145 if self._log_dir: 146 cmd += ['--log-file', os.path.join(self._log_dir, 'nebraska.log')] 147 if self._update_metadata_dir: 148 cmd += ['--update-metadata', self._update_metadata_dir] 149 if self._update_payloads_address: 150 cmd += ['--update-payloads-address', self._update_payloads_address] 151 if self._install_metadata_dir: 152 cmd += ['--install-metadata', self._install_metadata_dir] 153 if self._install_payloads_address: 154 cmd += ['--install-payloads-address', 155 self._install_payloads_address] 156 157 logging.info('Starting nebraska.py with command: %s', cmd) 158 159 try: 160 self._nebraska_server = subprocess.Popen(cmd, 161 stdout=subprocess.PIPE, 162 stderr=subprocess.STDOUT) 163 164 # Wait for port file to appear. 165 port_file = os.path.join(runtime_root, 'port') 166 utils.poll_for_condition(lambda: os.path.exists(port_file), 167 timeout=5) 168 169 with open(port_file, 'r') as f: 170 self._port = int(f.read()) 171 172 # Send a health_check request to it to make sure its working. 173 requests.get('http://127.0.0.1:%d/health_check' % self._port) 174 175 except Exception as e: 176 raise error.TestError('Failed to start Nebraska %s' % e) 177 178 def stop(self): 179 """Stops the Nebraska server.""" 180 if not self._nebraska_server: 181 return 182 try: 183 self._nebraska_server.terminate() 184 stdout, _ = self._nebraska_server.communicate() 185 logging.info('Stopping nebraska.py with stdout %s', stdout) 186 self._nebraska_server.wait() 187 except subprocess.TimeoutExpired: 188 logging.error('Failed to stop Nebraska. Ignoring...') 189 finally: 190 self._nebraska_server = None 191 192 def get_update_url(self, **kwargs): 193 """ 194 Returns a URL for getting updates from this Nebraska instance. 195 196 @param kwargs: A set of key/values to form a search query to instruct 197 Nebraska to do a set of activities. See 198 nebraska.py::ResponseProperties for examples key/values. 199 """ 200 201 query = '&'.join('%s=%s' % (k, v) for k, v in kwargs.items()) 202 url = six.moves.urllib.parse.SplitResult(scheme='http', 203 netloc='127.0.0.1:%d' % 204 self._port, 205 path='/update', 206 query=query, 207 fragment='') 208 return six.moves.urllib.parse.urlunsplit(url) 209 210 def get_payload_properties_file(self, payload_url, target_dir, **kwargs): 211 """ 212 Downloads the payload properties file into a directory. 213 214 @param payload_url: The URL to the update payload file. 215 @param target_dir: The directory to download the file into. 216 @param kwargs: A dictionary of key/values that needs to be overridden on 217 the payload properties file. 218 219 """ 220 payload_props_url = payload_url + '.json' 221 _, _, file_name = payload_props_url.rpartition('/') 222 try: 223 response = json.loads(requests.get(payload_props_url).text) 224 # Override existing keys if any. 225 for k, v in six.iteritems(kwargs): 226 # Don't set default None values. We don't want to override good 227 # values to None. 228 if v is not None: 229 response[k] = v 230 with open(os.path.join(target_dir, file_name), 'w') as fp: 231 json.dump(response, fp) 232 233 except (requests.exceptions.RequestException, 234 IOError, 235 ValueError) as err: 236 raise error.TestError( 237 'Failed to get update payload properties: %s with error: %s' % 238 (payload_props_url, err)) 239 240 def update_config(self, **kwargs): 241 """ 242 Updates the current running nebraska's config. 243 244 @param kwargs: A dictionary of key/values to update the nebraska's 245 config. See platform/dev/nebraska/nebraska.py for more 246 information. 247 248 """ 249 requests.post('http://127.0.0.1:%d/update_config' % self._port, 250 json=kwargs) 251 252 def _create_nebraska_dir(self, metadata=True): 253 """ 254 Creates /usr/local/nebraska for storing the startup conf and 255 persistent metadata files. 256 257 @param metadata: True to create a subdir for metadata. 258 259 """ 260 dir_to_make = NEBRASKA_DIR 261 if metadata: 262 dir_to_make = NEBRASKA_METADATA_DIR 263 try: 264 os.makedirs(dir_to_make) 265 except OSError as e: 266 if errno.EEXIST != e.errno: 267 raise error.TestError('Failed to create %s with error: %s', 268 dir_to_make, e) 269 270 def create_startup_config(self, **kwargs): 271 """ 272 Creates a nebraska startup config file. If this file is present, nebraska 273 will start before update_engine does during system startup. 274 275 @param kwargs: A dictionary of key/values for nebraska config options. 276 See platform/dev/nebraska/nebraska.py for more info. 277 278 """ 279 conf = {} 280 if self._update_metadata_dir: 281 conf['update_metadata'] = self._update_metadata_dir 282 if self._update_payloads_address: 283 conf['update_payloads_address'] = self._update_payloads_address 284 if self._install_metadata_dir: 285 conf['install_metadata'] = self._install_metadata_dir 286 if self._install_payloads_address: 287 conf['install_payloads_address'] = self._install_payloads_address 288 289 for k, v in six.iteritems(kwargs): 290 conf[k] = v 291 292 self._create_nebraska_dir() 293 with open(NEBRASKA_CONFIG, 'w') as fp: 294 json.dump(conf, fp) 295