From 92abf184310322c73f2997150eea43f9f29593b7 Mon Sep 17 00:00:00 2001 From: Matthew Date: Tue, 5 Nov 2024 10:25:47 +1030 Subject: [PATCH] Removed aioboto3 and replaced with asyncio --- .../insnrg_chlorinator/config_flow.py | 131 +++++++++--------- .../insnrg_chlorinator/coordinator.py | 1 - .../insnrg_chlorinator/manifest.json | 2 +- 3 files changed, 66 insertions(+), 68 deletions(-) diff --git a/custom_components/insnrg_chlorinator/config_flow.py b/custom_components/insnrg_chlorinator/config_flow.py index 780cff4..baa1bcd 100644 --- a/custom_components/insnrg_chlorinator/config_flow.py +++ b/custom_components/insnrg_chlorinator/config_flow.py @@ -2,7 +2,8 @@ import logging import aiohttp import async_timeout -import aioboto3 +import asyncio +import boto3 from datetime import datetime, timedelta from homeassistant import config_entries from homeassistant.core import callback @@ -22,80 +23,78 @@ class InsnrgChlorinatorConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): async def async_step_user(self, user_input=None): errors = {} if user_input is not None: - # Log input for debugging #_LOGGER.debug("Received user input: %s", user_input) - username = user_input["Username"] password = user_input["Password"] - session = aioboto3.Session() - - try: - async with session.client('cognito-idp', region_name='us-east-2') as client: - # Log when starting SRP auth - _LOGGER.debug("Starting SRP authentication") - - aws_srp = AWSSRP( - username=username, - password=password, - pool_id=PoolId, - client_id=ClientId, - client=client + def initiate_auth_sync(username, password): + """Synchronously perform USER_SRP_AUTH and process challenges.""" + client = boto3.client('cognito-idp', region_name='us-east-2') + + # Start SRP authentication + aws_srp = AWSSRP( + username=username, + password=password, + pool_id=PoolId, + client_id=ClientId, + client=client + ) + auth_params = aws_srp.get_auth_params() + + # Initiate authentication + response = client.initiate_auth( + ClientId=ClientId, + AuthFlow='USER_SRP_AUTH', + AuthParameters=auth_params + ) + _LOGGER.debug("Received auth response") + #_LOGGER.debug("Auth response contents: %s", response) + + if response.get('ChallengeName') == 'PASSWORD_VERIFIER': + _LOGGER.debug("Processing password challenge") + challenge_responses = aws_srp.process_challenge( + response['ChallengeParameters'], + auth_params ) + #_LOGGER.debug("ChallengeParameters: %s", response['ChallengeParameters']) # This should include PASSWORD_CLAIM_SECRET_BLOCK, PASSWORD_CLAIM_SIGNATURE, TIMESTAMP, USERNAME (as a UUID) - auth_params = aws_srp.get_auth_params() - #_LOGGER.debug("Generated auth params: %s", auth_params) - - response = await client.initiate_auth( + # Respond to password challenge + response = client.respond_to_auth_challenge( ClientId=ClientId, - AuthFlow='USER_SRP_AUTH', - AuthParameters=auth_params + ChallengeName='PASSWORD_VERIFIER', + ChallengeResponses=challenge_responses ) + #_LOGGER.debug("Challenge response received: %s", response) + + # Extract tokens + auth_result = response['AuthenticationResult'] + return { + "access_token": auth_result['AccessToken'], + "expiry": timedelta(seconds=auth_result['ExpiresIn']) + datetime.now(), + "id_token": auth_result['IdToken'], + "refresh_token": auth_result['RefreshToken'] + } + _LOGGER.debug("Authentication successful, tokens retrieved. Getting System ID") - _LOGGER.debug("Received auth response") - #_LOGGER.debug("Auth response contents: %s", response) - - if response.get('ChallengeName') == 'PASSWORD_VERIFIER': - _LOGGER.debug("Processing password challenge") - - challenge_responses = aws_srp.process_challenge( - response['ChallengeParameters'], # This should include PASSWORD_CLAIM_SECRET_BLOCK, PASSWORD_CLAIM_SIGNATURE, TIMESTAMP, USERNAME (as a UUID) - auth_params - ) - #_LOGGER.debug("ChallengeParameters: %s", response['ChallengeParameters']) - - response = await client.respond_to_auth_challenge( - ClientId=ClientId, - ChallengeName='PASSWORD_VERIFIER', - ChallengeResponses=challenge_responses - ) - - #_LOGGER.debug("Challenge response received: %s", response) - - # Extract tokens - auth_result = response['AuthenticationResult'] - access_token = auth_result['AccessToken'] - expiry = timedelta(seconds=auth_result['ExpiresIn']) + datetime.now() - id_token = auth_result['IdToken'] - refresh_token = auth_result['RefreshToken'] - - # Log success - _LOGGER.debug("Authentication successful, tokens retrieved. Getting System ID") - - system_id = await self._get_system_id(id_token) - - # Store tokens and additional data - return self.async_create_entry( - title="INSNRG Chlorinator", - data={ - "Username": username, - "access_token": access_token, - "expiry": expiry, - "id_token": id_token, - "refresh_token": refresh_token, - "system_id": system_id, - } - ) + try: + # Run the synchronous authentication function in the executor + _LOGGER.debug("Starting SRP authentication") + auth_result = await self.hass.async_add_executor_job( + initiate_auth_sync, username, password + ) + + # Use the id_token to retrieve the system ID asynchronously + system_id = await self._get_system_id(auth_result['id_token']) + + # Store tokens and additional data in the configuration entry + return self.async_create_entry( + title="INSNRG Chlorinator", + data={ + "Username": username, + **auth_result, + "system_id": system_id, + } + ) except ClientError as e: _LOGGER.error(f"Authentication failed: {e}") errors["base"] = "auth_failed" diff --git a/custom_components/insnrg_chlorinator/coordinator.py b/custom_components/insnrg_chlorinator/coordinator.py index bb7e303..02587ed 100644 --- a/custom_components/insnrg_chlorinator/coordinator.py +++ b/custom_components/insnrg_chlorinator/coordinator.py @@ -3,7 +3,6 @@ import asyncio import async_timeout import boto3 -from botocore.exceptions import ClientError import json from datetime import datetime, timedelta from botocore.exceptions import ClientError diff --git a/custom_components/insnrg_chlorinator/manifest.json b/custom_components/insnrg_chlorinator/manifest.json index 2f33fd7..fd43711 100644 --- a/custom_components/insnrg_chlorinator/manifest.json +++ b/custom_components/insnrg_chlorinator/manifest.json @@ -7,6 +7,6 @@ "documentation": "https://github.com/Mattat01/ha-insnrg-chlorinator", "iot_class": "cloud_polling", "issue_tracker": "https://github.com/Mattat01/ha-insnrg-chlorinator/issues", - "requirements": [ "aioboto3" ], + "requirements": [], "version": "1.0.1" } \ No newline at end of file