From e005fcfe524fcba1f281ab826111eb92ff8f50e3 Mon Sep 17 00:00:00 2001
From: "Alan D. Tse" <alandtse@gmail.com>
Date: Mon, 1 Mar 2021 19:23:50 -0800
Subject: [PATCH] fix: fix mfa code handling

---
 teslajsonpy/connection.py |  9 ++++++---
 teslajsonpy/controller.py | 11 +++++++++--
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/teslajsonpy/connection.py b/teslajsonpy/connection.py
index 378ee4b3..22e9a23e 100644
--- a/teslajsonpy/connection.py
+++ b/teslajsonpy/connection.py
@@ -75,6 +75,7 @@ def __init__(
             self.__sethead(access_token=self.access_token, expiration=self.expiration)
             _LOGGER.debug("Connecting with existing access token")
         self.websocket = None
+        self.mfa_code: Text = ""
 
     async def get(self, command):
         """Get data from API."""
@@ -101,7 +102,7 @@ async def post(self, command, method="post", data=None):
                 if self.email and self.password:
                     _LOGGER.debug("Getting sso auth code using credentials")
                     self.code = await self.get_authorization_code(
-                        self.email, self.password
+                        self.email, self.password, mfa_code=self.mfa_code
                     )
                 else:
                     _LOGGER.debug("Using existing authorization code")
@@ -396,6 +397,7 @@ async def get_authorization_code(
             if not resp.history:
                 html = await resp.text()
                 if "/mfa/verify" in html:
+                    _LOGGER.debug("Detected MFA request")
                     mfa_resp = await self.websession.get(
                         "https://auth.tesla.com/oauth2/v3/authorize/mfa/factors",
                         params={"transaction_id": transaction_id},
@@ -416,7 +418,7 @@ async def get_authorization_code(
                     #     ]
                     # }
                     mfa_json = await mfa_resp.json()
-                    if len(mfa_json.get("data", [])) > 1:
+                    if len(mfa_json.get("data", [])) >= 1:
                         factor_id = mfa_json["data"][mfa_device]["id"]
                     if not mfa_code:
                         _LOGGER.debug("No MFA provided")
@@ -445,10 +447,11 @@ async def get_authorization_code(
                     resp = await self.websession.post(url, data=data)
                     _process_resp(resp)
             await asyncio.sleep(3)
-        if not (resp.history):
+        if not resp.history or not URL(resp.history[-1].url).query.get("code"):
             _LOGGER.debug("Failed to authenticate")
             raise IncompleteCredentials("Unable to login with credentials")
         code_url = URL(resp.history[-1].url)
+        _LOGGER.debug("Found code %s", code_url.query.get("code"))
         return code_url.query.get("code")
 
     def get_authorization_code_link(self, new=False) -> yarl.URL:
diff --git a/teslajsonpy/controller.py b/teslajsonpy/controller.py
index d1680ab8..c5835b99 100644
--- a/teslajsonpy/controller.py
+++ b/teslajsonpy/controller.py
@@ -10,7 +10,7 @@
 import asyncio
 import logging
 import time
-from typing import Callable, Dict, Optional, Text
+from typing import Callable, Dict, List, Optional, Text
 
 from aiohttp import ClientConnectorError
 import backoff
@@ -260,7 +260,11 @@ def __init__(
         self.enable_websocket = enable_websocket
 
     async def connect(
-        self, test_login=False, wake_if_asleep=False, filtered_vins=None
+        self,
+        test_login: bool = False,
+        wake_if_asleep: bool = False,
+        filtered_vins: Optional[List[Text]] = None,
+        mfa_code: Text = "",
     ) -> Dict[Text, Text]:
         """Connect controller to Tesla.
 
@@ -268,12 +272,15 @@ async def connect(
             test_login (bool, optional): Whether to test credentials only. Defaults to False.
             wake_if_asleep (bool, optional): Whether to wake up any sleeping cars to update state. Defaults to False.
             filtered_vins (list, optional): If not empty, filters the cars by the provided VINs.
+            mfa_code (Text, optional): MFA code to use for connection
 
         Returns
             Dict[Text, Text]: Returns the refresh_token, access_token, and expires_in time
 
         """
 
+        if mfa_code:
+            self.__connection.mfa_code = mfa_code
         cars = await self.get_vehicles()
         self._last_attempted_update_time = time.time()
         self.__update_lock = asyncio.Lock()