From bb7d262ea39473a42d8c7e08b14521ca33400937 Mon Sep 17 00:00:00 2001
From: Edward <73746306+WangEdward@users.noreply.github.com>
Date: Thu, 15 Feb 2024 16:39:14 +0000
Subject: [PATCH] feat: m-team 2fa

---
 app/helper/cookie.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/app/helper/cookie.py b/app/helper/cookie.py
index 8187256e..150f7e91 100644
--- a/app/helper/cookie.py
+++ b/app/helper/cookie.py
@@ -52,7 +52,8 @@ class CookieHelper:
         ],
         "twostep": [
             '//input[@name="two_step_code"]',
-            '//input[@name="2fa_secret"]'
+            '//input[@name="2fa_secret"]',
+            '//input[@name="otp"]'
         ]
     }
 
@@ -179,6 +180,24 @@ class CookieHelper:
             except Exception as e:
                 logger.error(f"仿真登录失败：{str(e)}")
                 return None, None, f"仿真登录失败：{str(e)}"
+            # 对于某二次验证码为单页面的站点，输入二次验证码
+            if "verify" in page.url:
+                if not otp_code:
+                    return None, None, "需要二次验证码"
+                html = etree.HTML(page.content())
+                for xpath in self._SITE_LOGIN_XPATH.get("twostep"):
+                    if html.xpath(xpath):
+                        try:
+                            # 刷新一下 2fa code
+                            otp_code = TwoFactorAuth(two_step_code).get_code()
+                            page.fill(xpath, otp_code)
+                            # 登录按钮 xpath 理论上相同，不再重复查找
+                            page.click(submit_xpath)
+                            page.wait_for_load_state("networkidle", timeout=30 * 1000)
+                        except Exception as e:
+                            logger.error(f"二次验证码输入失败：{str(e)}")
+                            return None, None, f"二次验证码输入失败：{str(e)}"
+                        break
             # 登录后的源码
             html_text = page.content()
             if not html_text: