From 8beb2cf442d41f75bc9cf759b9d036105dd4525c Mon Sep 17 00:00:00 2001
From: Karako <70872201+karakoo@users.noreply.github.com>
Date: Thu, 8 Dec 2022 22:22:59 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix=20cookie=20parser?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 plugins/genshin/cookies.py | 74 +++++++++++++-------------------------
 1 file changed, 25 insertions(+), 49 deletions(-)

diff --git a/plugins/genshin/cookies.py b/plugins/genshin/cookies.py
index a510542..a2f9b9c 100644
--- a/plugins/genshin/cookies.py
+++ b/plugins/genshin/cookies.py
@@ -1,9 +1,8 @@
 import contextlib
-import re
-from http.cookies import CookieError, SimpleCookie
 from typing import Dict, Optional
 
 import genshin
+from arkowrapper import ArkoWrapper
 from genshin import DataNotPublic, GenshinException, InvalidCookies, types
 from genshin.models import GenshinAccount
 from telegram import ReplyKeyboardMarkup, ReplyKeyboardRemove, TelegramObject, Update
@@ -45,43 +44,26 @@ class SetUserCookies(Plugin.Conversation, BasePlugin.Conversation):
         self.cookies_service = cookies_service
         self.user_service = user_service
 
+    # noinspection SpellCheckingInspection
     @staticmethod
-    def parse_cookie(cookie: SimpleCookie) -> Dict[str, str]:
+    def parse_cookie(cookie: Dict[str, str]) -> Dict[str, str]:
         cookies = {}
-        ltoken = cookie.get("ltoken")
-        if ltoken:
-            cookies["ltoken"] = ltoken.value
-        ltuid = cookie.get("ltuid")
-        login_uid = cookie.get("login_uid")
-        if ltuid:
-            cookies["ltuid"] = ltuid.value
-            cookies["account_id"] = ltuid.value
-        if login_uid:
-            cookies["ltuid"] = login_uid.value
-            if ltuid:
-                cookies["account_id"] = ltuid.value
-        cookie_token = cookie.get("cookie_token")
-        cookie_token_v2 = cookie.get("cookie_token_v2")
-        if cookie_token:
-            cookies["cookie_token"] = cookie_token.value
-        if cookie_token_v2:
-            cookies["cookie_token"] = cookie_token_v2.value
-        account_mid_v2 = cookie.get("account_mid_v2")
-        if account_mid_v2:
-            cookies["account_mid_v2"] = account_mid_v2.value
-        cookie_token_v2 = cookie.get("cookie_token_v2")
-        if cookie_token_v2:
-            cookies["cookie_token_v2"] = cookie_token_v2.value
-        ltoken_v2 = cookie.get("ltoken_v2")
-        if ltoken_v2:
-            cookies["ltoken_v2"] = ltoken_v2.value
-        ltmid_v2 = cookie.get("ltmid_v2")
-        if ltmid_v2:
-            cookies["ltmid_v2"] = ltmid_v2.value
-        login_ticket = cookie.get("login_ticket")
-        if login_ticket:
-            cookies["login_ticket"] = login_ticket.value
-        return cookies
+
+        v1_keys = ["ltoken", "ltuid", "login_uid", "cookie_token"]
+        v2_keys = ["ltoken_v2", "ltmid_v2", "account_mid_v2", "cookie_token", "cookie_token_v2", "login_ticket"]
+
+        cookie_is_v1 = None
+
+        for k in v1_keys + v2_keys:
+            v = cookie.get(k)
+            if v is not None and cookie_is_v1 is None:
+                cookie_is_v1 = k not in v2_keys
+            cookies[k] = cookie.get(k)
+
+        if cookie_is_v1:
+            cookies["account_id"] = cookies["ltuid"]
+
+        return {k: v for k, v in cookies.items() if v is not None}
 
     @conversation.entry_point
     @handler.command(command="setcookie", filters=filters.ChatType.PRIVATE, block=True)
@@ -278,19 +260,13 @@ class SetUserCookies(Plugin.Conversation, BasePlugin.Conversation):
         if message.text == "退出":
             await message.reply_text("退出任务", reply_markup=ReplyKeyboardRemove())
             return ConversationHandler.END
-        str_cookies = re.sub(r"(\s*\S*?=\{.*?};?\s*)", " ", message.text, 0)
 
-        cookie = SimpleCookie()
-        try:
-            cookie.load(str_cookies)
-        except CookieError:
-            logger.info("用户 %s[%s] Cookies格式有误", user.full_name, user.id)
-            await message.reply_text("Cookies格式有误，请检查", reply_markup=ReplyKeyboardRemove())
-            return ConversationHandler.END
-        if len(cookie) == 0:
-            logger.info("用户 %s[%s] Cookies格式有误", user.full_name, user.id)
-            await message.reply_text("Cookies格式有误，请检查", reply_markup=ReplyKeyboardRemove())
-            return ConversationHandler.END
+        # cookie str to dict
+        wrapped = (
+            ArkoWrapper(message.text.split(";")).map(lambda x: x.strip()).map(lambda x: ((y := x.split("="))[0], y[1]))
+        )
+        cookie = {x[0]: x[1] for x in wrapped}
+
         try:
             cookies = self.parse_cookie(cookie)
         except (AttributeError, ValueError) as exc: