xref: /aosp_15_r20/external/autotest/frontend/afe/json_rpc/proxy.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1# Lint as: python2, python3
2"""
3  Copyright (c) 2007 Jan-Klaas Kollhof
4
5  This file is part of jsonrpc.
6
7  jsonrpc is free software; you can redistribute it and/or modify
8  it under the terms of the GNU Lesser General Public License as published by
9  the Free Software Foundation; either version 2.1 of the License, or
10  (at your option) any later version.
11
12  This software is distributed in the hope that it will be useful,
13  but WITHOUT ANY WARRANTY; without even the implied warranty of
14  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15  GNU Lesser General Public License for more details.
16
17  You should have received a copy of the GNU Lesser General Public License
18  along with this software; if not, write to the Free Software
19  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
20"""
21
22from __future__ import absolute_import
23from __future__ import division
24from __future__ import print_function
25
26import os
27import socket
28import subprocess
29from six.moves import urllib
30import six
31from six.moves import urllib
32from autotest_lib.client.common_lib import error as exceptions
33from autotest_lib.client.common_lib import global_config
34
35from json import decoder
36
37from json import encoder as json_encoder
38json_encoder_class = json_encoder.JSONEncoder
39
40
41# Try to upgrade to the Django JSON encoder. It uses the standard json encoder
42# but can handle DateTime
43try:
44    # See http://crbug.com/418022 too see why the try except is needed here.
45    from django import conf as django_conf
46    # The serializers can't be imported if django isn't configured.
47    # Using try except here doesn't work, as test_that initializes it's own
48    # django environment (setup_django_lite_environment) which raises import
49    # errors if the django dbutils have been previously imported, as importing
50    # them leaves some state behind.
51    # This the variable name must not be undefined or empty string.
52    if os.environ.get(django_conf.ENVIRONMENT_VARIABLE, None):
53        from django.core.serializers import json as django_encoder
54        json_encoder_class = django_encoder.DjangoJSONEncoder
55except ImportError:
56    pass
57
58
59class JSONRPCException(Exception):
60    pass
61
62
63class ValidationError(JSONRPCException):
64    """Raised when the RPC is malformed."""
65    def __init__(self, error, formatted_message):
66        """Constructor.
67
68        @param error: a dict of error info like so:
69                      {error['name']: 'ErrorKind',
70                       error['message']: 'Pithy error description.',
71                       error['traceback']: 'Multi-line stack trace'}
72        @formatted_message: string representation of this exception.
73        """
74        self.problem_keys = eval(error['message'])
75        self.traceback = error['traceback']
76        super(ValidationError, self).__init__(formatted_message)
77
78
79def BuildException(error):
80    """Exception factory.
81
82    Given a dict of error info, determine which subclass of
83    JSONRPCException to build and return.  If can't determine the right one,
84    just return a JSONRPCException with a pretty-printed error string.
85
86    @param error: a dict of error info like so:
87                  {error['name']: 'ErrorKind',
88                   error['message']: 'Pithy error description.',
89                   error['traceback']: 'Multi-line stack trace'}
90    """
91    error_message = '%(name)s: %(message)s\n%(traceback)s' % error
92    for cls in JSONRPCException.__subclasses__():
93        if error['name'] == cls.__name__:
94            return cls(error, error_message)
95    for cls in (exceptions.CrosDynamicSuiteException.__subclasses__() +
96                exceptions.RPCException.__subclasses__()):
97        if error['name'] == cls.__name__:
98            return cls(error_message)
99    return JSONRPCException(error_message)
100
101
102class ServiceProxy(object):
103    def __init__(self, serviceURL, serviceName=None, headers=None):
104        """
105        @param serviceURL: The URL for the service we're proxying.
106        @param serviceName: Name of the REST endpoint to hit.
107        @param headers: Extra HTTP headers to include.
108        """
109        self.__serviceURL = serviceURL
110        self.__serviceName = serviceName
111        self.__headers = headers or {}
112
113        # TODO(pprabhu) We are reading this config value deep in the stack
114        # because we don't want to update all tools with a new command line
115        # argument. Once this has been proven to work, flip the switch -- use
116        # sso by default, and turn it off internally in the lab via
117        # shadow_config.
118        self.__use_sso_client = global_config.global_config.get_config_value(
119            'CLIENT', 'use_sso_client', type=bool, default=False)
120
121
122    def __getattr__(self, name):
123        if self.__serviceName is not None:
124            name = "%s.%s" % (self.__serviceName, name)
125        return ServiceProxy(self.__serviceURL, name, self.__headers)
126
127    def __call__(self, *args, **kwargs):
128        # Caller can pass in a minimum value of timeout to be used for urlopen
129        # call. Otherwise, the default socket timeout will be used.
130        min_rpc_timeout = kwargs.pop('min_rpc_timeout', None)
131        postdata = json_encoder_class().encode({
132                'method': self.__serviceName,
133                'params': args + (kwargs, ),
134                'id': 'jsonrpc'
135        }).encode('utf-8')
136        url_with_args = self.__serviceURL + '?' + urllib.parse.urlencode(
137                {'method': self.__serviceName})
138        if self.__use_sso_client:
139            respdata = _sso_request(url_with_args, self.__headers, postdata,
140                                    min_rpc_timeout)
141        else:
142            respdata = _raw_http_request(url_with_args, self.__headers,
143                                         postdata, min_rpc_timeout)
144
145        if isinstance(respdata, bytes):
146            respdata = respdata.decode('utf-8')
147
148        try:
149            resp = decoder.JSONDecoder().decode(respdata)
150        except ValueError:
151            raise JSONRPCException('Error decoding JSON reponse:\n' + respdata)
152        if resp['error'] is not None:
153            raise BuildException(resp['error'])
154        else:
155            return resp['result']
156
157
158def _raw_http_request(url_with_args, headers, postdata, timeout):
159    """Make a raw HTPP request.
160
161    @param url_with_args: url with the GET params formatted.
162    @headers: Any extra headers to include in the request.
163    @postdata: data for a POST request instead of a GET.
164    @timeout: timeout to use (in seconds).
165
166    @returns: the response from the http request.
167    """
168    request = urllib.request.Request(url_with_args,
169                                     data=postdata,
170                                     headers=headers)
171    default_timeout = socket.getdefaulttimeout()
172    if not default_timeout:
173        # If default timeout is None, socket will never time out.
174        return urllib.request.urlopen(request).read()
175    else:
176        return urllib.request.urlopen(
177                request,
178                timeout=max(timeout, default_timeout),
179        ).read()
180
181
182def _sso_request(url_with_args, headers, postdata, timeout):
183    """Make an HTTP request via sso_client.
184
185    @param url_with_args: url with the GET params formatted.
186    @headers: Any extra headers to include in the request.
187    @postdata: data for a POST request instead of a GET.
188    @timeout: timeout to use (in seconds).
189
190    @returns: the response from the http request.
191    """
192    headers_str = '; '.join(
193            ['%s: %s' % (k, v) for k, v in six.iteritems(headers)])
194    cmd = [
195        'sso_client',
196        '-url', url_with_args,
197    ]
198    if headers_str:
199        cmd += [
200                '-header_sep', '";"',
201                '-headers', headers_str,
202        ]
203    if postdata:
204        cmd += [
205                '-method', 'POST',
206                '-data', postdata,
207        ]
208    if timeout:
209        cmd += ['-request_timeout', str(timeout)]
210    else:
211        # sso_client has a default timeout of 5 seconds. To mimick the raw
212        # behaviour of never timing out, we force a large timeout.
213        cmd += ['-request_timeout', '3600']
214
215    try:
216        return subprocess.check_output(cmd, stderr=subprocess.STDOUT)
217    except subprocess.CalledProcessError as e:
218        if _sso_creds_error(e.output):
219            raise JSONRPCException('RPC blocked by uberproxy. Have your run '
220                                   '`prodaccess`')
221
222        raise JSONRPCException(
223                'Error (code: %s) retrieving url (%s): %s' %
224                (e.returncode, url_with_args, e.output)
225        )
226
227
228def _sso_creds_error(output):
229    return 'No user creds available' in output
230