text: # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  # SPDX-License-Identifier: Apache-2.0.
  
  from test import NativeResourceTest, TIMEOUT
  from awscrt.http import HttpProxyOptions, HttpProxyAuthenticationType, HttpProxyConnectionType, HttpClientConnection, HttpClientStream, HttpRequest
  from awscrt.io import init_logging, LogLevel, ClientTlsContext, TlsContextOptions, ClientBootstrap, ClientTlsContext, DefaultHostResolver, EventLoopGroup
  from awscrt.mqtt import Client, Connection
  import os
  import unittest
  from test.test_http_client import Response
  from test.test_mqtt import create_client_id
  
  
  """
  
  # AWS_TEST_HTTP_PROXY_HOST - host address of the proxy to use for tests that make open connections to the proxy
  # AWS_TEST_HTTP_PROXY_PORT - port to use for tests that make open connections to the proxy
  # AWS_TEST_HTTPS_PROXY_HOST - host address of the proxy to use for tests that make tls-protected connections to the
      proxy
  # AWS_TEST_HTTPS_PROXY_PORT - port to use for tests that make tls-protected connections to the proxy
  # AWS_TEST_HTTP_PROXY_BASIC_HOST - host address of the proxy to use for tests that make open connections to the proxy
      with basic authentication
  # AWS_TEST_HTTP_PROXY_BASIC_PORT - port to use for tests that make open connections to the proxy with basic
      authentication
  
  # AWS_TEST_BASIC_AUTH_USERNAME - username to use when using basic authentication to the proxy
  # AWS_TEST_BASIC_AUTH_PASSWORD - password to use when using basic authentication to the proxy
  
  # AWS_TEST_TLS_CERT_PATH - file path to certificate used to initialize the tls context of the mqtt connection
  # AWS_TEST_TLS_KEY_PATH - file path to the key used to initialize the tls context of the mqtt connection
  # AWS_TEST_TLS_ROOT_CERT_PATH - file path to the root CA used to initialize the tls context of the mqtt connection
  
  # AWS_TEST_IOT_MQTT_ENDPOINT - AWS account-specific endpoint to connect to IoT core by
  
  """
  
  
  class ProxyTestType:
      FORWARDING = 0
      TUNNELING_HTTP = 1
      TUNNELING_HTTPS = 2
      TUNNELING_DOUBLE_TLS = 3
      LEGACY_HTTP = 4
      LEGACY_HTTPS = 5
  
  
  class ProxyTestConfiguration():
      HTTP_PROXY_HOST = os.environ.get('AWS_TEST_HTTP_PROXY_HOST')
      HTTP_PROXY_PORT = int(os.environ.get('AWS_TEST_HTTP_PROXY_PORT', '0'))
      HTTPS_PROXY_HOST = os.environ.get('AWS_TEST_HTTPS_PROXY_HOST')
      HTTPS_PROXY_PORT = int(os.environ.get('AWS_TEST_HTTPS_PROXY_PORT', '0'))
      HTTP_PROXY_BASIC_HOST = os.environ.get('AWS_TEST_HTTP_PROXY_BASIC_HOST')
      HTTP_PROXY_BASIC_PORT = int(os.environ.get('AWS_TEST_HTTP_PROXY_BASIC_PORT', '0'))
  
      HTTP_PROXY_BASIC_AUTH_USERNAME = os.environ.get('AWS_TEST_BASIC_AUTH_USERNAME')
      HTTP_PROXY_BASIC_AUTH_PASSWORD = os.environ.get('AWS_TEST_BASIC_AUTH_PASSWORD')
  
      HTTP_PROXY_TLS_CERT_PATH = os.environ.get('AWS_TEST_TLS_CERT_PATH')
      HTTP_PROXY_TLS_KEY_PATH = os.environ.get('AWS_TEST_TLS_KEY_PATH')
      HTTP_PROXY_TLS_ROOT_CA_PATH = os.environ.get('AWS_TEST_TLS_ROOT_CERT_PATH')
  
      HTTP_PROXY_MQTT_ENDPOINT = os.environ.get('AWS_TEST_IOT_MQTT_ENDPOINT')
  
      @staticmethod
      def is_proxy_environment_initialized():
          return ProxyTestConfiguration.HTTP_PROXY_HOST is not None and \
              ProxyTestConfiguration.HTTP_PROXY_PORT > 0 and \
              ProxyTestConfiguration.HTTPS_PROXY_HOST is not None and \
              ProxyTestConfiguration.HTTPS_PROXY_PORT > 0 and \
              ProxyTestConfiguration.HTTP_PROXY_BASIC_HOST is not None and \
              ProxyTestConfiguration.HTTP_PROXY_BASIC_PORT > 0 and \
              ProxyTestConfiguration.HTTP_PROXY_BASIC_AUTH_USERNAME is not None and \
              ProxyTestConfiguration.HTTP_PROXY_BASIC_AUTH_PASSWORD is not None
  
      @staticmethod
      def is_mqtt_env_initialized():
          return ProxyTestConfiguration.is_proxy_environment_initialized() and \
              ProxyTestConfiguration.HTTP_PROXY_MQTT_ENDPOINT is not None and \
              ProxyTestConfiguration.HTTP_PROXY_TLS_CERT_PATH is not None and \
              ProxyTestConfiguration.HTTP_PROXY_TLS_KEY_PATH is not None and \
              ProxyTestConfiguration.HTTP_PROXY_TLS_KEY_PATH is not None
  
      @staticmethod
      def get_proxy_host_for_test(test_type, auth_type):
          if auth_type == HttpProxyAuthenticationType.Basic:
              return ProxyTestConfiguration.HTTP_PROXY_BASIC_HOST
  
          if test_type == ProxyTestType.TUNNELING_DOUBLE_TLS:
              return ProxyTestConfiguration.HTTPS_PROXY_HOST
  
          return ProxyTestConfiguration.HTTP_PROXY_HOST
  
      @staticmethod
      def get_proxy_port_for_test(test_type, auth_type):
          if auth_type == HttpProxyAuthenticationType.Basic:
              return ProxyTestConfiguration.HTTP_PROXY_BASIC_PORT
  
          if test_type == ProxyTestType.TUNNELING_DOUBLE_TLS:
              return ProxyTestConfiguration.HTTPS_PROXY_PORT
  
          return ProxyTestConfiguration.HTTP_PROXY_PORT
  
      @staticmethod
      def get_proxy_connection_type_for_test(test_type):
          if test_type == ProxyTestType.FORWARDING:
              return HttpProxyConnectionType.Forwarding
          elif test_type == ProxyTestType.TUNNELING_DOUBLE_TLS or \
                  test_type == ProxyTestType.TUNNELING_HTTP or \
                  test_type == ProxyTestType.TUNNELING_HTTPS:
              return HttpProxyConnectionType.Tunneling
          else:
              return HttpProxyConnectionType.Legacy
  
      @staticmethod
      def get_proxy_tls_connection_options_for_test(test_type):
          if test_type == ProxyTestType.TUNNELING_DOUBLE_TLS:
              tls_ctx_opt = TlsContextOptions()
              tls_ctx_opt.verify_peer = False
              tls_ctx = ClientTlsContext(tls_ctx_opt)
              tls_conn_opt = tls_ctx.new_connection_options()
              tls_conn_opt.set_server_name("localhost")
              return tls_conn_opt
          else:
              return None
  
      @staticmethod
      def create_http_proxy_options_from_environment(test_type, auth_type):
          return HttpProxyOptions(
              ProxyTestConfiguration.get_proxy_host_for_test(
                  test_type,
                  auth_type),
              ProxyTestConfiguration.get_proxy_port_for_test(
                  test_type,
                  auth_type),
              tls_connection_options=ProxyTestConfiguration.get_proxy_tls_connection_options_for_test(test_type),
              auth_type=auth_type,
              auth_username=ProxyTestConfiguration.HTTP_PROXY_BASIC_AUTH_USERNAME,
              auth_password=ProxyTestConfiguration.HTTP_PROXY_BASIC_AUTH_PASSWORD,
              connection_type=ProxyTestConfiguration.get_proxy_connection_type_for_test(test_type))
  
      @staticmethod
      def get_tls_connection_options_for_test(test_type, host_name):
          if test_type == ProxyTestType.FORWARDING or test_type == ProxyTestType.LEGACY_HTTP or test_type == ProxyTestType.TUNNELING_HTTP:
              return None
          else:
              tls_ctx_opt = TlsContextOptions()
              tls_ctx = ClientTlsContext(tls_ctx_opt)
              tls_conn_opt = tls_ctx.new_connection_options()
              tls_conn_opt.set_server_name(host_name)
              return tls_conn_opt
  
      @staticmethod
      def get_uri_from_test_type(test_type):
          if test_type == ProxyTestType.FORWARDING or test_type == ProxyTestType.LEGACY_HTTP or test_type == ProxyTestType.TUNNELING_HTTP:
              return "www.example.com"
          else:
              return "www.amazon.com"
  
      @staticmethod
      def get_port_from_test_type(test_type):
          if test_type == ProxyTestType.FORWARDING or test_type == ProxyTestType.LEGACY_HTTP or test_type == ProxyTestType.TUNNELING_HTTP:
              return 80
          else:
              return 443
  
  
  class ProxyHttpTest(NativeResourceTest):
  
      def _establish_http_connection(self, test_type, uri, proxy_options):
          event_loop_group = EventLoopGroup()
          host_resolver = DefaultHostResolver(event_loop_group)
          bootstrap = ClientBootstrap(event_loop_group, host_resolver)
  
          connection_future = HttpClientConnection.new(
              host_name=uri,
              port=ProxyTestConfiguration.get_port_from_test_type(test_type),
              bootstrap=bootstrap,
              tls_connection_options=ProxyTestConfiguration.get_tls_connection_options_for_test(
                  test_type,
                  uri),
              proxy_options=proxy_options)
          return connection_future.result(TIMEOUT)
  
      def _do_proxy_http_test(self, test_type, auth_type):
          uri = ProxyTestConfiguration.get_uri_from_test_type(test_type)
          proxy_options = ProxyTestConfiguration.create_http_proxy_options_from_environment(test_type, auth_type)
          connection = self._establish_http_connection(test_type, uri, proxy_options)
  
          request = HttpRequest('GET', '/')
          request.headers.add('host', uri)
          response = Response()
          stream = connection.request(request, response.on_response, response.on_body)
          stream.activate()
  
          # wait for stream to complete
          stream_completion_result = stream.completion_future.result(TIMEOUT)
  
          self.assertEqual(200, response.status_code)
          self.assertEqual(200, stream_completion_result)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_forwarding_proxy_no_auth(self):
          self._do_proxy_http_test(ProxyTestType.FORWARDING, HttpProxyAuthenticationType.Nothing)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_forwarding_proxy_legacy_http_no_auth(self):
          self._do_proxy_http_test(ProxyTestType.LEGACY_HTTP, HttpProxyAuthenticationType.Nothing)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_proxy_legacy_https_no_auth(self):
          self._do_proxy_http_test(ProxyTestType.LEGACY_HTTPS, HttpProxyAuthenticationType.Nothing)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_tunneling_proxy_http_no_auth(self):
          self._do_proxy_http_test(ProxyTestType.TUNNELING_HTTP, HttpProxyAuthenticationType.Nothing)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_tunneling_proxy_https_no_auth(self):
          self._do_proxy_http_test(ProxyTestType.TUNNELING_HTTPS, HttpProxyAuthenticationType.Nothing)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_tunneling_proxy_double_tls_no_auth(self):
          self._do_proxy_http_test(ProxyTestType.TUNNELING_DOUBLE_TLS, HttpProxyAuthenticationType.Nothing)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_forwarding_proxy_basic_auth(self):
          self._do_proxy_http_test(ProxyTestType.FORWARDING, HttpProxyAuthenticationType.Basic)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_forwarding_proxy_legacy_http_basic_auth(self):
          self._do_proxy_http_test(ProxyTestType.LEGACY_HTTP, HttpProxyAuthenticationType.Basic)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_proxy_legacy_https_basic_auth(self):
          self._do_proxy_http_test(ProxyTestType.LEGACY_HTTPS, HttpProxyAuthenticationType.Basic)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_tunneling_proxy_http_basic_auth(self):
          self._do_proxy_http_test(ProxyTestType.TUNNELING_HTTP, HttpProxyAuthenticationType.Basic)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_proxy_environment_initialized(), 'requires proxy test env vars')
      def test_tunneling_proxy_https_basic_auth(self):
          self._do_proxy_http_test(ProxyTestType.TUNNELING_HTTPS, HttpProxyAuthenticationType.Basic)
  
      def _establish_mqtt_connection(self, proxy_options):
          event_loop_group = EventLoopGroup()
          host_resolver = DefaultHostResolver(event_loop_group)
          bootstrap = ClientBootstrap(event_loop_group, host_resolver)
  
          tls_opts = TlsContextOptions.create_client_with_mtls_from_path(
              ProxyTestConfiguration.HTTP_PROXY_TLS_CERT_PATH,
              ProxyTestConfiguration.HTTP_PROXY_TLS_KEY_PATH)
          tls_opts.override_default_trust_store_from_path(ca_filepath=ProxyTestConfiguration.HTTP_PROXY_TLS_ROOT_CA_PATH)
          tls = ClientTlsContext(tls_opts)
  
          client = Client(bootstrap, tls)
          connection = Connection(
              client=client,
              client_id=create_client_id(),
              host_name=ProxyTestConfiguration.HTTP_PROXY_MQTT_ENDPOINT,
              port=8883,
              proxy_options=proxy_options)
          connection.connect().result(TIMEOUT)
          return connection
  
      def _do_proxy_mqtt_test(self, test_type, auth_type):
          proxy_options = ProxyTestConfiguration.create_http_proxy_options_from_environment(test_type, auth_type)
          connection = self._establish_mqtt_connection(proxy_options)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_mqtt_env_initialized(), 'requires proxy and MQTT test env vars')
      def test_tunneling_http_proxy_mqtt_no_auth(self):
          self._do_proxy_mqtt_test(ProxyTestType.TUNNELING_HTTP, HttpProxyAuthenticationType.Nothing)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_mqtt_env_initialized(), 'requires proxy and MQTT test env vars')
      def test_tunneling_http_proxy_mqtt_basic_auth(self):
          self._do_proxy_mqtt_test(ProxyTestType.TUNNELING_HTTP, HttpProxyAuthenticationType.Basic)
  
      @unittest.skipIf(not ProxyTestConfiguration.is_mqtt_env_initialized(), 'requires proxy and MQTT test env vars')
      def test_tunneling_http_proxy_mqtt_double_tls(self):
          self._do_proxy_mqtt_test(ProxyTestType.TUNNELING_DOUBLE_TLS, HttpProxyAuthenticationType.Nothing)
  
  
  if __name__ == '__main__':
      unittest.main()
  