1# Lint as: python2, python3 2# Copyright 2022 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 json 11import logging 12import os 13import requests 14import six 15 16from autotest_lib.client.common_lib import autotemp 17from autotest_lib.client.common_lib import error 18from autotest_lib.client.cros.update_engine import nebraska_wrapper 19 20 21class NebraskaService: 22 """ 23 Remotely sets up nebraska on the DUT. 24 25 This service is different from 26 `autotest_lib.client.cros.update_engine.nebraska_wrapper.NebraskaWrapper` in 27 that it is used by server-only tests to remotely launch nebraska on the DUT. 28 29 """ 30 31 def __init__(self, test, host, payload_url=None, **props_to_override): 32 """ 33 Initializes the NebraskaService. 34 35 @param test: Instance of the test using the service. 36 @param host: The DUT we will be running on. 37 @param payload_url: The payload that will be returned in responses for 38 update requests. This can be a single URL string or a list of URLs 39 to return multiple payload URLs (such as a platform payload + DLC 40 payloads) in the responses. 41 @param props_to_override: Dictionary of key/values to use in responses 42 instead of the default values in payload_url's properties file. 43 44 """ 45 self._host = host 46 self._test = test 47 48 # _update_metadata_dir is the directory for storing the json metadata 49 # files associated with the payloads. 50 # _update_payloads_address is the address of the update server where 51 # the payloads are staged. 52 self._update_metadata_dir = None 53 self._update_payloads_address = None 54 55 if payload_url: 56 # Normalize payload_url to be a list. 57 if not isinstance(payload_url, list): 58 payload_url = [payload_url] 59 60 self._update_metadata_dir = self._host.get_tmp_dir() 61 self._update_payloads_address = ''.join( 62 payload_url[0].rpartition('/')[0:2]) 63 64 # Download the metadata files and save them in a tempdir for general 65 # use. 66 for url in payload_url: 67 self.get_payload_properties_file(url, 68 self._update_metadata_dir, 69 **props_to_override) 70 71 def get_payload_properties_file(self, payload_url, target_dir, **kwargs): 72 """ 73 Downloads the payload properties file into a directory on the DUT. 74 75 @param payload_url: The URL to the update payload file. 76 @param target_dir: The directory on the DUT to download the file into. 77 @param kwargs: A dictionary of key/values that needs to be overridden on 78 the payload properties file. 79 80 """ 81 payload_props_url = payload_url + '.json' 82 _, _, file_name = payload_props_url.rpartition('/') 83 try: 84 response = json.loads(requests.get(payload_props_url).text) 85 # Override existing keys if any. 86 for k, v in six.iteritems(kwargs): 87 # Don't set default None values. We don't want to override good 88 # values to None. 89 if v is not None: 90 response[k] = v 91 self._write_remote_file(os.path.join(target_dir, file_name), 92 json.dumps(response)) 93 94 except (requests.exceptions.RequestException, IOError, 95 ValueError) as err: 96 raise error.TestError( 97 'Failed to get update payload properties: %s with error: %s' 98 % (payload_props_url, err)) 99 100 def start(self, **kwargs): 101 """Launch nebraska on DUT.""" 102 # Generate nebraska configuration. 103 self._write_remote_file( 104 nebraska_wrapper.NEBRASKA_CONFIG, 105 json.dumps(self._create_startup_config(**kwargs)), 106 ) 107 logging.info('Start nebraska service') 108 self._host.upstart_restart('nebraska') 109 self._host.wait_for_service('nebraska') 110 111 def stop(self): 112 """Stop Nebraska service.""" 113 logging.info('Stop nebraska service') 114 self._host.upstart_stop('nebraska') 115 self._host.run('rm', args=('-f', nebraska_wrapper.NEBRASKA_CONFIG)) 116 117 def _create_startup_config(self, **kwargs): 118 """ 119 Creates a nebraska startup config file. If this file is present, nebraska 120 can be started by upstart. 121 122 @param kwargs: A dictionary of key/values for nebraska config options. 123 See platform/dev/nebraska/nebraska.py for more info. 124 125 @return: A dictionary of nebraska config options. 126 127 """ 128 conf = {} 129 if self._update_metadata_dir: 130 conf['update_metadata'] = self._update_metadata_dir 131 if self._update_payloads_address: 132 conf['update_payloads_address'] = self._update_payloads_address 133 134 for k, v in six.iteritems(kwargs): 135 conf[k] = v 136 return conf 137 138 def _create_remote_dir(self, remote_dir, owner=None): 139 """ 140 Create directory on DUT. 141 142 @param remote_dir: The directory to create. 143 @param owner: Set owner of the remote directory. 144 145 """ 146 permission = '1777' 147 if owner: 148 permission = '1770' 149 self._host.run(['mkdir', '-p', '-m', permission, remote_dir]) 150 if owner: 151 self._host.run('chown', args=(owner, remote_dir)) 152 153 def _write_remote_file(self, 154 filepath, 155 content, 156 permission=None, 157 owner=None): 158 """ 159 Write content to filepath on DUT. 160 161 @param permission: set permission to 0xxx octal number of remote file. 162 @param owner: set owner of remote file. 163 164 """ 165 tmpdir = autotemp.tempdir(unique_id='minios') 166 tmp_path = os.path.join(tmpdir.name, os.path.basename(filepath)) 167 with open(tmp_path, 'w') as f: 168 f.write(content) 169 if permission is not None: 170 os.chmod(tmp_path, permission) 171 self._create_remote_dir(os.path.dirname(filepath), owner) 172 self._host.send_file(tmp_path, filepath, delete_dest=True) 173 if owner is not None: 174 self._host.run('chown', args=(owner, filepath)) 175 tmpdir.clean() 176