more o365 oauth adjustments

custom_components/mail_and_packages/application_credentials.py

@@ -1,4 +1,5 @@
 """application_credentials platform for mail and packages."""
+
 from homeassistant.core import HomeAssistant
 from homeassistant.components.application_credentials import AuthorizationServer
 
@@ -7,9 +8,10 @@ async def async_get_authorization_server(hass: HomeAssistant) -> AuthorizationSe
     """Return authorization server."""
     return AuthorizationServer(
         authorize_url="https://accounts.google.com/o/oauth2/v2/auth",
-        token_url="https://oauth2.googleapis.com/token"
+        token_url="https://oauth2.googleapis.com/token",
     )
 
+
 async def async_get_description_placeholders(hass: HomeAssistant) -> dict[str, str]:
     """Return description placeholders for the credentials dialog."""
     return {
@@ -20,4 +22,4 @@ async def async_get_description_placeholders(hass: HomeAssistant) -> dict[str, s
             "https://github.com/moralmunky/Home-Assistant-Mail-And-Packages/wiki"
         ),
         "oauth_creds_url": "https://console.cloud.google.com/apis/credentials",
-    }
+    }

custom_components/mail_and_packages/config_flow.py

@@ -1,4 +1,5 @@
 """Adds config flow for Mail and Packages."""
+
 from collections.abc import Mapping
 import logging
 from os import path
@@ -8,6 +9,7 @@
 from googleapiclient.discovery import build
 
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResult
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers import config_entry_oauth2_flow
@@ -140,10 +142,11 @@ async def _validate_user_input(user_input: dict) -> tuple:
     return errors, user_input
 
 
-def _get_mailboxes(host: str, port: int, user: str , pwd: str | None, token: str | None) -> list:
+def _get_mailboxes(
+    host: str, port: int, user: str, pwd: str | None, token: str | None
+) -> list:
     """Get list of mailbox folders from mail server."""
     if token:
-        token = {"token": token}
         account = login(host, port, user, None, token)
     else:
         account = login(host, port, user, pwd, None)
@@ -187,9 +190,7 @@ def _get_default(key: str, fallback_default: Any = None) -> None:
             vol.Required(
                 CONF_CLIENT_ID, default=_get_default(CONF_CLIENT_ID)
             ): cv.string,
-            vol.Required(
-                CONF_SECRET, default=_get_default(CONF_SECRET)
-            ): cv.string,
+            vol.Required(CONF_SECRET, default=_get_default(CONF_SECRET)): cv.string,
         }
     )
 
@@ -213,7 +214,9 @@ def _get_default(key: str, fallback_default: Any = None) -> None:
     )
 
 
-def _get_schema_step_2(data: list, user_input: list, default_dict: list) -> Any:
+async def _get_schema_step_2(
+    data: list, user_input: list, default_dict: list, hass: HomeAssistant | None = None
+) -> Any:
     """Get a schema using the default_dict as a backup."""
     if user_input is None:
         user_input = {}
@@ -224,14 +227,10 @@ def _get_default(key: str, fallback_default: Any = None) -> None:
 
     # No password, likely oAuth login
     if data[CONF_METHOD] == "o365":
-        app = O365Auth(data)
-        app.client()
+        app = O365Auth(hass, data)
+        await app.client()
         mailboxes = _get_mailboxes(
-            data[CONF_HOST],
-            data[CONF_PORT],
-            data[CONF_USERNAME],
-            None,
-            app.token
+            data[CONF_HOST], data[CONF_PORT], data[CONF_USERNAME], None, app.token
         )
 
     else:
@@ -296,6 +295,7 @@ class OAuth2FlowHandler(
     config_entry_oauth2_flow.AbstractOAuth2FlowHandler, domain=DOMAIN
 ):
     """Config flow to handle Google Mail OAuth2 authentication."""
+
     DOMAIN = DOMAIN
     reauth_entry: ConfigEntry | None = None
     _data = {}
@@ -309,23 +309,21 @@ def extra_authorize_data(self) -> dict[str, Any]:
             "access_type": "offline",
             "prompt": "consent",
         }
-    
+
     async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
         """Perform reauth upon an API authentication error."""
         self.reauth_entry = self.hass.config_entries.async_get_entry(
             self.context["entry_id"]
         )
         return await self.async_step_reauth_confirm()
