mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 1 | # |
| 2 | # Copyright 2008 Google Inc. All Rights Reserved. |
| 3 | # |
| 4 | |
Gregory Nisbet | 8d60b1c | 2020-07-14 16:26:46 -0700 | [diff] [blame] | 5 | from __future__ import print_function |
| 6 | |
jamesren | 1a2914a | 2010-02-12 00:44:31 +0000 | [diff] [blame] | 7 | import os |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 8 | from autotest_lib.frontend.afe import rpc_client_lib |
| 9 | from autotest_lib.frontend.afe.json_rpc import proxy |
showard | f4a6899 | 2010-02-03 20:29:59 +0000 | [diff] [blame] | 10 | from autotest_lib.client.common_lib import global_config, utils |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 11 | |
mbligh | 37eceaa | 2008-12-15 22:56:37 +0000 | [diff] [blame] | 12 | GLOBAL_CONFIG = global_config.global_config |
| 13 | DEFAULT_SERVER = 'autotest' |
Scott Zawalski | 347aaf4 | 2012-04-03 16:33:00 -0400 | [diff] [blame] | 14 | AFE_RPC_PATH = '/afe/server/noauth/rpc/' |
| 15 | TKO_RPC_PATH = '/new_tko/server/noauth/rpc/' |
mbligh | 37eceaa | 2008-12-15 22:56:37 +0000 | [diff] [blame] | 16 | |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 17 | |
mbligh | b68405d | 2010-03-11 18:32:39 +0000 | [diff] [blame] | 18 | class AuthError(Exception): |
| 19 | pass |
| 20 | |
| 21 | |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 22 | def 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: |
Prathmesh Prabhu | 02163aa | 2017-12-12 19:13:21 -0800 | [diff] [blame] | 27 | web_server = GLOBAL_CONFIG.get_config_value( |
mbligh | 37eceaa | 2008-12-15 22:56:37 +0000 | [diff] [blame] | 28 | 'SERVER', 'hostname', default=DEFAULT_SERVER) |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 29 | |
Prathmesh Prabhu | 2892ff6 | 2017-12-19 10:21:31 -0800 | [diff] [blame] | 30 | web_server = rpc_client_lib.add_protocol(web_server) |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 31 | return web_server |
| 32 | |
| 33 | |
mbligh | 759b6d0 | 2008-10-21 16:28:18 +0000 | [diff] [blame] | 34 | class rpc_comm(object): |
| 35 | """Shared AFE/TKO RPC class stuff""" |
mbligh | 33a8a7e | 2008-11-18 14:57:11 +0000 | [diff] [blame] | 36 | def __init__(self, web_server, rpc_path, username): |
| 37 | self.username = username |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 38 | self.web_server = get_autotest_server(web_server) |
mbligh | b68405d | 2010-03-11 18:32:39 +0000 | [diff] [blame] | 39 | try: |
| 40 | self.proxy = self._connect(rpc_path) |
Gregory Nisbet | 8d60b1c | 2020-07-14 16:26:46 -0700 | [diff] [blame] | 41 | except rpc_client_lib.AuthError as s: |
mbligh | b68405d | 2010-03-11 18:32:39 +0000 | [diff] [blame] | 42 | raise AuthError(s) |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 43 | |
mbligh | 11efd23 | 2008-11-27 00:20:46 +0000 | [diff] [blame] | 44 | |
showard | 5557957 | 2008-09-19 00:44:35 +0000 | [diff] [blame] | 45 | def _connect(self, rpc_path): |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 46 | # This does not fail even if the address is wrong. |
| 47 | # We need to wait for an actual RPC to fail |
jamesren | 1a2914a | 2010-02-12 00:44:31 +0000 | [diff] [blame] | 48 | headers = rpc_client_lib.authorization_headers(self.username, |
| 49 | self.web_server) |
showard | 5557957 | 2008-09-19 00:44:35 +0000 | [diff] [blame] | 50 | rpc_server = self.web_server + rpc_path |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 51 | return rpc_client_lib.get_proxy(rpc_server, headers=headers) |
| 52 | |
| 53 | |
showard | 5557957 | 2008-09-19 00:44:35 +0000 | [diff] [blame] | 54 | def run(self, op, *args, **data): |
mbligh | e249072 | 2008-12-22 14:55:50 +0000 | [diff] [blame] | 55 | if 'AUTOTEST_CLI_DEBUG' in os.environ: |
Gregory Nisbet | 8d60b1c | 2020-07-14 16:26:46 -0700 | [diff] [blame] | 56 | print(self.web_server, op, args, data) |
mbligh | be630eb | 2008-08-01 16:41:48 +0000 | [diff] [blame] | 57 | function = getattr(self.proxy, op) |
mbligh | 9196763 | 2009-03-16 21:45:52 +0000 | [diff] [blame] | 58 | result = function(*args, **data) |
| 59 | if 'AUTOTEST_CLI_DEBUG' in os.environ: |
Gregory Nisbet | 8d60b1c | 2020-07-14 16:26:46 -0700 | [diff] [blame] | 60 | print('result:', result) |
mbligh | 9196763 | 2009-03-16 21:45:52 +0000 | [diff] [blame] | 61 | return result |
mbligh | 759b6d0 | 2008-10-21 16:28:18 +0000 | [diff] [blame] | 62 | |
| 63 | |
| 64 | class afe_comm(rpc_comm): |
| 65 | """Handles the AFE setup and communication through RPC""" |
mbligh | 33a8a7e | 2008-11-18 14:57:11 +0000 | [diff] [blame] | 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) |
mbligh | 759b6d0 | 2008-10-21 16:28:18 +0000 | [diff] [blame] | 68 | |
| 69 | |
| 70 | class tko_comm(rpc_comm): |
| 71 | """Handles the TKO setup and communication through RPC""" |
mbligh | 33a8a7e | 2008-11-18 14:57:11 +0000 | [diff] [blame] | 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) |