Merge pull request #30 from lyylyylyylyy/master

Fix login with password needing vertification code

custom_components/china_southern_power_grid_stat/config_flow.py

@@ -8,6 +8,7 @@
 """
 from __future__ import annotations
 
+import asyncio
 import logging
 import time
 from typing import Any
@@ -27,6 +28,7 @@
     ABORT_NO_ACCOUNT_TO_DELETE,
     CONF_ACCOUNT_NUMBER,
     CONF_ACTION,
+    CONF_VERIFICATION_CODE,
     CONF_AUTH_TOKEN,
     CONF_ELE_ACCOUNTS,
     CONF_GENERAL_ERROR,
@@ -50,12 +52,11 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-
-def authenticate_csg(username: str, password: str) -> CSGClient:
+def authenticate_csg(username: str, password: str, code: str) -> CSGClient:
     """use username and password combination to authenticate"""
     client = CSGClient()
     try:
-        client.authenticate(username, password)
+        client.authenticate(username, password, code)
     except InvalidCredentials as exc:
         _LOGGER.error("Authentication failed: %s", exc)
         raise InvalidAuth from exc
@@ -70,11 +71,13 @@ async def validate_input(
     """Validate the credentials (login)"""
 
     client = await hass.async_add_executor_job(
-        authenticate_csg, data[CONF_USERNAME], data[CONF_PASSWORD]
+        authenticate_csg, data[CONF_USERNAME], data[CONF_PASSWORD], data[CONF_VERIFICATION_CODE]
     )
     return client.dump()
 
 
+
+
 class CSGConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow for China Southern Power Grid Statistics."""
 
@@ -88,6 +91,7 @@ def async_get_options_flow(
     ) -> config_entries.OptionsFlow:
         """Create the options flow."""
         return CSGOptionsFlowHandler(config_entry)
