Black

homeassistant/auth/mfa_modules/totp.py

@@ -9,23 +9,26 @@ import voluptuous as vol
 from homeassistant.auth.models import User
 from homeassistant.core import HomeAssistant
 
-from . import MultiFactorAuthModule, MULTI_FACTOR_AUTH_MODULES, \
-    MULTI_FACTOR_AUTH_MODULE_SCHEMA, SetupFlow
+from . import (
+    MultiFactorAuthModule,
+    MULTI_FACTOR_AUTH_MODULES,
+    MULTI_FACTOR_AUTH_MODULE_SCHEMA,
+    SetupFlow,
+)
 
-REQUIREMENTS = ['pyotp==2.2.7', 'PyQRCode==1.2.1']
+REQUIREMENTS = ["pyotp==2.2.7", "PyQRCode==1.2.1"]
 
-CONFIG_SCHEMA = MULTI_FACTOR_AUTH_MODULE_SCHEMA.extend({
-}, extra=vol.PREVENT_EXTRA)
+CONFIG_SCHEMA = MULTI_FACTOR_AUTH_MODULE_SCHEMA.extend({}, extra=vol.PREVENT_EXTRA)
 
 STORAGE_VERSION = 1
-STORAGE_KEY = 'auth_module.totp'
-STORAGE_USERS = 'users'
-STORAGE_USER_ID = 'user_id'
-STORAGE_OTA_SECRET = 'ota_secret'
+STORAGE_KEY = "auth_module.totp"
+STORAGE_USERS = "users"
+STORAGE_USER_ID = "user_id"
+STORAGE_OTA_SECRET = "ota_secret"
 
-INPUT_FIELD_CODE = 'code'
+INPUT_FIELD_CODE = "code"
 
-DUMMY_SECRET = 'FPPTH34D4E3MI2HG'
+DUMMY_SECRET = "FPPTH34D4E3MI2HG"
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -38,10 +41,15 @@ def _generate_qr_code(data: str) -> str:
 
     with BytesIO() as buffer:
         qr_code.svg(file=buffer, scale=4)
-        return '{}'.format(
-            buffer.getvalue().decode("ascii").replace('\n', '')
-            .replace('<?xml version="1.0" encoding="UTF-8"?>'
-                     '<svg xmlns="http://www.w3.org/2000/svg"', '<svg')
+        return "{}".format(
+            buffer.getvalue()
+            .decode("ascii")
+            .replace("\n", "")
+            .replace(
+                '<?xml version="1.0" encoding="UTF-8"?>'
+                '<svg xmlns="http://www.w3.org/2000/svg"',
+                "<svg",
+            )
         )
 
 
@@ -51,16 +59,17 @@ def _generate_secret_and_qr_code(username: str) -> Tuple[str, str, str]:
 
     ota_secret = pyotp.random_base32()
     url = pyotp.totp.TOTP(ota_secret).provisioning_uri(
-        username, issuer_name="Home Assistant")
+        username, issuer_name="Home Assistant"
+    )
     image = _generate_qr_code(url)
     return ota_secret, url, image
 
 
-@MULTI_FACTOR_AUTH_MODULES.register('totp')
+@MULTI_FACTOR_AUTH_MODULES.register("totp")
 class TotpAuthModule(MultiFactorAuthModule):
     """Auth module validate time-based one time password."""
 
-    DEFAULT_TITLE = 'Time-based One Time Password'
+    DEFAULT_TITLE = "Time-based One Time Password"
     MAX_RETRY_TIME = 5
 
     def __init__(self, hass: HomeAssistant, config: Dict[str, Any]) -> None:
@@ -68,7 +77,8 @@ class TotpAuthModule(MultiFactorAuthModule):
         super().__init__(hass, config)
         self._users = None  # type: Optional[Dict[str, str]]
         self._user_store = hass.helpers.storage.Store(
-            STORAGE_VERSION, STORAGE_KEY, private=True)
+            STORAGE_VERSION, STORAGE_KEY, private=True
+        )
         self._init_lock = asyncio.Lock()
 
     @property
@@ -93,14 +103,13 @@ class TotpAuthModule(MultiFactorAuthModule):
         """Save data."""
         await self._user_store.async_save({STORAGE_USERS: self._users})
 
-    def _add_ota_secret(self, user_id: str,
-                        secret: Optional[str] = None) -> str:
+    def _add_ota_secret(self, user_id: str, secret: Optional[str] = None) -> str:
         """Create a ota_secret for user."""
         import pyotp
 
         ota_secret = secret or pyotp.random_base32()  # type: str
 
-        self._users[user_id] = ota_secret   # type: ignore
+        self._users[user_id] = ota_secret  # type: ignore
         return ota_secret
 
     async def async_setup_flow(self, user_id: str) -> SetupFlow:
