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 logging 6import os 7import requests 8import subprocess 9 10from autotest_lib.client.bin import utils 11from autotest_lib.client.common_lib import error 12 13 14# JSON attributes used in payload properties. Look at nebraska.py for more 15# information. 16KEY_PUBLIC_KEY='public_key' 17KEY_METADATA_SIZE='metadata_size' 18KEY_SHA256='sha256_hex' 19 20 21class NebraskaWrapper(object): 22 """ 23 A wrapper around nebraska.py 24 25 This wrapper is used to start a nebraska.py service and allow the 26 update_engine to interact with it. 27 28 """ 29 30 def __init__(self, log_dir=None, update_metadata_dir=None, 31 update_payloads_address=None): 32 """ 33 Initializes the NebraskaWrapper module. 34 35 @param log_dir: The directory to write nebraska.log into. 36 @param update_metadata_dir: The directory containing payload properties 37 files. Look at nebraska.py 38 @param update_payloads_address: The base URL for the update payload. 39 40 """ 41 self._nebraska_server = None 42 self._port = None 43 self._log_dir = log_dir 44 self._update_metadata_dir = update_metadata_dir 45 self._update_payloads_address = update_payloads_address 46 47 def __enter__(self): 48 """So that NebraskaWrapper can be used as a Context Manager.""" 49 self.start() 50 return self 51 52 def __exit__(self, *exception_details): 53 """ 54 So that NebraskaWrapper can be used as a Context Manager. 55 56 @param exception_details: Details of exceptions happened in the 57 ContextManager. 58 59 """ 60 self.stop() 61 62 def start(self): 63 """ 64 Starts the Nebraska server. 65 66 @raise error.TestError: If fails to start the Nebraska server. 67 68 """ 69 # Any previously-existing files (port, pid and log files) will be 70 # overriden by Nebraska during bring up. 71 runtime_root = '/tmp/nebraska' 72 cmd = ['nebraska.py', '--runtime-root', runtime_root] 73 if self._log_dir: 74 cmd += ['--log-file', os.path.join(self._log_dir, 'nebraska.log')] 75 if self._update_metadata_dir: 76 cmd += ['--update-metadata', self._update_metadata_dir] 77 if self._update_payloads_address: 78 cmd += ['--update-payloads-address', self._update_payloads_address] 79 80 logging.info('Starting nebraska.py with command: %s', cmd) 81 82 try: 83 self._nebraska_server = subprocess.Popen(cmd, 84 stdout=subprocess.PIPE, 85 stderr=subprocess.STDOUT) 86 87 # Wait for port file to appear. 88 port_file = os.path.join(runtime_root, 'port') 89 utils.poll_for_condition(lambda: os.path.exists(port_file), 90 timeout=5) 91 92 with open(port_file, 'r') as f: 93 self._port = int(f.read()) 94 95 # Send a health_check request to it to make sure its working. 96 requests.get('http://127.0.0.1:%d/health_check' % self._port) 97 98 except Exception as e: 99 raise error.TestError('Failed to start Nebraska %s' % e) 100 101 def stop(self): 102 """Stops the Nebraska server.""" 103 if not self._nebraska_server: 104 return 105 try: 106 self._nebraska_server.terminate() 107 stdout, _ = self._nebraska_server.communicate() 108 logging.info('Stopping nebraska.py with stdout %s', stdout) 109 self._nebraska_server.wait() 110 except subprocess.TimeoutExpired: 111 logging.error('Failed to stop Nebraska. Ignoring...') 112 finally: 113 self._nebraska_server = None 114 115 def get_port(self): 116 """Returns the port which Nebraska is running.""" 117 return self._port 118 119 def get_update_url(self, **kwargs): 120 """ 121 Returns a URL for getting updates from this Nebraska instance. 122 123 @param kwargs: A set of key/values to form a search query to instruct 124 Nebraska to do a set of activities. See 125 nebraska.py::ResponseProperties for examples key/values. 126 """ 127 query = '&'.join('%s=%s' % (k, v) for k, v in kwargs.items()) 128 return 'http://127.0.0.1:%d/update?%s' % (self._port, query) 129