feat(app):

- add email service
- fix security problem

src/retk/config.py

@@ -19,7 +19,6 @@ class Settings(BaseSettings):
     VERIFY_REFERER: bool = Field(default=False, env='VERIFY_REFERER')
     PLUGINS: bool = Field(default=False, env='PLUGINS')
     RETHINK_LOCAL_STORAGE_PATH: Optional[DirectoryPath] = Field(env='RETHINK_LOCAL_STORAGE_PATH', default=None)
-    CAPTCHA_SALT: str = Field(env='CAPTCHA_SALT', default="")
     MD_BACKUP_INTERVAL: int = Field(env="MD_BACKUP_INTERVAL", default=60 * 5)  # 5 minutes
 
     # database settings: MongoDB
@@ -85,7 +84,8 @@ class Settings(BaseSettings):
 
     # Email client settings
     RETHINK_EMAIL: str = Field(env='RETHINK_EMAIL', default="")
-    RETHINK_EMAIL_PASSWORD: str = Field(env='RETHINK_EMAIL_PASSWORD', default="")
+    RETHINK_EMAIL_SECRET_ID: str = Field(env='RETHINK_EMAIL_SECRET_ID', default="")
+    RETHINK_EMAIL_SECRET_KEY: str = Field(env='RETHINK_EMAIL_SECRET_KEY', default="")
 
     # OAuth settings
     OAUTH_REDIRECT_URL: str = Field(env='OAUTH_REDIRECT_URL', default="")

src/retk/controllers/account.py

@@ -74,7 +74,7 @@ async def signup(
 ) -> JSONResponse:
     if not const.LanguageEnum.is_valid(req.language):
         req.language = const.LanguageEnum.EN.value
