test_mqtt.py

fromawscrt.authimportAwsCredentialsProvider
fromawscrt.ioimportClientBootstrap, DefaultHostResolver, EventLoopGroup
fromawsiotimportmqtt_connection_builder
importboto3
importbotocore.exceptions
importos
importunittest
importshutil
importtempfile
importuuid
importwarnings

TIMEOUT=100.0
PROXY_HOST=os.environ.get('proxyhost')
PROXY_PORT=int(os.environ.get('proxyport', '0'))

CUSTOM_AUTHORIZER_ENDPOINT=os.environ.get("CUSTOM_AUTHORIZER_ENDPOINT")
CUSTOM_AUTHORIZER_NAME_SIGNED=os.environ.get("CUSTOM_AUTHORIZER_NAME_SIGNED")
CUSTOM_AUTHORIZER_NAME_UNSIGNED=os.environ.get("CUSTOM_AUTHORIZER_NAME_UNSIGNED")
CUSTOM_AUTHORIZER_PASSWORD=os.environ.get("CUSTOM_AUTHORIZER_PASSWORD")
CUSTOM_AUTHORIZER_SIGNATURE=os.environ.get("CUSTOM_AUTHORIZER_SIGNATURE")
CUSTOM_AUTHORIZER_SIGNATURE_UNENCODED=os.environ.get("CUSTOM_AUTHORIZER_SIGNATURE_UNENCODED")
CUSTOM_AUTHORIZER_TOKEN_KEY_NAME=os.environ.get("CUSTOM_AUTHORIZER_TOKEN_KEY_NAME")
CUSTOM_AUTHORIZER_TOKEN_VALUE=os.environ.get("CUSTOM_AUTHORIZER_TOKEN_VALUE")
AWS_DEFAULT_REGION=os.environ.get("AWS_DEFAULT_REGION")

defhas_custom_auth_environment():
return (CUSTOM_AUTHORIZER_ENDPOINTisnotNone) and (CUSTOM_AUTHORIZER_NAME_SIGNEDisnotNone) and \
 (CUSTOM_AUTHORIZER_NAME_UNSIGNEDisnotNone) and (CUSTOM_AUTHORIZER_PASSWORDisnotNone) and \
 (CUSTOM_AUTHORIZER_SIGNATUREisnotNone) and (CUSTOM_AUTHORIZER_TOKEN_KEY_NAMEisnotNone) and \
 (CUSTOM_AUTHORIZER_TOKEN_VALUEisnotNone) and (CUSTOM_AUTHORIZER_SIGNATURE_UNENCODEDisnotNone)

classConfig:
cache=None

def__init__(self, endpoint, cert, key, region, cognito_creds):
self.cert=cert
self.key=key
self.endpoint=endpoint
self.region=region
self.cognito_creds=cognito_creds

@staticmethod
defget():
"""Raises SkipTest if credentials aren't set up correctly"""
ifConfig.cache:
returnConfig.cache

# boto3 caches the HTTPS connection for the API calls, which appears to the unit test
# framework as a leak, so ignore it, that's not what we're testing here
warnings.simplefilter('ignore', ResourceWarning)

try:
secrets=boto3.client('secretsmanager', region_name=AWS_DEFAULT_REGION)
response=secrets.get_secret_value(SecretId='unit-test/endpoint')
endpoint=response['SecretString']
response=secrets.get_secret_value(SecretId='ci/mqtt5/us/mqtt5_thing/cert')
cert=response['SecretString'].encode('utf8')
response=secrets.get_secret_value(SecretId='ci/mqtt5/us/mqtt5_thing/key')
key=response['SecretString'].encode('utf8')
region=secrets.meta.region_name
response=secrets.get_secret_value(SecretId='ci/Cognito/identity_id')
cognito_id=response['SecretString']

cognito=boto3.client('cognito-identity')
response=cognito.get_credentials_for_identity(IdentityId=cognito_id)
cognito_creds=response['Credentials']

Config.cache=Config(endpoint, cert, key, region, cognito_creds)
except (botocore.exceptions.BotoCoreError, botocore.exceptions.ClientError) asex:
print(ex)
raiseunittest.SkipTest("No credentials")

