1# 2# Copyright 2008 Google Inc. All Rights Reserved. 3# 4 5from __future__ import print_function 6 7import os 8from autotest_lib.frontend.afe import rpc_client_lib 9from autotest_lib.frontend.afe.json_rpc import proxy 10from autotest_lib.client.common_lib import global_config, utils 11 12GLOBAL_CONFIG = global_config.global_config 13DEFAULT_SERVER = 'autotest' 14AFE_RPC_PATH = '/afe/server/noauth/rpc/' 15TKO_RPC_PATH = '/new_tko/server/noauth/rpc/' 16 17 18class AuthError(Exception): 19 pass 20 21 22def get_autotest_server(web_server=None): 23 if not web_server: 24 if 'AUTOTEST_WEB' in os.environ: 25 web_server = os.environ['AUTOTEST_WEB'] 26 else: 27 web_server = GLOBAL_CONFIG.get_config_value( 28 'SERVER', 'hostname', default=DEFAULT_SERVER) 29 30 web_server = rpc_client_lib.add_protocol(web_server) 31 return web_server 32 33 34class rpc_comm(object): 35 """Shared AFE/TKO RPC class stuff""" 36 def __init__(self, web_server, rpc_path, username): 37 self.username = username 38 self.web_server = get_autotest_server(web_server) 39 try: 40 self.proxy = self._connect(rpc_path) 41 except rpc_client_lib.AuthError as s: 42 raise AuthError(s) 43 44 45 def _connect(self, rpc_path): 46 # This does not fail even if the address is wrong. 47 # We need to wait for an actual RPC to fail 48 headers = rpc_client_lib.authorization_headers(self.username, 49 self.web_server) 50 rpc_server = self.web_server + rpc_path 51 return rpc_client_lib.get_proxy(rpc_server, headers=headers) 52 53 54 def run(self, op, *args, **data): 55 if 'AUTOTEST_CLI_DEBUG' in os.environ: 56 print(self.web_server, op, args, data) 57 function = getattr(self.proxy, op) 58 result = function(*args, **data) 59 if 'AUTOTEST_CLI_DEBUG' in os.environ: 60 print('result:', result) 61 return result 62 63 64class afe_comm(rpc_comm): 65 """Handles the AFE setup and communication through RPC""" 66 def __init__(self, web_server=None, rpc_path=AFE_RPC_PATH, username=None): 67 super(afe_comm, self).__init__(web_server, rpc_path, username) 68 69 70class tko_comm(rpc_comm): 71 """Handles the TKO setup and communication through RPC""" 72 def __init__(self, web_server=None, rpc_path=TKO_RPC_PATH, username=None): 73 super(tko_comm, self).__init__(web_server, rpc_path, username) 74