| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # |
| from __future__ import absolute_import, print_function, unicode_literals |
| import base64 |
| from collections import namedtuple |
| import datetime |
| from io import BytesIO |
| import errno |
| import os |
| import os.path |
| import socket |
| import sys |
| |
| import six |
| from six.moves import http_client, urllib |
| from thrift.transport.TTransport import TTransportBase |
| |
| from impala_shell.cookie_util import ( |
| get_all_cookies, |
| get_all_matching_cookies, |
| get_cookie_expiry, |
| ) |
| from impala_shell.shell_exceptions import AuthenticationException, HttpError |
| |
| # Declare namedtuple for Cookie with named fields - cookie and expiry_time |
| Cookie = namedtuple('Cookie', ['cookie', 'expiry_time']) |
| |
| |
| # This was taken from THttpClient.py in Thrift to allow making changes Impala needs. |
| # The current changes that have been applied: |
| # - Added logic for the 'Expect: 100-continue' header on large requests |
| # - If an error code is received back in flush(), an exception is thrown. |
| # Note there is a copy of this code in Impyla. |
| class ImpalaHttpClient(TTransportBase): |
| """Http implementation of TTransport base.""" |
| |
| # When sending requests larger than this size, include the 'Expect: 100-continue' header |
| # to indicate to the server to validate the request before reading the contents. This |
| # value was chosen to match curl's behavior. See Section 8.2.3 of RFC2616. |
| MIN_REQUEST_SIZE_FOR_EXPECT = 1024 |
| |
| def __init__(self, uri_or_host, ssl_context=None, http_cookie_names=None, |
| socket_timeout_s=None, verbose=False, reuse_connection=True, |
| connect_timeout_ms=None): |
| """To properly authenticate against an HTTPS server, provide an ssl_context created |
| with ssl.create_default_context() to validate the server certificate. |
| |
| http_cookie_names is used to specify a comma-separated list of possible cookie names |
| used for cookie-based authentication or session management. If a cookie with one of |
| these names is returned in an http response by the server or an intermediate proxy |
| then it will be included in each subsequent request for the same connection. If it |
| is set as wildcards, all cookies in an http response will be preserved. |
| |
| If reuse_connection is set to True, the underlying HTTP connection will be reused |
| for multiple requests; it will retry establishing a connection on socket error in case |
| the server closed it while idle. If set to False, the connection will be closed and |
| reopened for each request. |
| """ |
| parsed = urllib.parse.urlparse(uri_or_host) |
| self.scheme = parsed.scheme |
| assert self.scheme in ('http', 'https') |
| if self.scheme == 'http': |
| self.port = parsed.port or http_client.HTTP_PORT |
| elif self.scheme == 'https': |
| self.port = parsed.port or http_client.HTTPS_PORT |
| self.context = ssl_context |
| self.host = parsed.hostname |
| self.path = parsed.path |
| if parsed.query: |
| self.path += '?%s' % parsed.query |
| try: |
| proxy = urllib.request.getproxies()[self.scheme] |
| except KeyError: |
| proxy = None |
| else: |
| if urllib.request.proxy_bypass(self.host): |
| proxy = None |
| if proxy: |
| parsed = urllib.parse.urlparse(proxy) |
| self.realhost = self.host |
| self.realport = self.port |
| self.host = parsed.hostname |
| self.port = parsed.port |
| self.proxy_auth = self.basic_proxy_auth_header(parsed) |
| else: |
| self.realhost = self.realport = self.proxy_auth = None |
| if (not http_cookie_names) or (str(http_cookie_names).strip() == ""): |
| self.__preserve_all_cookies = False |
| self.__http_cookie_dict = None |
| self.__auth_cookie_names = None |
| elif str(http_cookie_names).strip() == "*": |
| self.__preserve_all_cookies = True |
| self.__http_cookie_dict = dict() |
| self.__auth_cookie_names = set() |
| else: |
| self.__preserve_all_cookies = False |
| # Build a dictionary that maps cookie name to namedtuple. |
| cookie_names = http_cookie_names.split(',') |
| self.__http_cookie_dict = \ |
| {cn: Cookie(cookie=None, expiry_time=None) for cn in cookie_names} |
| # Store the auth cookie names in __auth_cookie_names. |
| # Assume auth cookie names end with ".auth". |
| self.__auth_cookie_names = {cn for cn in cookie_names if cn.endswith(".auth")} |
| # Set __are_matching_cookies_found as True if matching cookies are found in response. |
| self.__are_matching_cookies_found = False |
| self.__wbuf = BytesIO() |
| self.__http = None |
| self.__http_response = None |
| self.__connect_timeout_s = connect_timeout_ms / 1000.0 if connect_timeout_ms else None |
| self.__socket_timeout_s = socket_timeout_s |
| # __custom_headers is used to store HTTP headers which are generated in runtime for |
| # new request. |
| self.__custom_headers = None |
| self.__get_custom_headers_func = None |
| self.__basic_auth = None |
| self.__kerb_service = None |
| self.__add_custom_headers_funcs = [] |
| self.__verbose = verbose |
| self.__reuse_connection = reuse_connection |
| |
| @staticmethod |
| def basic_proxy_auth_header(proxy): |
| if proxy is None or not proxy.username: |
| return None |
| ap = "%s:%s" % (urllib.parse.unquote(proxy.username), |
| urllib.parse.unquote(proxy.password)) |
| cr = base64.b64encode(ap).strip() |
| return "Basic " + cr |
| |
| def using_proxy(self): |
| return self.realhost is not None |
| |
| def open(self): |
| if self.scheme == 'http': |
| self.__http = http_client.HTTPConnection(self.host, self.port, |
| timeout=self.__connect_timeout_s) |
| elif self.scheme == 'https': |
| self.__http = http_client.HTTPSConnection(self.host, self.port, |
| timeout=self.__connect_timeout_s, |
| context=self.context) |
| if self.using_proxy(): |
| self.__http.set_tunnel(self.realhost, self.realport, |
| {"Proxy-Authorization": self.proxy_auth}) |
| # Initiate connection with configured timeout. |
| self.__http.connect() |
| # Set socket_timeout_s for the remainder of the session. |
| self.__http.sock.settimeout(self.__socket_timeout_s) |
| |
| def close(self): |
| self.__http.close() |
| self.__http = None |
| self.__http_response = None |
| |
| def isOpen(self): |
| return self.__http is not None |
| |
| def getCustomHeadersWithBasicAuth(self, cookie_header, has_auth_cookie): |
| custom_headers = {} |
| if cookie_header: |
| # Add cookies to HTTP header. |
| custom_headers['Cookie'] = cookie_header |
| # Add the 'Authorization' header to request even if the auth cookie is |
| # present to avoid a round trip in case the cookie is expired when server |
| # receive the request. Since the 'auth' value is calculated once, so it |
| # won't cause a performance issue. |
| custom_headers['Authorization'] = "Basic " + self.__basic_auth |
| return custom_headers |
| |
| def getCustomHeadersWithBearerAuth(self, cookie_header, has_auth_cookie): |
| custom_headers = {} |
| if cookie_header: |
| # Add cookies to HTTP header. |
| custom_headers['Cookie'] = cookie_header |
| # Add the 'Authorization' header to request even if the auth cookie is |
| # present to avoid a round trip in case the cookie is expired when server |
| # receive the request. Since the 'auth' value is calculated once, so it |
| # won't cause a performance issue. |
| custom_headers['Authorization'] = "Bearer {0}".format(self.__bearer_token) |
| return custom_headers |
| |
| def getCustomHeadersWithNegotiateAuth(self, cookie_header, has_auth_cookie): |
| import kerberos |
| custom_headers = {} |
| if cookie_header: |
| # Add cookies to HTTP header. |
| custom_headers['Cookie'] = cookie_header |
| # For GSSAPI over http we need to dynamically generate custom request headers. |
| if not has_auth_cookie: |
| try: |
| _, krb_context = kerberos.authGSSClientInit(self.__kerb_service) |
| kerberos.authGSSClientStep(krb_context, "") |
| negotiate_details = kerberos.authGSSClientResponse(krb_context) |
| custom_headers['Authorization'] = "Negotiate " + negotiate_details |
| except kerberos.GSSError: |
| raise AuthenticationException("Kerberos authentication failure.") |
| return custom_headers |
| |
| def getCustomHeadersWithoutAuth(self, cookie_header, has_auth_cookie): |
| custom_headers = {} |
| if cookie_header: |
| # Add cookies to HTTP header. |
| custom_headers['Cookie'] = cookie_header |
| return custom_headers |
| |
| # Set function to generate customized HTTP headers for Ldap authorization. |
| def setLdapAuth(self, basic_auth): |
| # auth mechanism: LDAP |
| self.__basic_auth = basic_auth |
| self.__get_custom_headers_func = self.getCustomHeadersWithBasicAuth |
| |
| # Set function to generate customized HTTP headers for JWT authorization. |
| def setJwtAuth(self, jwt): |
| # auth mechanism: JWT |
| self.__bearer_token = jwt |
| self.__get_custom_headers_func = self.getCustomHeadersWithBearerAuth |
| |
| # Set function to generate customized HTTP headers for OAuth authorization. |
| def setOAuthAuth(self, oauth): |
| # auth mechanism: Oauth |
| self.__bearer_token = oauth |
| self.__get_custom_headers_func = self.getCustomHeadersWithBearerAuth |
| |
| # Set function to generate customized HTTP headers for Kerberos authorization. |
| def setKerberosAuth(self, kerb_service): |
| # auth mechanism: GSSAPI |
| self.__kerb_service = kerb_service |
| self.__get_custom_headers_func = self.getCustomHeadersWithNegotiateAuth |
| |
| # Set function to generate customized HTTP headers without authorization. |
| def setNoneAuth(self): |
| # auth mechanism: None |
| self.__get_custom_headers_func = self.getCustomHeadersWithoutAuth |
| |
| # Whenever http(s) calls are made to the backend impala, each function |
| # added through this method will be called. Thus, arbitrary custom |
| # headers can be set on each request. |
| # parameters: |
| # funcs - tuple of functions where each takes no arguments and returns |
| # a dict of http headers |
| # Note: if the custom function returns a http header with a name that |
| # does not start with "X-" or "x-", it will cause an error to be thrown |
| def addCustomHeaderFunc(self, *funcs): |
| if funcs is None: |
| return |
| |
| for f in funcs: |
| self.__add_custom_headers_funcs.append(f) |
| |
| # Update HTTP headers based on the saved cookies and auth mechanism. |
| def refreshCustomHeaders(self): |
| self.__custom_headers = {} |
| |
| if self.__get_custom_headers_func: |
| cookie_header, has_auth_cookie = self.getHttpCookieHeaderForRequest() |
| self.__custom_headers = \ |
| self.__get_custom_headers_func(cookie_header, has_auth_cookie) |
| |
| for f in self.__add_custom_headers_funcs: |
| headers = f() |
| if headers is not None: |
| for key in headers: |
| assert key[0:2].lower() == "x-", \ |
| "header '{0}' is not valid, all custom headers must start with "\ |
| "'X-' or 'x-'".format(key) |
| assert key not in self.__custom_headers, \ |
| "header '{0}' already exists in custom headers dictionary".format(key) |
| self.__custom_headers[key] = headers[key] |
| |
| # Return first value as a cookie list for Cookie header. It's a list of name-value |
| # pairs in the form of <cookie-name>=<cookie-value>. Pairs in the list are separated by |
| # a semicolon and a space ('; '). |
| # Return second value as True if the cookie list contains auth cookie. |
| def getHttpCookieHeaderForRequest(self): |
| if (self.__http_cookie_dict is None) or not self.__are_matching_cookies_found: |
| return None, False |
| cookie_headers = [] |
| has_auth_cookie = False |
| for cn, c_tuple in self.__http_cookie_dict.items(): |
| if c_tuple.cookie: |
| if c_tuple.expiry_time and c_tuple.expiry_time <= datetime.datetime.now(): |
| self.__http_cookie_dict[cn] = Cookie(cookie=None, expiry_time=None) |
| else: |
| cookie_header = c_tuple.cookie.output(attrs=['value'], header='').strip() |
| cookie_headers.append(cookie_header) |
| if not has_auth_cookie and self.__auth_cookie_names \ |
| and cn in self.__auth_cookie_names: |
| has_auth_cookie = True |
| if not cookie_headers: |
| self.__are_matching_cookies_found = False |
| return None, False |
| else: |
| return '; '.join(cookie_headers), has_auth_cookie |
| |
| # Extract cookies from response and save those cookies for which the cookie names |
| # are in the cookie name list specified in the __init__(). |
| def extractHttpCookiesFromResponse(self): |
| if self.__preserve_all_cookies: |
| matching_cookies = get_all_cookies(self.path, self.headers) |
| elif self.__http_cookie_dict: |
| matching_cookies = get_all_matching_cookies( |
| self.__http_cookie_dict.keys(), self.path, self.headers) |
| else: |
| matching_cookies = None |
| |
| if matching_cookies: |
| if self.__verbose: |
| names = sorted([c.key for c in matching_cookies]) |
| print("Preserving cookies: " + ', '.join(names), file=sys.stderr) |
| sys.stderr.flush() |
| |
| self.__are_matching_cookies_found = True |
| for c in matching_cookies: |
| self.__http_cookie_dict[c.key] = Cookie(c, get_cookie_expiry(c)) |
| if c.key.endswith(".auth"): |
| self.__auth_cookie_names.add(c.key) |
| |
| # Return True if there are any saved cookies which are sent in previous request. |
| def areHttpCookiesSaved(self): |
| return self.__are_matching_cookies_found |
| |
| # Clean all saved cookies. |
| def cleanHttpCookies(self): |
| if (self.__http_cookie_dict is not None) and self.__are_matching_cookies_found: |
| self.__are_matching_cookies_found = False |
| self.__http_cookie_dict = \ |
| {cn: Cookie(cookie=None, expiry_time=None) for cn in self.__http_cookie_dict} |
| |
| def read(self, sz): |
| return self.__http_response.read(sz) |
| |
| def readBody(self): |
| return self.__http_response.read() |
| |
| def write(self, buf): |
| self.__wbuf.write(buf) |
| |
| def flush(self): |
| # Send HTTP request headers. This is repeatable, so if there's a connection error |
| # like when the connection has been closed it's safe to retry. |
| def sendRequestHeaders(data_len): |
| if not self.__reuse_connection and self.isOpen(): |
| self.close() |
| if not self.isOpen(): |
| self.open() |
| |
| # HTTP request |
| if self.using_proxy() and self.scheme == "http": |
| # need full URL of real host for HTTP proxy here (HTTPS uses CONNECT tunnel) |
| self.__http.putrequest('POST', "http://%s:%s%s" % |
| (self.realhost, self.realport, self.path)) |
| else: |
| self.__http.putrequest('POST', self.path) |
| |
| # Write headers |
| self.__http.putheader('Content-Type', 'application/x-thrift') |
| self.__http.putheader('Content-Length', str(data_len)) |
| if data_len > ImpalaHttpClient.MIN_REQUEST_SIZE_FOR_EXPECT: |
| # Add the 'Expect' header to large requests. Note that we do not explicitly wait |
| # for the '100 continue' response before sending the data - HTTPConnection simply |
| # ignores these types of responses, but we'll get the right behavior anyways. |
| self.__http.putheader("Expect", "100-continue") |
| if self.using_proxy() and self.scheme == "http" and self.proxy_auth is not None: |
| self.__http.putheader("Proxy-Authorization", self.proxy_auth) |
| |
| self.refreshCustomHeaders() |
| if not self.__custom_headers or 'User-Agent' not in self.__custom_headers: |
| user_agent = 'Python/ImpalaHttpClient' |
| script = os.path.basename(sys.argv[0]) |
| if script: |
| user_agent = '%s (%s)' % (user_agent, urllib.parse.quote(script)) |
| self.__http.putheader('User-Agent', user_agent) |
| |
| if self.__custom_headers: |
| for key, val in six.iteritems(self.__custom_headers): |
| self.__http.putheader(key, val) |
| |
| self.__http.endheaders() |
| |
| # Complete the request by sending data and getting the response. Return True if the |
| # client should retry this method due to a '401 Unauthorized' response. |
| def sendDataRecvResp(data): |
| # Write payload |
| self.__http.send(data) |
| |
| # Get reply to flush the request |
| self.__http_response = self.__http.getresponse() |
| self.code = self.__http_response.status |
| self.message = self.__http_response.reason |
| self.headers = self.__http_response.msg |
| # A '401 Unauthorized' response might mean that we tried cookie-based |
| # authentication with one or more expired cookies. |
| # Delete the cookies and try again. |
| if self.code == 401 and self.areHttpCookiesSaved(): |
| self.cleanHttpCookies() |
| return True |
| else: |
| self.extractHttpCookiesFromResponse() |
| return False |
| |
| # Pull data out of buffer |
| data = self.__wbuf.getvalue() |
| data_len = len(data) |
| self.__wbuf = BytesIO() |
| |
| # Send the request headers, retrying once if the connection was closed. Sending |
| # headers is the earliest point that http_client allows us to detect a closed |
| # connection. |
| retry_because_disconnected = False |
| try: |
| sendRequestHeaders(data_len) |
| except http_client.CannotSendRequest: |
| retry_because_disconnected = True |
| except socket.error as e: |
| if e.errno not in [errno.EPIPE, errno.ECONNRESET]: |
| raise |
| retry_because_disconnected = True |
| |
| if retry_because_disconnected and self.__reuse_connection: |
| if self.__verbose: |
| print('Connection closed, reconnecting...', file=sys.stderr) |
| # The underlying socket is broken. Try to reconnect and then retry. |
| self.close() |
| sendRequestHeaders(data_len) |
| |
| # Send the data and receive the response. We no longer retry on socket errors as the |
| # request may have already been partially processed by the server and we don't want to |
| # repeat it. We do retry on 401 Unauthorized if we sent cookies as they may have |
| # expired and we don't send the Authorization header with cookies for Kerberos. |
| retry_because_401 = sendDataRecvResp(data) |
| if retry_because_401: |
| if self.__verbose: |
| print('Cookies expired, restarting authentication...', file=sys.stderr) |
| # Received "401 Unauthorized" response and cookies have been cleaned. Retry with |
| # the same connection. |
| sendRequestHeaders(data_len) |
| sendDataRecvResp(data) |
| |
| if self.code >= 300: |
| # Report any http response code that is not 1XX (informational response) or |
| # 2XX (successful). |
| body = self.readBody().decode('utf-8') |
| raise HttpError(self.code, self.message, body, self.headers) |