+
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
@@ -106,52 +110,78 @@ async def async_step_user(
         if user_input is None:
             return self.async_show_form(step_id=STEP_USER, data_schema=schema)
 
-        errors = {}
-        unique_id = f"CSG-{user_input[CONF_USERNAME]}"
-        await self.async_set_unique_id(unique_id)
-        self._abort_if_unique_id_configured()
+        # Initialize the client
+        client = CSGClient()
 
-        # noinspection PyBroadException
+        errors = {}
         try:
-            session_data = await validate_input(self.hass, user_input)
-        except CannotConnect:
+            # Send SMS verification code after receiving username and password
+            await asyncio.to_thread(client.api_send_login_sms, user_input[CONF_USERNAME])
+
+            # Save username and password for the next step (verification)
+            self.context["user_data"] = {
+                CONF_USERNAME: user_input[CONF_USERNAME],
+                CONF_PASSWORD: user_input[CONF_PASSWORD],
+            }
+
+            # Move to the verification step
+            return await self.async_step_verification()
+
+        except RequestException:
             errors[CONF_GENERAL_ERROR] = ERROR_CANNOT_CONNECT
         except InvalidAuth:
             errors[CONF_GENERAL_ERROR] = ERROR_INVALID_AUTH
-        except Exception:  # pylint: disable=broad-except
-            _LOGGER.exception("Unexpected exception")
+        except Exception:
+            _LOGGER.exception("Unexpected exception during login")
             errors[CONF_GENERAL_ERROR] = ERROR_UNKNOWN
-        else:
-            if self.reauth_entry:
-                new_data = self.reauth_entry.data.copy()
-                if user_input[CONF_USERNAME] != new_data[CONF_USERNAME]:
-                    _LOGGER.warning(
-                        "Account name changed: previous: %s, now: %s",
-                        new_data[CONF_USERNAME],
-                        user_input[CONF_USERNAME],
-                    )
-                new_data[CONF_USERNAME] = user_input[CONF_USERNAME]
-                new_data[CONF_PASSWORD] = user_input[CONF_PASSWORD]
-                new_data[CONF_AUTH_TOKEN] = session_data[CONF_AUTH_TOKEN]
-                new_data[CONF_UPDATED_AT] = str(int(time.time() * 1000))
-                self.hass.config_entries.async_update_entry(
-                    self.reauth_entry,
-                    data=new_data,
-                    title=f"CSG-{user_input[CONF_USERNAME]}",
-                )
-                await self.hass.config_entries.async_reload(self.reauth_entry.entry_id)
-                _LOGGER.info(
-                    "Reauth of account %s is successful!", user_input[CONF_USERNAME]
-                )
-                return self.async_abort(reason="reauth_successful")
-
-            _LOGGER.info("Adding csg account %s", user_input[CONF_USERNAME])
+
+        return self.async_show_form(step_id=STEP_USER, data_schema=schema, errors=errors)
+
+    async def async_step_verification(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """
+        Second step: input verification code and verify it with the server.
+        """
+        schema = vol.Schema(
+            {
+                vol.Required(CONF_VERIFICATION_CODE): str,
+            }
+        )
+
+        if user_input is None:
+            # First time in this step
+            return self.async_show_form(step_id="verification", data_schema=schema)
+
+        # Retrieve saved username and password from the context
+        username = self.context["user_data"][CONF_USERNAME]
+        password = self.context["user_data"][CONF_PASSWORD]
+        verification_code = user_input[CONF_VERIFICATION_CODE]
+
+        self.context["user_data"] = {
+            CONF_USERNAME: username,
+            CONF_PASSWORD: password,
+            CONF_VERIFICATION_CODE: verification_code,
+        }
+
+        client = CSGClient()
+
+        errors = {}
+        try:
+            # Attempt to authenticate using the verification code
+            session_data = await self.hass.async_add_executor_job(
+                client.api_login_with_password, username, password, verification_code
+            )
+
+            _LOGGER.debug("session data: %s", session_data)
+
+            # Successful authentication, create the entry
             return self.async_create_entry(
-                title=f"CSG-{user_input[CONF_USERNAME]}",
+                title=f"CSG-{username}",
                 data={
-                    CONF_USERNAME: user_input[CONF_USERNAME],
-                    CONF_PASSWORD: user_input[CONF_PASSWORD],
-                    CONF_AUTH_TOKEN: session_data[CONF_AUTH_TOKEN],
+                    CONF_USERNAME: username,
+                    CONF_PASSWORD: password,
+                    CONF_AUTH_TOKEN: session_data,
                     CONF_ELE_ACCOUNTS: {},
                     CONF_SETTINGS: {
                         CONF_UPDATE_INTERVAL: DEFAULT_UPDATE_INTERVAL,
@@ -160,9 +190,15 @@ async def async_step_user(
                 },
             )
 
-        return self.async_show_form(
-            step_id=STEP_USER, data_schema=schema, errors=errors
-        )
+        except CannotConnect:
+            errors[CONF_GENERAL_ERROR] = ERROR_CANNOT_CONNECT
+        except InvalidAuth:
+            errors[CONF_GENERAL_ERROR] = ERROR_INVALID_AUTH
+        except Exception:
+            _LOGGER.exception("Unexpected exception during verification")
+            errors[CONF_GENERAL_ERROR] = ERROR_UNKNOWN
+
+        return self.async_show_form(step_id="verification", data_schema=schema, errors=errors)
 
     async def async_step_reauth(self, user_input=None):
         """Perform reauth upon an API authentication error."""
@@ -267,6 +303,7 @@ async def async_step_add_account(
                 client.authenticate,
                 self.config_entry.data[CONF_USERNAME],
                 self.config_entry.data[CONF_PASSWORD],
+                self.config_entry.data[CONF_VERIFICATION_CODE],
             )
         await self.hass.async_add_executor_job(client.initialize)
 

custom_components/china_southern_power_grid_stat/const.py

@@ -14,6 +14,7 @@
 CONF_SETTINGS = "settings"
 CONF_UPDATED_AT = "updated_at"
 CONF_ACTION = "action"
+CONF_VERIFICATION_CODE = "code"
 
 STEP_USER = "user"
 STEP_INIT = "init"

custom_components/china_southern_power_grid_stat/csg_client/__init__.py

@@ -238,7 +238,11 @@ def _make_request(
                     "API call %s returned status code %d", path, response.status_code
                 )
                 raise CSGHTTPError(response.status_code)
-            response_data = response.json()
+
+            json_str = response.content.decode('utf-8', errors='ignore')
+            json_str = json_str[json_str.find('{'):json_str.rfind('}')+1]
+            json_data = json.loads(json_str)
+            response_data = json_data
             _LOGGER.debug("_make_request: response: %s", json.dumps(response_data))
 
             # headers need to be returned since they may contain additional data
@@ -297,15 +301,17 @@ def api_login_with_sms_code(self, phone_no: str, code: str):
             return resp_header[HEADER_X_AUTH_TOKEN]
         self._handle_unsuccessful_response(path, resp_data)
 
-    def api_login_with_password(self, phone_no: str, password: str):
+    def api_login_with_password(self, phone_no: str, password: str, code: str):
         """Login with phone number and password"""
-        path = "center/login"
+        path = "center/loginByPwdAndMsg"
         payload = {
             JSON_KEY_AREA_CODE: AREACODE_FALLBACK,
             JSON_KEY_ACCT_ID: phone_no,
             JSON_KEY_LOGON_CHAN: LOGON_CHANNEL_HANDHELD_HALL,
             JSON_KEY_CRED_TYPE: LOGIN_TYPE_PHONE_PWD,
             "credentials": encrypt_credential(password),
+            "code": code,
+            "checkPwd": True
         }
         payload = {"param": encrypt_params(payload)}
         resp_header, resp_data = self._make_request(
@@ -509,12 +515,12 @@ def set_authentication_params(self, auth_token: str, login_type: LoginType):
         self.auth_token = auth_token
         self.login_type = login_type
 
-    def authenticate(self, phone_no: str, password: str):
+    def authenticate(self, phone_no: str, password: str, code: str):
         """
         Authenticate the client using phone number and password
         Will set session parameters
         """
-        auth_token = self.api_login_with_password(phone_no, password)
+        auth_token = self.api_login_with_password(phone_no, password, code)
         self.set_authentication_params(auth_token, LoginType.LOGIN_TYPE_PWD)
 
     def initialize(self):
@@ -693,3 +699,4 @@ def get_yesterday_kwh(self, account: CSGElectricityAccount) -> float:
         return float(resp_data["power"])
 
     # end high-level api wrappers
+

custom_components/china_southern_power_grid_stat/csg_client/const.py

@@ -39,7 +39,7 @@ class LoginType(Enum):
 # however they're not programmatically linked in the source code
 # use them as seperated parameters for now
 LOGIN_TYPE_PHONE_CODE = "11"
-LOGIN_TYPE_PHONE_PWD = "10"
+LOGIN_TYPE_PHONE_PWD = "1011"
 SEND_MSG_TYPE_VERIFICATION_CODE = "1"
 VERIFICATION_CODE_TYPE_LOGIN = "1"
 

custom_components/china_southern_power_grid_stat/strings.json

@@ -6,7 +6,9 @@
         "description": "使用手机号和密码登录",
         "data": {
           "username": "手机号",
-          "password": "密码"
+          "password": "密码",
+          "code": "验证码",
+          "send_code": "发送验证码"
         }
       },
       "reauth_confirm": {

custom_components/china_southern_power_grid_stat/translations/en.json

@@ -1,52 +1,54 @@
 {
-    "config": {
-        "abort": {
-            "already_configured": "\u8d26\u53f7\u5df2\u6dfb\u52a0\uff0c\u8bf7\u52ff\u91cd\u590d\u6dfb\u52a0",
-            "reauth_successful": "\u8d26\u53f7\u5bc6\u7801\u5237\u65b0\u6210\u529f"
-        },
-        "error": {
-            "cannot_connect": "\u7f51\u7edc\u5f02\u5e38",
-            "invalid_auth": "\u7528\u6237\u540d\u6216\u5bc6\u7801\u9519\u8bef",
-            "unknown": "\u672a\u77e5\u9519\u8bef"
+  "config": {
+    "abort": {
+      "already_configured": "\u8d26\u53f7\u5df2\u6dfb\u52a0\uff0c\u8bf7\u52ff\u91cd\u590d\u6dfb\u52a0",
+      "reauth_successful": "\u8d26\u53f7\u5bc6\u7801\u5237\u65b0\u6210\u529f"
+    },
+    "error": {
+      "cannot_connect": "\u7f51\u7edc\u5f02\u5e38",
+      "invalid_auth": "\u7528\u6237\u540d\u6216\u5bc6\u7801\u9519\u8bef",
+      "unknown": "\u672a\u77e5\u9519\u8bef"
+    },
+    "step": {
+      "reauth_confirm": {
+        "description": "\u81ea\u52a8\u767b\u9646\u5931\u8d25\uff08\u7528\u6237\u540d\u6216\u5bc6\u7801\u9519\u8bef\uff09\uff0c\u8bf7\u91cd\u65b0\u8f93\u5165",
+        "title": "\u91cd\u65b0\u767b\u9646\u8d26\u53f7"
+      },
+      "user": {
+        "data": {
+          "password": "\u5bc6\u7801",
+          "username": "\u624b\u673a\u53f7",
+          "code": "\u9a8c\u8bc1\u7801",
+          "send_code": "\u53d1\u9001\u9a8c\u8bc1\u7801"
         },
-        "step": {
-            "reauth_confirm": {
-                "description": "\u81ea\u52a8\u767b\u9646\u5931\u8d25\uff08\u7528\u6237\u540d\u6216\u5bc6\u7801\u9519\u8bef\uff09\uff0c\u8bf7\u91cd\u65b0\u8f93\u5165",
-                "title": "\u91cd\u65b0\u767b\u9646\u8d26\u53f7"
-            },
-            "user": {
-                "data": {
-                    "password": "\u5bc6\u7801",
-                    "username": "\u624b\u673a\u53f7"
-                },
-                "description": "\u4f7f\u7528\u624b\u673a\u53f7\u548c\u5bc6\u7801\u767b\u5f55",
-                "title": "\u767b\u9646\u5357\u65b9\u7535\u7f51"
-            }
-        }
+        "description": "\u4f7f\u7528\u624b\u673a\u53f7\u548c\u5bc6\u7801\u767b\u5f55",
+        "title": "\u767b\u9646\u5357\u65b9\u7535\u7f51"
+      }
+    }
+  },
+  "options": {
+    "abort": {
+      "all_added": "\u6ca1\u6709\u53ef\u6dfb\u52a0\u5230\u6b64\u96c6\u6210\u7684\u7f34\u8d39\u53f7",
+      "no_account": "\u8d26\u6237\u672a\u7ed1\u5b9a\u7f34\u8d39\u53f7",
+      "no_account_to_delete": "\u6b64\u96c6\u6210\u672a\u7ed1\u5b9a\u7f34\u8d39\u53f7"
     },
-    "options": {
-        "abort": {
-            "all_added": "\u6ca1\u6709\u53ef\u6dfb\u52a0\u5230\u6b64\u96c6\u6210\u7684\u7f34\u8d39\u53f7",
-            "no_account": "\u8d26\u6237\u672a\u7ed1\u5b9a\u7f34\u8d39\u53f7",
-            "no_account_to_delete": "\u6b64\u96c6\u6210\u672a\u7ed1\u5b9a\u7f34\u8d39\u53f7"
+    "step": {
+      "add_account": {
+        "title": "\u9009\u62e9\u7f34\u8d39\u53f7"
+      },
+      "init": {
+        "title": "\u9009\u62e9\u64cd\u4f5c"
+      },
+      "remove_account": {
+        "title": "\u79fb\u9664\u7f34\u8d39\u53f7"
+      },
+      "settings": {
+        "data": {
+          "update_interval": "\u66f4\u65b0\u95f4\u9694\uff08\u79d2\uff09",
+          "update_timeout": "\u8bf7\u6c42\u8d85\u65f6\uff08\u79d2\uff09"
         },
-        "step": {
-            "add_account": {
-                "title": "\u9009\u62e9\u7f34\u8d39\u53f7"
-            },
-            "init": {
-                "title": "\u9009\u62e9\u64cd\u4f5c"
-            },
-            "remove_account": {
-                "title": "\u79fb\u9664\u7f34\u8d39\u53f7"
-            },
-            "settings": {
-                "data": {
-                    "update_interval": "\u66f4\u65b0\u95f4\u9694\uff08\u79d2\uff09",
-                    "update_timeout": "\u8bf7\u6c42\u8d85\u65f6\uff08\u79d2\uff09"
-                },
-                "title": "\u53c2\u6570\u8bbe\u7f6e"
-            }
-        }
+        "title": "\u53c2\u6570\u8bbe\u7f6e"
+      }
     }
-}
+  }
+}

custom_components/china_southern_power_grid_stat/utils.py

@@ -7,7 +7,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
 
-from .const import CONF_AUTH_TOKEN, CONF_ELE_ACCOUNTS, CONF_UPDATED_AT
+from .const import CONF_AUTH_TOKEN, CONF_ELE_ACCOUNTS, CONF_UPDATED_AT, CONF_VERIFICATION_CODE
 from .csg_client import CSGClient, InvalidCredentials
 
 _LOGGER = logging.getLogger(__name__)
@@ -22,6 +22,7 @@ async def async_refresh_login_and_update_config(
             client.authenticate,
             entry.data[CONF_USERNAME],
             entry.data[CONF_PASSWORD],
+            entry.data[CONF_VERIFICATION_CODE],
         )
     except InvalidCredentials as err:
         raise ConfigEntryAuthFailed(str(err)) from err