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