-
 
     async def async_step_reauth_confirm(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
         """Confirm reauth dialog."""
         if user_input is None:
             return self.async_show_form(step_id="reauth_confirm")
-        return await self.async_step_user()    
-
+        return await self.async_step_user()
 
     async def async_oauth_create_entry(self, data: dict[str, Any]) -> FlowResult:
         """Create an entry for the flow, or update existing entry."""
@@ -339,7 +337,6 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> FlowResult:
             await self.hass.config_entries.async_reload(self.reauth_entry.entry_id)
             return self.async_abort(reason="reauth_successful")
 
-
         def _get_profile() -> str:
             """Get profile from inside the executor."""
             users = build("gmail", "v1", credentials=credentials).users()
@@ -352,7 +349,7 @@ def _get_profile() -> str:
         self._data[CONF_METHOD] = "gmail"
 
         return await self._show_config_2(None)
-    
+
     async def async_step_config_2(self, user_input=None):
         """Configure form step 2."""
         self._errors = {}
@@ -388,7 +385,9 @@ async def _show_config_2(self, user_input):
 
         return self.async_show_form(
             step_id="config_2",
-            data_schema=_get_schema_step_2(self._data, user_input, defaults),
+            data_schema=await _get_schema_step_2(
+                self._data, user_input, defaults, self.hass
+            ),
             errors=self._errors,
         )
 
@@ -417,7 +416,8 @@ async def _show_config_3(self, user_input):
             step_id="config_3",
             data_schema=_get_schema_step_3(user_input, defaults),
             errors=self._errors,
-        )    
+        )
+
 
 @config_entries.HANDLERS.register(DOMAIN)
 class MailAndPackagesFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
@@ -454,18 +454,24 @@ async def async_step_o365(
 
             try:
                 await app.client()
-                valid = True
+                valid = test_login(
+                    user_input[CONF_HOST],
+                    user_input[CONF_PORT],
+                    user_input[CONF_USERNAME],
+                    None,
+                    app.token,
+                )
             except TokenError:
                 _LOGGER.error("Problems obtaining oAuth token.")
-                self._problem = "token"
+                self._errors["base"] = "token"
             except MissingTenantID:
                 _LOGGER.error("Missing tenant ID.")
-                self._problem = "tenant"
+                self._errors["base"] = "tenant"
 
             if not valid:
-                self._errors["base"] = self._problem
-            else:
+                self._errors["base"] = "communication"
 
+            if not self._errors:
                 return await self.async_step_config_2()
 
             return await self._show_config_o365(user_input)
@@ -482,7 +488,7 @@ async def _show_config_o365(self, user_input):
             data_schema=_get_schema_step_o365(user_input, defaults),
             errors=self._errors,
         )
-    
+
     async def async_step_manual(self, user_input=None):
         """Handle a flow initialized by the user."""
         self._errors = {}
@@ -553,7 +559,9 @@ async def _show_config_2(self, user_input):
 
         return self.async_show_form(
             step_id="config_2",
-            data_schema=_get_schema_step_2(self._data, user_input, defaults),
+            data_schema=await _get_schema_step_2(
+                self._data, user_input, defaults, self.hass
+            ),
             errors=self._errors,
         )
 
@@ -662,7 +670,7 @@ async def _show_step_options_2(self, user_input):
 
         return self.async_show_form(
             step_id="options_2",
-            data_schema=_get_schema_step_2(self._data, user_input, defaults),
+            data_schema=await _get_schema_step_2(self._data, user_input, defaults),
             errors=self._errors,
         )
 

custom_components/mail_and_packages/helpers.py

@@ -112,7 +112,9 @@ async def check_ffmpeg() -> bool:
     return which("ffmpeg")
 
 
-async def test_login(host: str, port: int, user: str | None, pwd: str) -> bool:
+def test_login(
+    host: str, port: int, user: str | None, pwd: str, token: str | None
+) -> bool:
     """Test IMAP login to specified server.
 
     Returns success boolean
@@ -125,7 +127,8 @@ async def test_login(host: str, port: int, user: str | None, pwd: str) -> bool:
         _LOGGER.error("Error connecting into IMAP Server: %s", str(err))
         return False
     # Validate we can login to mail server
-    if user is None:
+    if token:
+        pwd = generate_auth_string(user, token)
         try:
             account.authenticate("XOAUTH2", lambda x: pwd.encode("utf-8"))
         except Exception as err:
@@ -439,7 +442,7 @@ def fetch(
 
 
 def login(
-    host: str, port: int, user: str , pwd: str | None, token: str | None
+    host: str, port: int, user: str, pwd: str | None, token: str | None
 ) -> Union[bool, Type[imaplib.IMAP4_SSL]]:
     """Login to IMAP server.
 

custom_components/mail_and_packages/oauth.py

@@ -17,7 +17,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-def generate_auth_string(user, token) -> str:
+def generate_auth_string(user: str, token: str) -> str:
     return f"user={user}\x01auth=Bearer {token}\x01\x01"
 
 
@@ -58,17 +58,17 @@ async def client(self) -> None:
         )
 
         result = await self.hass.async_add_executor_job(
-                app.acquire_token_silent,
-                self._scope, 
-                None,
-            )
+            app.acquire_token_silent,
+            self._scope,
+            None,
+        )
 
         if not result:
             _LOGGER.debug("No token cached, getting new token.")
             result = await self.hass.async_add_executor_job(
-                    app.acquire_token_for_client,
-                    self._scope,
-                )
+                app.acquire_token_for_client,
+                self._scope,
+            )
 
         if CONF_TOKEN in result:
             self.token = result[CONF_TOKEN]
@@ -81,6 +81,7 @@ async def client(self) -> None:
             )
             raise TokenError
 
+
 class MissingTenantID(Exception):
     """Exception for missing tenant ID."""