1# Copyright 2020 The Chromium OS Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5import json 6import logging 7import os 8import requests 9import subprocess 10import urlparse 11 12from autotest_lib.client.bin import utils 13from autotest_lib.client.common_lib import autotemp 14from autotest_lib.client.common_lib import error 15 16 17# JSON attributes used in payload properties. Look at nebraska.py for more 18# information. 19KEY_PUBLIC_KEY='public_key' 20KEY_METADATA_SIZE='metadata_size' 21KEY_SHA256='sha256_hex' 22 23 24class NebraskaWrapper(object): 25 """ 26 A wrapper around nebraska.py 27 28 This wrapper is used to start a nebraska.py service and allow the 29 update_engine to interact with it. 30 31 """ 32 33 def __init__(self, log_dir=None, payload_url=None, **props_to_override): 34 """ 35 Initializes the NebraskaWrapper module. 36 37 @param log_dir: The directory to write nebraska.log into. 38 @param payload_url: The payload that will be returned in responses for 39 update requests. This can be a single URL string 40 or a list of URLs to return multiple payload URLs 41 (such as a platform payload + DLC payloads) in the 42 responses. 43 @param props_to_override: Dictionary of key/values to use in responses 44 instead of the default values in payload_url's properties file. 45 46 """ 47 self._nebraska_server = None 48 self._port = None 49 self._log_dir = log_dir 50 51 # _update_metadata_dir is the directory for storing the json metadata 52 # files associated with the payloads. 53 # _update_payloads_address is the address of the update server where 54 # the payloads are staged. 55 # The _install variables serve the same purpose for payloads intended 56 # for DLC install requests. 57 self._update_metadata_dir = None 58 self._update_payloads_address = None 59 self._install_metadata_dir = None 60 self._install_payloads_address = None 61 62 # Create a temporary directory for the metadata and download the 63 # metadata files. 64 if payload_url: 65 # Normalize payload_url to be a list. 66 if not isinstance(payload_url, list): 67 payload_url = [payload_url] 68 69 self._update_metadata_dir = autotemp.tempdir() 70 self._update_payloads_address = ''.join( 71 payload_url[0].rpartition('/')[0:2]) 72 # We can reuse _update_metadata_dir and _update_payloads_address 73 # for the DLC-specific install values for N-N tests, since the 74 # install and update versions will be the same. For the delta 75 # payload case, Nebraska will always use a full payload for 76 # installation and prefer a delta payload for update, so both full 77 # and delta payload metadata files can occupy the same 78 # metadata_dir. The payloads_address can be shared as well, 79 # provided all payloads have the same base URL. 80 self._install_metadata_dir = self._update_metadata_dir 81 self._install_payloads_address = self._update_payloads_address 82 83 for url in payload_url: 84 self.get_payload_properties_file( 85 url, self._update_metadata_dir.name, 86 **props_to_override) 87 88 def __enter__(self): 89 """So that NebraskaWrapper can be used as a Context Manager.""" 90 self.start() 91 return self 92 93 def __exit__(self, *exception_details): 94 """ 95 So that NebraskaWrapper can be used as a Context Manager. 96 97 @param exception_details: Details of exceptions happened in the 98 ContextManager. 99 100 """ 101 self.stop() 102 103 def start(self): 104 """ 105 Starts the Nebraska server. 106 107 @raise error.TestError: If fails to start the Nebraska server. 108 109 """ 110 # Any previously-existing files (port, pid and log files) will be 111 # overriden by Nebraska during bring up. 112 runtime_root = '/tmp/nebraska' 113 cmd = ['nebraska.py', '--runtime-root', runtime_root] 114 if self._log_dir: 115 cmd += ['--log-file', os.path.join(self._log_dir, 'nebraska.log')] 116 if self._update_metadata_dir: 117 cmd += ['--update-metadata', self._update_metadata_dir.name] 118 if self._update_payloads_address: 119 cmd += ['--update-payloads-address', self._update_payloads_address] 120 if self._install_metadata_dir: 121 cmd += ['--install-metadata', self._install_metadata_dir.name] 122 if self._install_payloads_address: 123 cmd += ['--install-payloads-address', 124 self._install_payloads_address] 125 126 logging.info('Starting nebraska.py with command: %s', cmd) 127 128 try: 129 self._nebraska_server = subprocess.Popen(cmd, 130 stdout=subprocess.PIPE, 131 stderr=subprocess.STDOUT) 132 133 # Wait for port file to appear. 134 port_file = os.path.join(runtime_root, 'port') 135 utils.poll_for_condition(lambda: os.path.exists(port_file), 136 timeout=5) 137 138 with open(port_file, 'r') as f: 139 self._port = int(f.read()) 140 141 # Send a health_check request to it to make sure its working. 142 requests.get('http://127.0.0.1:%d/health_check' % self._port) 143 144 except Exception as e: 145 raise error.TestError('Failed to start Nebraska %s' % e) 146 147 def stop(self): 148 """Stops the Nebraska server.""" 149 if not self._nebraska_server: 150 return 151 try: 152 self._nebraska_server.terminate() 153 stdout, _ = self._nebraska_server.communicate() 154 logging.info('Stopping nebraska.py with stdout %s', stdout) 155 self._nebraska_server.wait() 156 except subprocess.TimeoutExpired: 157 logging.error('Failed to stop Nebraska. Ignoring...') 158 finally: 159 self._nebraska_server = None 160 161 def get_update_url(self, **kwargs): 162 """ 163 Returns a URL for getting updates from this Nebraska instance. 164 165 @param kwargs: A set of key/values to form a search query to instruct 166 Nebraska to do a set of activities. See 167 nebraska.py::ResponseProperties for examples key/values. 168 """ 169 170 query = '&'.join('%s=%s' % (k, v) for k, v in kwargs.items()) 171 url = urlparse.SplitResult(scheme='http', 172 netloc='127.0.0.1:%d' % self._port, 173 path='/update', 174 query=query, 175 fragment='') 176 return urlparse.urlunsplit(url) 177 178 def get_payload_properties_file(self, payload_url, target_dir, **kwargs): 179 """ 180 Downloads the payload properties file into a directory. 181 182 @param payload_url: The URL to the update payload file. 183 @param target_dir: The directory to download the file into. 184 @param kwargs: A dictionary of key/values that needs to be overridden on 185 the payload properties file. 186 187 """ 188 payload_props_url = payload_url + '.json' 189 _, _, file_name = payload_props_url.rpartition('/') 190 try: 191 response = json.loads(requests.get(payload_props_url).text) 192 # Override existing keys if any. 193 for k, v in kwargs.iteritems(): 194 # Don't set default None values. We don't want to override good 195 # values to None. 196 if v is not None: 197 response[k] = v 198 with open(os.path.join(target_dir, file_name), 'w') as fp: 199 json.dump(response, fp) 200 201 except (requests.exceptions.RequestException, 202 IOError, 203 ValueError) as err: 204 raise error.TestError( 205 'Failed to get update payload properties: %s with error: %s' % 206 (payload_props_url, err)) 207