returnConfig.cache


defcreate_client_id():
return'test-aws-iot-device-sdk-python-v2-unit-test-{0}'.format(uuid.uuid4())


classMqttBuilderTest(unittest.TestCase):
def_test_connection(self, connection):
connection.connect().result(TIMEOUT)
connection.disconnect().result(TIMEOUT)

deftest_mtls_from_bytes(self):
config=Config.get()
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)
connection=mqtt_connection_builder.mtls_from_bytes(
cert_bytes=config.cert,
pri_key_bytes=config.key,
endpoint=config.endpoint,
client_id=create_client_id(),
client_bootstrap=bootstrap)
self._test_connection(connection)

deftest_mtls_from_path(self):
config=Config.get()
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)

# test "from path" builder by writing secrets to tempfiles
tmp_dirpath=tempfile.mkdtemp()
try:
cert_filepath=os.path.join(tmp_dirpath, 'cert')
withopen(cert_filepath, 'wb') ascert_file:
cert_file.write(config.cert)

key_filepath=os.path.join(tmp_dirpath, 'key')
withopen(key_filepath, 'wb') askey_file:
key_file.write(config.key)

connection=mqtt_connection_builder.mtls_from_path(
cert_filepath=cert_filepath,
pri_key_filepath=key_filepath,
endpoint=config.endpoint,
client_id=create_client_id(),
client_bootstrap=bootstrap)

finally:
shutil.rmtree(tmp_dirpath)

self._test_connection(connection)

deftest_websockets_default(self):
"""Websocket connection with default credentials provider"""
config=Config.get()
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)
cred_provider=AwsCredentialsProvider.new_default_chain(bootstrap)
connection=mqtt_connection_builder.websockets_with_default_aws_signing(
region=config.region,
credentials_provider=cred_provider,
endpoint=config.endpoint,
client_id=create_client_id(),
client_bootstrap=bootstrap)
self._test_connection(connection)

deftest_websockets_sts(self):
"""Websocket connection with X-Amz-Security-Token query param"""
config=Config.get()
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)
cred_provider=AwsCredentialsProvider.new_static(
access_key_id=config.cognito_creds['AccessKeyId'],
secret_access_key=config.cognito_creds['SecretKey'],
session_token=config.cognito_creds['SessionToken'])
connection=mqtt_connection_builder.websockets_with_default_aws_signing(
region=config.region,
credentials_provider=cred_provider,
endpoint=config.endpoint,
client_id=create_client_id(),
client_bootstrap=bootstrap)
self._test_connection(connection)

@unittest.skipIf(PROXY_HOSTisNone, 'requires "proxyhost" and "proxyport" env vars')
deftest_websockets_proxy(self):
config=Config.get()
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)
cred_provider=AwsCredentialsProvider.new_default_chain(bootstrap)
connection=mqtt_connection_builder.websockets_with_default_aws_signing(
credentials_provider=cred_provider,
websocket_proxy_options=HttpProxyOptions(PROXY_HOST, PROXY_PORT),
endpoint=config.endpoint,
region=config.region,
client_id=create_client_id(),
client_bootstrap=bootstrap)
self._test_connection(connection)

@unittest.skipIf(nothas_custom_auth_environment(), 'requires custom authentication env vars')
deftest_mqtt311_builder_direct_signed_custom_authorizer(self):
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)

connection=mqtt_connection_builder.direct_with_custom_authorizer(
auth_username="",
auth_authorizer_name=CUSTOM_AUTHORIZER_NAME_SIGNED,
auth_authorizer_signature=CUSTOM_AUTHORIZER_SIGNATURE,
auth_password=CUSTOM_AUTHORIZER_PASSWORD,
auth_token_key_name=CUSTOM_AUTHORIZER_TOKEN_KEY_NAME,
auth_token_value=CUSTOM_AUTHORIZER_TOKEN_VALUE,
endpoint=CUSTOM_AUTHORIZER_ENDPOINT,
client_id=create_client_id(),
client_bootstrap=bootstrap)

self._test_connection(connection)