-    code = account.app_captcha.verify_captcha(token=req.verificationToken, code_str=req.verification)
+    code = account.email.verify_number(cid=req.verificationToken, user_code=req.verification)
     if code != const.CodeEnum.OK:
         raise json_exception(
             request_id=req_id,
@@ -221,7 +221,7 @@ async def forget(
         req_id: str,
         req: schemas.account.ForgetPasswordRequest
 ) -> schemas.RequestIdResponse:
-    code = account.email.verify_number(token=req.verificationToken, number_str=req.verification)
+    code = account.email.verify_number(cid=req.verificationToken, user_code=req.verification)
     if code != const.CodeEnum.OK:
         raise json_exception(
             request_id=req_id,
@@ -263,24 +263,24 @@ async def forget(
 
 
 def get_captcha_img():
-    token, data = account.app_captcha.generate(length=4, sound=False)
+    cid, data = account.app_captcha.generate(length=4, sound=False)
     return StreamingResponse(
         data["img"],
         headers={
-            "X-Captcha-Token": token
+            "X-Captcha-Token": cid
         },
         media_type="image/png",
     )
 
 
 def __check_and_send_email(
         email: str,
-        token: str,
+        cid: str,
         code_str: str,
         language: str,
 ) -> Tuple[str, int, const.CodeEnum]:
     # verify captcha code in image
-    code = account.app_captcha.verify_captcha(token=token, code_str=code_str)
+    code = account.app_captcha.verify_captcha(cid=cid, code_str=code_str)
 
     if code != const.CodeEnum.OK:
         return "", 0, code
@@ -315,7 +315,7 @@ async def email_send_code(
 
     numbers, expired_min, code = __check_and_send_email(
         email=req.email,
-        token=req.captchaToken,
+        cid=req.captchaToken,
         code_str=req.captchaCode,
         language=req.language,
     )
@@ -327,10 +327,10 @@ async def email_send_code(
             language=req.language,
         )
 
-    token = account.email.encode_number(number=numbers, expired_min=expired_min)
+    cid = account.email.encode_number(number=numbers, expired_min=expired_min)
     return schemas.account.TokenResponse(
         requestId=req_id,
-        token=token,
+        token=cid,
     )
 
 

src/retk/core/account/app_captcha.py

@@ -1,14 +1,14 @@
-from datetime import timedelta
+from collections import OrderedDict
+from datetime import datetime
 from io import BytesIO
 from random import choices
 from typing import Tuple, Dict
 
-import jwt
 from captcha.audio import AudioCaptcha
 from captcha.image import ImageCaptcha
 
-from retk import const, config
-from retk.utils import jwt_encode, jwt_decode
+from retk import const
+from retk.core.utils import cached_verification
 
 DEFAULT_CAPTCHA_EXPIRE_SECOND = 60
 
@@ -19,6 +19,8 @@
 alphabet_len = len(alphabet)
 code_idx_range = list(range(0, alphabet_len - 1))
 
+cache_captcha: OrderedDict[str, Tuple[datetime, str]] = OrderedDict()
+
 
 def generate(length: int = 4, sound: bool = False) -> Tuple[str, Dict[str, BytesIO]]:
     code = [alphabet[i] for i in choices(code_idx_range, k=length)]
@@ -28,22 +30,17 @@ def generate(length: int = 4, sound: bool = False) -> Tuple[str, Dict[str, Bytes
     }
     if sound:
         data["sound"] = audio_captcha.generate(code_str)
-    token = jwt_encode(
-        exp_delta=timedelta(seconds=DEFAULT_CAPTCHA_EXPIRE_SECOND),
-        data={"code": code_str.lower() + config.get_settings().CAPTCHA_SALT}
+    cid = cached_verification.add_to_cache(
+        cached=cache_captcha,
+        code=code_str.lower(),
+        expired_seconds=DEFAULT_CAPTCHA_EXPIRE_SECOND
+    )
+    return cid, data
+
+
+def verify_captcha(cid: str, code_str: str) -> const.CodeEnum:
+    return cached_verification.verify_captcha(
+        cached=cache_captcha,
+        cid=cid,
+        user_code=code_str
     )
-    # TODO: 以后有 redis 之后要考虑 token 的一次性校验。校验过一次后，需要主动失效
-    return token, data
-
-
-def verify_captcha(token: str, code_str: str) -> const.CodeEnum:
-    code = const.CodeEnum.CAPTCHA_ERROR
-    try:
-        data = jwt_decode(token)
-        if data["code"] == code_str.lower() + config.get_settings().CAPTCHA_SALT:
-            code = const.CodeEnum.OK
-    except jwt.ExpiredSignatureError:
-        code = const.CodeEnum.CAPTCHA_EXPIRED
-    except (jwt.DecodeError, Exception):  # pylint: disable=broad-except
-        code = const.CodeEnum.INVALID_AUTH
-    return code

src/retk/core/account/email.py

@@ -1,14 +1,10 @@
-import email.header
-from datetime import timedelta
-from email.mime.multipart import MIMEMultipart
-from email.mime.text import MIMEText
-from textwrap import dedent
-from typing import List, Tuple
-
-import jwt
+from collections import OrderedDict
+from datetime import datetime
+from typing import List, Dict, Tuple
 
 from retk import const, config, regex, utils
 from retk.core import scheduler
+from retk.core.utils import cached_verification
 
 
 class EmailServer:
@@ -18,55 +14,24 @@ class EmailServer:
         const.LanguageEnum.EN.value: "Rethink: Security Code",
         const.LanguageEnum.ZH.value: "Rethink: 安全密码",
     }
-    lang_content = {
-        const.LanguageEnum.EN.value: dedent("""\
-        Please use the following security code for your Rethink account {email}:
-        <br><br>
-        Security Code: <br><br>
-        <strong style="font-size:26px;">{numbers}</strong>
-        <br><br>
-        Valid for {expire} minutes, please do not tell others to prevent personal information leakage.
-        <br><br>
-        If you did not request this code, you can safely ignore this email, \
-        someone may have entered your email address by mistake.
-        <br><br>
-        Thank you!<br>
-        Rethink Team
-        """),
-        const.LanguageEnum.ZH.value: dedent("""\
-        请使用以下用于 Rethink 账户 {email} 的安全代码：
-        <br><br>
-        安全代码：<br><br>
-        <strong style="font-size:26px;">{numbers}</strong>
-        <br><br>
-        有效期 {expire} 分钟，请勿告知他人，以防个人信息泄露。
-        <br><br>
-        若您并未要求此代码，可以安全地忽视此邮件，可能有人误输入了您的电子邮件地址。
-        <br><br>
-        谢谢！<br>
-        Rethink 团队
-        """),
+    lang_template_id = {
+        const.LanguageEnum.EN.value: 127743,
+        const.LanguageEnum.ZH.value: 127744,
     }
 
-    def get_subject_content(self, recipient: str, numbers: str, expire: int, language: str) -> Tuple[str, str]:
+    def send(self, recipient: str, numbers: str, expire: int, language: str) -> const.CodeEnum:
         try:
             subject = self.lang_subject[language]
-            content_temp = self.lang_content[language]
+            template_id = self.lang_template_id[language]
         except KeyError:
             subject = self.lang_subject[self.default_language]
-            content_temp = self.lang_content[self.default_language]
-        content = content_temp.format(email=utils.mask_email(recipient), numbers=numbers, expire=expire)
-        return subject, content
-
-    def send(self, recipient: str, numbers: str, expire: int, language: str) -> const.CodeEnum:
-        subject, content = self.get_subject_content(
-            recipient=recipient, numbers=numbers, expire=expire, language=language
-        )
+            template_id = self.lang_template_id[self.default_language]
 
         return self._send(
-            subject=subject,
             recipients=[recipient],
-            html_message=content
+            subject=subject,
+            template_id=template_id,
+            values={"email": utils.mask_email(recipient), "numbers": numbers, "expire": expire},
         )
 
     @staticmethod
@@ -75,45 +40,44 @@ def email_ok(email_addr: str) -> bool:
             return False
         return True
 
-    def _send(self, recipients: List[str], subject: str, html_message: str) -> const.CodeEnum:
+    def _send(self, recipients: List[str], subject: str, template_id: int, values: Dict) -> const.CodeEnum:
         for recipient in recipients:
             if not self.email_ok(recipient):
                 return const.CodeEnum.INVALID_EMAIL
         conf = config.get_settings()
-        msg = MIMEMultipart('alternative')
-        msg['Subject'] = email.header.Header(subject, 'utf-8')
-        msg['From'] = conf.RETHINK_EMAIL
-        msg['To'] = ", ".join(recipients)
-        html_body = MIMEText(html_message, 'html', 'utf-8')
-        msg.attach(html_body)
 
         _, code = scheduler.run_once_now(
-            job_id=f"send_email_{conf.RETHINK_EMAIL}_{recipients}_{subject}_{html_message}",
-            func=scheduler.tasks.email.send,
-            kwargs={"recipients": recipients, "subject": msg.as_string()},
+            job_id=f"send_email_{conf.RETHINK_EMAIL}_{recipients}_{subject}_{template_id}_{values}",
+            func=scheduler.tasks.email.send_verification_code,
+            kwargs={
+                "from_email": f"Rethink <{conf.RETHINK_EMAIL}>",
+                "to_emails": recipients,
+                "subject": subject,
+                "values": values,
+                "template_id": template_id,
+                "secret_id": conf.RETHINK_EMAIL_SECRET_ID,
+                "secret_key": conf.RETHINK_EMAIL_SECRET_KEY,
+            },
         )
         return code
 
 
 email_server = EmailServer()
 
+cache_email: OrderedDict[str, Tuple[datetime, str]] = OrderedDict()
+
 
 def encode_number(number: str, expired_min: int) -> str:
-    token = utils.jwt_encode(
-        exp_delta=timedelta(minutes=expired_min),
-        data={"code": number + config.get_settings().CAPTCHA_SALT}
+    return cached_verification.add_to_cache(
+        cached=cache_email,
+        code=number,
+        expired_seconds=expired_min * 60
+    )
+
+
+def verify_number(cid: str, user_code: str) -> const.CodeEnum:
+    return cached_verification.verify_captcha(
+        cached=cache_email,
+        cid=cid,
+        user_code=user_code
     )
-    return token
-
-
-def verify_number(token: str, number_str: str) -> const.CodeEnum:
-    code = const.CodeEnum.CAPTCHA_ERROR
-    try:
-        data = utils.jwt_decode(token)
-        if data["code"] == number_str + config.get_settings().CAPTCHA_SALT:
-            code = const.CodeEnum.OK
-    except jwt.ExpiredSignatureError:
-        code = const.CodeEnum.CAPTCHA_EXPIRED
-    except (jwt.DecodeError, Exception):  # pylint: disable=broad-except
-        code = const.CodeEnum.INVALID_AUTH
-    return code

src/retk/core/account/emailVerificationEn.html

@@ -0,0 +1,12 @@
+Please use the following security code for your Rethink account {{email}}:
+<br><br>
+Security Code: <br><br>
+<strong style="font-size:26px;">{{numbers}}</strong>
+<br><br>
+Valid for {{expire}} minutes, please do not tell others to prevent personal information leakage.
+<br><br>
+If you did not request this code, you can safely ignore this email,
+someone may have entered your email address by mistake.
+<br><br>
+Thank you!<br>
+Rethink Team

src/retk/core/account/emailVerificationZh.html

@@ -0,0 +1,11 @@
+请使用以下用于 Rethink 账户 {{email}} 的安全代码：
+<br><br>
+安全代码：<br><br>
+<strong style="font-size:26px;">{{numbers}}</strong>
+<br><br>
+有效期 {{expire}} 分钟，请勿告知他人，以防个人信息泄露。
+<br><br>
+若您并未要求此代码，可以安全地忽视此邮件，可能有人误输入了您的电子邮件地址。
+<br><br>
+谢谢！<br>
+Rethink 团队

src/retk/core/account/manager.py

@@ -90,6 +90,8 @@ async def __delete_post_process(uid: str):
     await client.coll.import_data.delete_many({"uid": uid})
     await client.coll.notice_system.delete_many({"recipientId": uid})
     await client.search.force_delete_all(uid=uid)
+    await client.coll.llm_extend_node_queue.delete_many({"uid": uid})
+    await client.coll.llm_extended_node.delete_many({"uid": uid})
 
 
 async def delete_by_uid(uid: str):