Merge branch 'fronzbot:dev' into dev

Author: austinc3030

blinkapp/blinkapp.py

@@ -5,7 +5,7 @@
 from datetime import datetime, timedelta
 from aiohttp import ClientSession
 from blinkpy.blinkpy import Blink
-from blinkpy.auth import Auth
+from blinkpy.auth import Auth, BlinkTwoFARequiredError
 from blinkpy.helpers.util import json_load
 
 CREDFILE = environ.get("CREDFILE")
@@ -26,7 +26,10 @@ async def start(session: ClientSession):
     """Startup blink app."""
     blink = Blink(session=session)
     blink.auth = Auth(await json_load(CREDFILE), session=session)
-    await blink.start()
+    try:
+        await blink.start()
+    except BlinkTwoFARequiredError:
+        await blink.prompt_2fa()
     return blink
 
 

blinkpy/api.py

@@ -62,7 +62,7 @@ async def request_login(
 
     if is_refresh:
         form_data["grant_type"] = OAUTH_GRANT_TYPE_REFRESH_TOKEN
-        form_data["refresh_token"] = auth._refresh_token
+        form_data["refresh_token"] = auth.refresh_token
     else:
         form_data["grant_type"] = OAUTH_GRANT_TYPE_PASSWORD
         form_data["password"] = login_data["password"]

blinkpy/auth.py

@@ -32,6 +32,7 @@ def __init__(
         session=None,
         agent=DEFAULT_USER_AGENT,
         app_build=APP_BUILD,
+        callback=None,
     ):
         """
         Initialize auth handler.
@@ -48,8 +49,8 @@ def __init__(
         self.data = login_data
         self.token = login_data.get("token", None)
         self.expires_in = login_data.get("expires_in", None)
-        self.expiration_date = None
-        self._refresh_token = login_data.get("refresh_token", None)
+        self.expiration_date = login_data.get("expiration_date", None)
+        self.refresh_token = login_data.get("refresh_token", None)
         self.host = login_data.get("host", None)
         self.region_id = login_data.get("region_id", None)
         self.client_id = login_data.get("client_id", None)
@@ -63,12 +64,16 @@ def __init__(
         self._app_build = app_build
         self.session = session if session else ClientSession()
 
+        # Callback to notify on token refresh
+        self.callback = callback
+
     @property
     def login_attributes(self):
         """Return a dictionary of login attributes."""
         self.data["token"] = self.token
         self.data["expires_in"] = self.expires_in
-        self.data["refresh_token"] = self._refresh_token
+        self.data["expiration_date"] = self.expiration_date
+        self.data["refresh_token"] = self.refresh_token
         self.data["host"] = self.host
         self.data["region_id"] = self.region_id
         self.data["client_id"] = self.client_id
@@ -96,12 +101,6 @@ def validate_login(self):
             self.data = util.prompt_login_data(self.data)
         self.data = util.validate_login_data(self.data)
 
-    def validate_2fa(self):
-        """Check 2FA information and prompt if not available."""
-        self.data["2fa_code"] = self.data.get("2fa_code", None)
-        if not self.no_prompt:
-            self.data = util.prompt_2fa_data(self.data)
-
     async def login(self, login_url=LOGIN_ENDPOINT, refresh=False):
         """Attempt OAuth login to blink servers."""
         self.validate_login()
@@ -116,8 +115,7 @@ async def login(self, login_url=LOGIN_ENDPOINT, refresh=False):
             if response.status == 200:
                 return await response.json()
             if response.status == 412:
-                self.validate_2fa()
-                return await self.login()
+                raise BlinkTwoFARequiredError
             raise LoginError
         except AttributeError as error:
             raise LoginError from error
@@ -130,16 +128,24 @@ def logout(self, blink):
         """Log out."""
         return api.request_logout(blink)
 
-    async def refresh_token(self):
-        """Refresh auth token."""
+    async def refresh_tokens(self, refresh=False):
+        """Create or refresh access token."""
         self.is_errored = True
         try:
-            _LOGGER.info("Token expired, attempting automatic refresh.")
-            self.login_response = await self.login()
+            _LOGGER.info(
+                f"{'Refreshing' if refresh else 'Obtaining'} authentication token."
+            )
+            self.login_response = await self.login(refresh=refresh)
             self.extract_login_info()
-            self.tier_info = await self.get_tier_info()
-            self.extract_tier_info()
+
+            if not refresh:
+                self.tier_info = await self.get_tier_info()
+                self.extract_tier_info()
+
             self.is_errored = False
+        except BlinkTwoFARequiredError as error:
+            _LOGGER.error("Two-factor authentication required. Waiting for otp.")
+            raise BlinkTwoFARequiredError from error
         except LoginError as error:
             _LOGGER.error("Login endpoint failed. Try again later.")
             raise TokenRefreshFailed from error
@@ -153,7 +159,7 @@ def extract_login_info(self):
         self.token = self.login_response["access_token"]
         self.expires_in = self.login_response["expires_in"]
         self.expiration_date = time.time() + self.expires_in
-        self._refresh_token = self.login_response["refresh_token"]
+        self.refresh_token = self.login_response["refresh_token"]
 
     def extract_tier_info(self):
         """Extract tier info from tier info response."""
@@ -165,7 +171,7 @@ async def startup(self):
         """Initialize tokens for communication."""
         self.validate_login()
         if None in self.login_attributes.values():
-            await self.refresh_token()
+            await self.refresh_tokens()
 
     async def validate_response(self, response: ClientResponse, json_resp):
         """Check for valid response."""
@@ -191,7 +197,7 @@ async def validate_response(self, response: ClientResponse, json_resp):
     def need_refresh(self):
         """Check if token needs refresh."""
         if self.expiration_date is None:
-            return self._refresh_token is not None
+            return self.refresh_token is not None
 
         return self.expiration_date - time.time() < 60
 
@@ -219,7 +225,14 @@ async def query(
         """
         try:
             if not skip_refresh_check and self.need_refresh():
-                await self.login(refresh=True)
+                await self.refresh_tokens(refresh=True)
+
+                if "Authorization" in headers:
+                    # update the authorization header with the new token
+                    headers["Authorization"] = f"Bearer {self.token}"
+
+                if self.callback is not None:
+                    self.callback()
 
             if reqtype == "get":
                 response = await self.session.get(
@@ -253,7 +266,7 @@ async def query(
         except UnauthorizedError:
             try:
                 if not is_retry:
-                    await self.refresh_token()
+                    await self.refresh_tokens()
                     return await self.query(
                         url=url,
                         data=data,
@@ -282,5 +295,9 @@ class BlinkBadResponse(Exception):
     """Class to throw bad json response exception."""
 
 
+class BlinkTwoFARequiredError(Exception):
+    """Class to throw two-factor authentication required exception."""
+
+
 class UnauthorizedError(Exception):
     """Class to throw an unauthorized access error."""

blinkpy/blinkpy.py

@@ -32,7 +32,7 @@
     TIMEOUT_MEDIA,
 )
 from blinkpy.helpers.constants import __version__
-from blinkpy.auth import Auth, TokenRefreshFailed, LoginError
+from blinkpy.auth import Auth, BlinkTwoFARequiredError, TokenRefreshFailed, LoginError
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -90,6 +90,16 @@ def account_id(self):
         """Return the account id."""
         return self.auth.account_id
 
+    async def prompt_2fa(self):
+        """Prompt user for two-factor authentication code."""
+        code = input("Enter the two-factor authentication code: ")
+        await self.send_2fa_code(code)
+
+    async def send_2fa_code(self, code):
+        """Send the two-factor authentication code to complete login."""
+        self.auth.data["2fa_code"] = code
+        await self.start()
+
     @util.Throttle(seconds=MIN_THROTTLE_TIME)
     async def refresh(self, force=False, force_cache=False):
         """
@@ -127,6 +137,8 @@ async def start(self):
             _LOGGER.error("Cannot setup Blink platform.")
             self.available = False
             return False
+        except BlinkTwoFARequiredError:
+            raise
 
         if not self.last_refresh:
             # Initialize last_refresh to be just before the refresh delay period.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "blinkpy"
-version = "0.24.0b0"
+version = "0.24.0b1"
 license = {text = "MIT"}
 description = "A Blink camera Python Library."
 readme = "README.rst"

tests/test_auth.py

@@ -6,6 +6,7 @@
 from aiohttp import ClientConnectionError, ContentTypeError
 from blinkpy.auth import (
     Auth,
+    BlinkTwoFARequiredError,
     TokenRefreshFailed,
     BlinkBadResponse,
     UnauthorizedError,
@@ -146,13 +147,13 @@ def test_header_no_token(self):
         self.assertEqual(self.auth.header, None)
 
     @mock.patch("blinkpy.auth.Auth.validate_login")
-    @mock.patch("blinkpy.auth.Auth.refresh_token")
+    @mock.patch("blinkpy.auth.Auth.refresh_tokens")
     async def test_auth_startup(self, mock_validate, mock_refresh):
         """Test auth startup."""
         await self.auth.startup()
 
     @mock.patch("blinkpy.auth.Auth.query")
-    async def test_refresh_token(self, mock_resp):
+    async def test_refresh_tokens(self, mock_resp):
         """Test refresh token method."""
         mock_json = mock.AsyncMock(
             return_value={
@@ -165,34 +166,42 @@ async def test_refresh_token(self, mock_resp):
         )
         mock_resp.side_effect = [
             # first request simulates otp required
-            mock.AsyncMock(status=412),
             mock.AsyncMock(status=200, json=mock_json),
             {"tier": "test", "account_id": 5678},
             mock.AsyncMock(status=400),
             mock.AsyncMock(status=200, json=mock.AsyncMock(side_effect=AttributeError)),
         ]
 
         self.auth.no_prompt = True
-        self.assertTrue(await self.auth.refresh_token())
+        self.assertTrue(await self.auth.refresh_tokens())
         self.assertEqual(self.auth.region_id, "test")
         self.assertEqual(self.auth.token, "foobar")
-        self.assertEqual(self.auth._refresh_token, "baz1")
+        self.assertEqual(self.auth.refresh_token, "baz1")
         self.assertEqual(self.auth.account_id, 5678)
         self.assertEqual(self.auth.user_id, None)
 
         with self.assertRaises(TokenRefreshFailed):
-            await self.auth.refresh_token()
+            await self.auth.refresh_tokens()
 
         with self.assertRaises(TokenRefreshFailed):
-            await self.auth.refresh_token()
+            await self.auth.refresh_tokens()
+
+    @mock.patch("blinkpy.auth.Auth.query")
+    async def test_refresh_token_otp_required(self, mock_resp):
+        """Test refresh token method."""
+        mock_resp.side_effect = [mock.AsyncMock(status=412)]
+
+        self.auth.no_prompt = True
+        with self.assertRaises(BlinkTwoFARequiredError):
+            await self.auth.refresh_tokens()
 
     @mock.patch("blinkpy.auth.Auth.login")
     async def test_refresh_token_failed(self, mock_login):
         """Test refresh token failed."""
         mock_login.return_value = {}
         self.auth.is_errored = False
         with self.assertRaises(TokenRefreshFailed):
-            await self.auth.refresh_token()
+            await self.auth.refresh_tokens()
         self.assertTrue(self.auth.is_errored)
 
     @mock.patch("blinkpy.auth.api.request_logout")
@@ -206,15 +215,15 @@ async def test_logout(self, mock_req):
         "blinkpy.auth.Auth.validate_response",
         mock.AsyncMock(side_effect=[UnauthorizedError, "foobar"]),
     )
-    @mock.patch("blinkpy.auth.Auth.refresh_token", mock.AsyncMock(return_value=True))
+    @mock.patch("blinkpy.auth.Auth.refresh_tokens", mock.AsyncMock(return_value=True))
     @mock.patch("blinkpy.auth.Auth.query", mock.AsyncMock(return_value="foobar"))
     async def test_query_retry(self):  # , mock_refresh, mock_validate):
         """Check handling of request retry."""
         self.auth.session = MockSession()
         self.assertEqual(await self.auth.query(url="http://example.com"), "foobar")
 
     @mock.patch("blinkpy.auth.Auth.validate_response")
-    @mock.patch("blinkpy.auth.Auth.refresh_token")
+    @mock.patch("blinkpy.auth.Auth.refresh_tokens")
     async def test_query_retry_failed(self, mock_refresh, mock_validate):
         """Check handling of failed retry request."""
         self.auth.session = MockSession()
@@ -241,7 +250,7 @@ async def test_query(self, mock_validate):
         self.assertIsNone(await self.auth.query("URL", "data", "headers", "post"))
 
         mock_validate.side_effect = UnauthorizedError
-        self.auth.refresh_token = mock.AsyncMock()
+        self.auth.refresh_tokens = mock.AsyncMock()
         self.assertIsNone(await self.auth.query("URL", "data", "headers", "post"))
 
     @mock.patch("blinkpy.auth.Auth.validate_response")
@@ -255,7 +264,7 @@ async def test_query_refresh_token_exists(
             mock.AsyncMock(status=200),
             mock.AsyncMock(status=200),
         ]
-        self.auth._refresh_token = "baz1"
+        self.auth.refresh_token = "baz1"
         self.auth.data = {"username": "foo", "password": "bar"}
         await self.auth.query("URL", "data", "headers", "get")
 
@@ -285,7 +294,7 @@ async def test_query_auth_expired(
         ]
         self.auth.expiration_date = time.time() - 100
         self.auth.data = {"username": "foo", "password": "bar"}
-        self.auth._refresh_token = "baz1"
+        self.auth.refresh_token = "baz1"
 
         await self.auth.query("URL", "data", "headers", "get")