1# Lint as: python2, python3 2# Copyright 2021 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"""Communication with the TLS FakeOmaha Service.""" 6 7import logging 8 9import common 10 11from autotest_lib.server.hosts.tls_client import autotest_common_pb2 12from autotest_lib.client.common_lib import error 13 14PAYLOAD_TYPE = { 15 'TYPE_UNSPECIFIED': 16 autotest_common_pb2.FakeOmaha.Payload.TYPE_UNSPECIFIED, 17 'FULL': autotest_common_pb2.FakeOmaha.Payload.FULL, 18 'DELTA': autotest_common_pb2.FakeOmaha.Payload.DELTA 19} 20 21 22class TLSFakeOmaha(): 23 """Object for sending commands to a host, and getting the response.""" 24 25 def __init__(self, tlsconnection): 26 """Configure the grpc channel.""" 27 if tlsconnection.alive: 28 self.stub = tlsconnection.stub 29 else: 30 raise error.TLSConnectionError( 31 "TLS connection is not alive when try to creating" 32 " FakeOmaha client.") 33 34 self.tlsconnection = tlsconnection 35 self.fo_name = None 36 37 def _make_payloads(self, payloads): 38 """Serialize and return the list of payloads.""" 39 serialized_payloads = [] 40 for payload in payloads: 41 serialized_payloads.append( 42 autotest_common_pb2.FakeOmaha.Payload( 43 id=payload['payload_id'], 44 type=PAYLOAD_TYPE[payload['payload_type']])) 45 46 return serialized_payloads 47 48 def start_omaha(self, 49 hostname, 50 target_build, 51 payloads, 52 exposed_via_proxy=False, 53 critical_update=False, 54 return_noupdate_starting=0): 55 """Serialize and send the cmd to the TLS service. 56 57 @param hostname: hostname of dut. Normally 'hostname' or 'self.hostname' 58 @param target_build: full target build for the update. Example: 59 60 @param payloads: list of the payloads in the format: 61 [{'payload_id': <id>, 'payload_type': <type>}] 62 example: 63 [{'payload_id': 'ROOTFS', 'payload_type': 'FULL'},] 64 @param exposed_via_proxy: bool indicates that the fake Omaha service is 65 exposed to a DUT via a proxy server, instead of exposing to the DUT 66 directly. 67 @param critical_update:bool instructs the fake Omaha created that the 68 update is critical. 69 @param return_noupdate_starting: int indicates from which update check 70 to start returning noupdate. 71 72 @returns: the omaha_url 73 """ 74 payload = self._make_payloads(payloads) 75 76 target_build = autotest_common_pb2.ChromeOsImage( 77 gs_path_prefix=target_build) 78 fake_omaha = autotest_common_pb2.FakeOmaha( 79 dut=hostname, 80 target_build=target_build, 81 payloads=payload, 82 exposed_via_proxy=exposed_via_proxy, 83 critical_update=critical_update, 84 return_noupdate_starting=return_noupdate_starting) 85 86 req = autotest_common_pb2.CreateFakeOmahaRequest(fake_omaha=fake_omaha) 87 88 try: 89 result = self.stub.CreateFakeOmaha(req) 90 self.fo_name = result.name 91 return result.omaha_url 92 except Exception as e: 93 logging.error("TLS FakeOmaha Debug String: %s", 94 e.debug_error_string()) 95 raise error.TestError( 96 "Could not start FakeOmaha Server because %s", e.details()) 97 98 def stop_omaha(self): 99 """Delete the running Omaha Service.""" 100 if not self.fo_name: 101 raise error.TestWarn( 102 "No FakeOmaha name specified, has it been started?") 103 req = autotest_common_pb2.DeleteFakeOmahaRequest(name=self.fo_name) 104 try: 105 self.stub.DeleteFakeOmaha(req) 106 except Exception as e: 107 raise error.TestWarn("Could not stop FakeOmaha Server because %s", 108 e.details()) 109