@@ -108,7 +117,7 @@ class TotpAuthModule(MultiFactorAuthModule):
 
         Mfa module should extend SetupFlow
         """
-        user = await self.hass.auth.async_get_user(user_id)   # type: ignore
+        user = await self.hass.auth.async_get_user(user_id)  # type: ignore
         return TotpSetupFlow(self, self.input_schema, user)
 
     async def async_setup_user(self, user_id: str, setup_data: Any) -> str:
@@ -117,7 +126,8 @@ class TotpAuthModule(MultiFactorAuthModule):
             await self._async_load()
 
         result = await self.hass.async_add_executor_job(
-            self._add_ota_secret, user_id, setup_data.get('secret'))
+            self._add_ota_secret, user_id, setup_data.get("secret")
+        )
 
         await self._async_save()
         return result
@@ -127,7 +137,7 @@ class TotpAuthModule(MultiFactorAuthModule):
         if self._users is None:
             await self._async_load()
 
-        if self._users.pop(user_id, None):   # type: ignore
+        if self._users.pop(user_id, None):  # type: ignore
             await self._async_save()
 
     async def async_is_user_setup(self, user_id: str) -> bool:
@@ -135,10 +145,9 @@ class TotpAuthModule(MultiFactorAuthModule):
         if self._users is None:
             await self._async_load()
 
-        return user_id in self._users   # type: ignore
+        return user_id in self._users  # type: ignore
 
-    async def async_validate(
-            self, user_id: str, user_input: Dict[str, Any]) -> bool:
+    async def async_validate(self, user_id: str, user_input: Dict[str, Any]) -> bool:
         """Return True if validation passed."""
         if self._users is None:
             await self._async_load()
@@ -146,7 +155,8 @@ class TotpAuthModule(MultiFactorAuthModule):
         # user_input has been validate in caller
         # set INPUT_FIELD_CODE as vol.Required is not user friendly
         return await self.hass.async_add_executor_job(
-            self._validate_2fa, user_id, user_input.get(INPUT_FIELD_CODE, ''))
+            self._validate_2fa, user_id, user_input.get(INPUT_FIELD_CODE, "")
+        )
 
     def _validate_2fa(self, user_id: str, code: str) -> bool:
         """Validate two factor authentication code."""
@@ -165,9 +175,9 @@ class TotpAuthModule(MultiFactorAuthModule):
 class TotpSetupFlow(SetupFlow):
     """Handler for the setup flow."""
 
-    def __init__(self, auth_module: TotpAuthModule,
-                 setup_schema: vol.Schema,
-                 user: User) -> None:
+    def __init__(
+        self, auth_module: TotpAuthModule, setup_schema: vol.Schema, user: User
+    ) -> None:
         """Initialize the setup flow."""
         super().__init__(auth_module, setup_schema, user.id)
         # to fix typing complaint
@@ -178,8 +188,8 @@ class TotpSetupFlow(SetupFlow):
         self._image = None  # type Optional[str]
 
     async def async_step_init(
-            self, user_input: Optional[Dict[str, str]] = None) \
-            -> Dict[str, Any]:
+        self, user_input: Optional[Dict[str, str]] = None
+    ) -> Dict[str, Any]:
         """Handle the first step of setup flow.
 
         Return self.async_show_form(step_id='init') if user_input is None.
@@ -191,30 +201,31 @@ class TotpSetupFlow(SetupFlow):
 
         if user_input:
             verified = await self.hass.async_add_executor_job(  # type: ignore
-                pyotp.TOTP(self._ota_secret).verify, user_input['code'])
+                pyotp.TOTP(self._ota_secret).verify, user_input["code"]
+            )
             if verified:
                 result = await self._auth_module.async_setup_user(
-                    self._user_id, {'secret': self._ota_secret})
+                    self._user_id, {"secret": self._ota_secret}
+                )
                 return self.async_create_entry(
-                    title=self._auth_module.name,
-                    data={'result': result}
+                    title=self._auth_module.name, data={"result": result}
                 )
 
-            errors['base'] = 'invalid_code'
+            errors["base"] = "invalid_code"
 
         else:
             hass = self._auth_module.hass
-            self._ota_secret, self._url, self._image = \
-                await hass.async_add_executor_job(  # type: ignore
-                    _generate_secret_and_qr_code, str(self._user.name))
+            self._ota_secret, self._url, self._image = await hass.async_add_executor_job(  # type: ignore
+                _generate_secret_and_qr_code, str(self._user.name)
+            )
 
         return self.async_show_form(
-            step_id='init',
+            step_id="init",
             data_schema=self._setup_schema,
             description_placeholders={
-                'code': self._ota_secret,
-                'url': self._url,
-                'qr_code': self._image
+                "code": self._ota_secret,
+                "url": self._url,
+                "qr_code": self._image,
             },
-            errors=errors
+            errors=errors,
         )