"""Authentication handlers.""" import secrets import urllib.parse import re import logging import json from datetime import datetime, timedelta import aiohttp from .const import ( AUTH_MSO_STEP_1_URL, AUTH_MSO_SELF_ASSERTED_URL, AUTH_MSO_GET_TOKEN_URL, AUTH_AW_BASE, AUTH_MSO_CONFIRM_URL, AUTH_MSO_CLIENT_ID, AUTH_MSO_REDIR_URI, AW_APP_USER_AGENT, ) from .exceptions import ( ExpiredAccessTokenError, UnknownEndpointError, InvalidAccountIdError, SelfAssertedError, TokenRequestError ) from .utils import ( random_string, build_code_challenge, decode_oauth_redirect, decode_jwt, ) _LOGGER = logging.getLogger(__name__) class MSOB2CAuth: """Represent an instance of MSO Auth.""" _auth_session: aiohttp.ClientSession | None = None username: str = None _password: str = None next_refresh: datetime = None auth_data: dict = None _pkce_verifier = random_string(43, 128) _pkce_challenge = None _state = secrets.token_urlsafe(32) _csrf_token: str = "" _cookie_cache: dict = {} def __init__(self, username, password, session=None, refresh_token=None): if session: self._auth_session = session else: self._auth_session = aiohttp.ClientSession() self.username = username self._password = password self._refresh_token = refresh_token @property def business_partner_number(self) -> str | None: """Return business partner number.""" if self.auth_data is None: return None return self.auth_data.get("extension_business_partner_number", "") @property def access_token(self) -> str: """Return the access token.""" if self.auth_data is None: return None return self.auth_data.get("access_token") @property def refresh_token(self) -> str: """Return the access token.""" if self._refresh_token is not None: return self._refresh_token if self.auth_data is None: return None return self.auth_data.get("refresh_token") @property def authenticated_headers(self) -> dict: """Return authenticated headers.""" return { "Authorization": f"Bearer {self.access_token}", "Ocp-Apim-Subscription-Key": "adbc43b29a87404cbc297fe6d7a3d10e", "Accept": "application/json", "User-Agent": AW_APP_USER_AGENT, } async def _get_initial_auth_data(self): """Retrieves initial authentication data (CSRF token, transId).""" _LOGGER.debug("B2C Auth: Getting initial auth data") self._state = secrets.token_urlsafe(32) self._pkce_challenge = build_code_challenge(self._pkce_verifier) auth_response = await self._auth_session.get( AUTH_MSO_STEP_1_URL.format( CODE_CHALLENGE=self._pkce_challenge, EMAIL=self.username, STATE=self._state, ), allow_redirects=False, ) if not auth_response.ok: data = await auth_response.text() _LOGGER.error( "B2C Auth: Initial Authorization URL request failed %s: %s", auth_response.status, data, ) return None if auth_response.status == 302: location = auth_response.headers.get("Location") if not location: _LOGGER.error("B2C Auth: Location header not found") return None return location html = await auth_response.text() match = re.search(r"var SETTINGS = {([^;]+)};", html) if not match: _LOGGER.error("B2C Auth: SETTINGS object not found in HTML") return None settings_json_str = match.group(1) csrf_match = re.search(r'"csrf":\s*"([^"]+)"', settings_json_str) trans_id_match = re.search(r'"transId":\s*"([^"]+)"', settings_json_str) if not csrf_match or not trans_id_match: _LOGGER.error("B2C Auth: CSRF or Transaction ID not found in SETTINGS") return None _csrf_token = csrf_match.group(1) trans_id = trans_id_match.group(1) _LOGGER.debug("Got CSRF %s and Transaction ID %s", _csrf_token, trans_id) return _csrf_token, trans_id async def _submit_self_asserted_form(self, trans_id): """Submits the SelfAsserted form.""" _LOGGER.debug("B2C Auth: Submitting self-asserted form") data = { "request_type": "RESPONSE", "email": self.username, "password": self._password, } asserted_login_response = await self._auth_session.post( AUTH_MSO_SELF_ASSERTED_URL.format(STATE=trans_id), headers={ "X-CSRF-TOKEN": self._csrf_token, "Referer": AUTH_MSO_STEP_1_URL.format( CODE_CHALLENGE=self._pkce_challenge, EMAIL=self.username, STATE=self._state, ), "Origin": AUTH_AW_BASE, "Content-Type": "application/x-www-form-urlencoded; charset=UTF-8", "User-Agent": AW_APP_USER_AGENT, }, data=data, ) data = await asserted_login_response.json(content_type="text/json") status = int(data.get("status")) if asserted_login_response.status != 200 or status != 200: _LOGGER.error( "B2C Auth: SelfAsserted request failed %s: %s", asserted_login_response.status, data, ) if status == 400: raise SelfAssertedError( f"{data.get('errorCode', 'Unknown')} {data.get('message', 'Unknown')}" ) return None return asserted_login_response async def _get_confirmation_redirect(self): """Gets the confirmation redirect URL.""" _LOGGER.debug("B2C Auth: Getting confirmation redirect URL") confirm_login_response = await self._auth_session.get( AUTH_MSO_CONFIRM_URL.format(CSRF=self._csrf_token, STATE=self._state), allow_redirects=False, ) if confirm_login_response.status != 302: text = await confirm_login_response.text() _LOGGER.error( "B2C Auth: Confirm request failed %s: %s", confirm_login_response.status, text, ) return None location = confirm_login_response.headers.get("Location") if not location: _LOGGER.error("B2C Auth: Location header not found") return None return location async def _get_token(self, code): """Requests the access token.""" _LOGGER.debug("B2C Auth: Getting access token from authorization code") token_request_response = await self._auth_session.post( AUTH_MSO_GET_TOKEN_URL, data=urllib.parse.urlencode( { "grant_type": "authorization_code", "client_id": AUTH_MSO_CLIENT_ID, "redirect_uri": AUTH_MSO_REDIR_URI, "code": code, "code_verifier": self._pkce_verifier, } ), headers={ "Content-Type": "application/x-www-form-urlencoded; charset=UTF-8", "User-Agent": AW_APP_USER_AGENT, }, ) if token_request_response.status != 200: text = await token_request_response.text() _LOGGER.error( "B2C Auth: Token request failed %s: %s", token_request_response.status, text, ) return None try: token_data = await token_request_response.json() return token_data except (aiohttp.ClientError, json.JSONDecodeError) as e: _LOGGER.error("B2C Auth: Error processing token response: %s", e) return None async def send_refresh_request(self): """Send a request to refresh the access token.""" _LOGGER.debug("B2C Auth: Refreshing access token") if self.access_token is None and self.refresh_token is None: raise ValueError("Not logged in.") if self.next_refresh is not None: if self.next_refresh > datetime.now(): _LOGGER.debug("B2C Auth: Access token not yet expired") return token_request_response = await self._auth_session.post( AUTH_MSO_GET_TOKEN_URL, data=urllib.parse.urlencode( { "grant_type": "refresh_token", "client_id": AUTH_MSO_CLIENT_ID, "refresh_token": self.refresh_token, } ), headers={ "Content-Type": "application/x-www-form-urlencoded; charset=UTF-8", "User-Agent": AW_APP_USER_AGENT, }, ) if token_request_response.status != 200: text = await token_request_response.text() _LOGGER.error( "B2C Auth: Refresh token request failed %s: %s", token_request_response.status, text, ) raise TokenRequestError try: token_data = await token_request_response.json() self.auth_data = token_data self.next_refresh = datetime.now() + timedelta( seconds=token_data["expires_in"] ) self.auth_data = {**self.auth_data, **decode_jwt(self.access_token)} _LOGGER.debug( "B2C Auth: Access token refreshed successfully, new expiration time: %s", self.next_refresh, ) except (aiohttp.ClientError, json.JSONDecodeError) as e: _LOGGER.error("B2C Auth: Error processing refresh response: %s", e) async def send_login_request(self): """Send a request to MSO for Auth.""" if self._refresh_token is not None: _LOGGER.debug("B2C Auth: Using refresh token for authentication") # Attempt to use refresh token first try: await self.send_refresh_request() return except Exception as e: _LOGGER.warning( "B2C Auth: Refresh token failed, falling back to initial login: %s", e, ) auth_data = await self._get_initial_auth_data() if auth_data is None: return if isinstance(auth_data, tuple): self._csrf_token, trans_id = auth_data asserted_response = await self._submit_self_asserted_form(trans_id) if asserted_response is None: return redirect_location = await self._get_confirmation_redirect() if redirect_location is None: return elif "uk.co.anglianwater.myaccount://" in auth_data: redirect_location = auth_data _, code = decode_oauth_redirect(redirect_location) if not code: _LOGGER.error("B2C Auth: Code not found") return token_response = await self._get_token(code) if token_response is None: _LOGGER.error("B2C Auth: Token request failed") return self.auth_data = token_response self.next_refresh = datetime.now() + timedelta( seconds=token_response["expires_in"] ) self._refresh_token = token_response.get("refresh_token") self.auth_data = {**self.auth_data, **decode_jwt(self.access_token)} _LOGGER.debug( "B2C Auth: Access token obtained successfully, new expiration time: %s", self.next_refresh, ) async def send_request( self, method: str, url: str, body: dict | None, headers: dict ) -> dict: """Send a request to the API, and return the JSON response.""" await self.send_refresh_request() if self.access_token is None: _LOGGER.debug("Access token unavailable, not logged in.") raise ExpiredAccessTokenError() async with self._auth_session.request( method=method, url=url, headers=headers, json=body ) as _response: _LOGGER.debug( "Request to %s returned with status %s", url, _response.status ) if _response.ok and _response.content_type == "application/json": return await _response.json() if _response.status == 401: raise ExpiredAccessTokenError() if _response.status == 403: raise InvalidAccountIdError() raise UnknownEndpointError(_response.status, await _response.text())