@unittest.skipIf(nothas_custom_auth_environment(), 'requires custom authentication env vars')
deftest_mqtt311_builder_direct_signed_custom_authorizer_unencoded(self):
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)

connection=mqtt_connection_builder.direct_with_custom_authorizer(
auth_username="",
auth_authorizer_name=CUSTOM_AUTHORIZER_NAME_SIGNED,
auth_authorizer_signature=CUSTOM_AUTHORIZER_SIGNATURE_UNENCODED,
auth_password=CUSTOM_AUTHORIZER_PASSWORD,
auth_token_key_name=CUSTOM_AUTHORIZER_TOKEN_KEY_NAME,
auth_token_value=CUSTOM_AUTHORIZER_TOKEN_VALUE,
endpoint=CUSTOM_AUTHORIZER_ENDPOINT,
client_id=create_client_id(),
client_bootstrap=bootstrap)

self._test_connection(connection)

@unittest.skipIf(nothas_custom_auth_environment(), 'requires custom authentication env vars')
deftest_mqtt311_builder_direct_unsigned_custom_authorizer(self):
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)

connection=mqtt_connection_builder.direct_with_custom_authorizer(
auth_username="",
auth_authorizer_name=CUSTOM_AUTHORIZER_NAME_UNSIGNED,
auth_password=CUSTOM_AUTHORIZER_PASSWORD,
endpoint=CUSTOM_AUTHORIZER_ENDPOINT,
client_id=create_client_id(),
client_bootstrap=bootstrap)

self._test_connection(connection)

@unittest.skipIf(nothas_custom_auth_environment(), 'requires custom authentication env vars')
deftest_mqtt311_builder_websocket_unsigned_custom_authorizer(self):
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)

connection=mqtt_connection_builder.websockets_with_custom_authorizer(
auth_username="",
auth_authorizer_name=CUSTOM_AUTHORIZER_NAME_UNSIGNED,
auth_password=CUSTOM_AUTHORIZER_PASSWORD,
endpoint=CUSTOM_AUTHORIZER_ENDPOINT,
client_id=create_client_id(),
client_bootstrap=bootstrap)

self._test_connection(connection)

@unittest.skipIf(nothas_custom_auth_environment(), 'requires custom authentication env vars')
deftest_mqtt311_builder_websocket_signed_custom_authorizer(self):
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)

connection=mqtt_connection_builder.websockets_with_custom_authorizer(
auth_username="",
auth_authorizer_name=CUSTOM_AUTHORIZER_NAME_SIGNED,
auth_authorizer_signature=CUSTOM_AUTHORIZER_SIGNATURE,
auth_password=CUSTOM_AUTHORIZER_PASSWORD,
auth_token_key_name=CUSTOM_AUTHORIZER_TOKEN_KEY_NAME,
auth_token_value=CUSTOM_AUTHORIZER_TOKEN_VALUE,
endpoint=CUSTOM_AUTHORIZER_ENDPOINT,
client_id=create_client_id(),
client_bootstrap=bootstrap)

self._test_connection(connection)

@unittest.skipIf(nothas_custom_auth_environment(), 'requires custom authentication env vars')
deftest_mqtt311_builder_websocket_signed_custom_authorizer_unencoded(self):
elg=EventLoopGroup()
resolver=DefaultHostResolver(elg)
bootstrap=ClientBootstrap(elg, resolver)

connection=mqtt_connection_builder.websockets_with_custom_authorizer(
auth_username="",
auth_authorizer_name=CUSTOM_AUTHORIZER_NAME_SIGNED,
auth_authorizer_signature=CUSTOM_AUTHORIZER_SIGNATURE_UNENCODED,
auth_password=CUSTOM_AUTHORIZER_PASSWORD,
auth_token_key_name=CUSTOM_AUTHORIZER_TOKEN_KEY_NAME,
auth_token_value=CUSTOM_AUTHORIZER_TOKEN_VALUE,
endpoint=CUSTOM_AUTHORIZER_ENDPOINT,
client_id=create_client_id(),
client_bootstrap=bootstrap)

self._test_connection(connection)