Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>tags/1.8.0
| @@ -1,3 +1,5 @@ | |||
| from datetime import datetime | |||
| import pytz | |||
| from flask import request | |||
| from flask_login import current_user | |||
| @@ -327,6 +329,9 @@ class EducationVerifyApi(Resource): | |||
| class EducationApi(Resource): | |||
| status_fields = { | |||
| "result": fields.Boolean, | |||
| "is_student": fields.Boolean, | |||
| "expire_at": TimestampField, | |||
| "allow_refresh": fields.Boolean, | |||
| } | |||
| @setup_required | |||
| @@ -354,7 +359,11 @@ class EducationApi(Resource): | |||
| def get(self): | |||
| account = current_user | |||
| return BillingService.EducationIdentity.is_active(account.id) | |||
| res = BillingService.EducationIdentity.status(account.id) | |||
| # convert expire_at to UTC timestamp from isoformat | |||
| if res and "expire_at" in res: | |||
| res["expire_at"] = datetime.fromisoformat(res["expire_at"]).astimezone(pytz.utc) | |||
| return res | |||
| class EducationAutoCompleteApi(Resource): | |||
| @@ -1,27 +1,38 @@ | |||
| from flask_restful import ( | |||
| Resource, # type: ignore | |||
| reqparse, | |||
| ) | |||
| from flask_restful import Resource, reqparse | |||
| from controllers.console.wraps import setup_required | |||
| from controllers.inner_api import api | |||
| from controllers.inner_api.wraps import enterprise_inner_api_only | |||
| from services.enterprise.mail_service import DifyMail, EnterpriseMailService | |||
| from controllers.inner_api.wraps import billing_inner_api_only, enterprise_inner_api_only | |||
| from tasks.mail_inner_task import send_inner_email_task | |||
| _mail_parser = reqparse.RequestParser() | |||
| _mail_parser.add_argument("to", type=str, action="append", required=True) | |||
| _mail_parser.add_argument("subject", type=str, required=True) | |||
| _mail_parser.add_argument("body", type=str, required=True) | |||
| _mail_parser.add_argument("substitutions", type=dict, required=False) | |||
| class BaseMail(Resource): | |||
| """Shared logic for sending an inner email.""" | |||
| class EnterpriseMail(Resource): | |||
| @setup_required | |||
| @enterprise_inner_api_only | |||
| def post(self): | |||
| parser = reqparse.RequestParser() | |||
| parser.add_argument("to", type=str, action="append", required=True) | |||
| parser.add_argument("subject", type=str, required=True) | |||
| parser.add_argument("body", type=str, required=True) | |||
| parser.add_argument("substitutions", type=dict, required=False) | |||
| args = parser.parse_args() | |||
| EnterpriseMailService.send_mail(DifyMail(**args)) | |||
| args = _mail_parser.parse_args() | |||
| send_inner_email_task.delay( | |||
| to=args["to"], | |||
| subject=args["subject"], | |||
| body=args["body"], | |||
| substitutions=args["substitutions"], | |||
| ) | |||
| return {"message": "success"}, 200 | |||
| class EnterpriseMail(BaseMail): | |||
| method_decorators = [setup_required, enterprise_inner_api_only] | |||
| class BillingMail(BaseMail): | |||
| method_decorators = [setup_required, billing_inner_api_only] | |||
| api.add_resource(EnterpriseMail, "/enterprise/mail") | |||
| api.add_resource(BillingMail, "/billing/mail") | |||
| @@ -10,6 +10,22 @@ from extensions.ext_database import db | |||
| from models.model import EndUser | |||
| def billing_inner_api_only(view): | |||
| @wraps(view) | |||
| def decorated(*args, **kwargs): | |||
| if not dify_config.INNER_API: | |||
| abort(404) | |||
| # get header 'X-Inner-Api-Key' | |||
| inner_api_key = request.headers.get("X-Inner-Api-Key") | |||
| if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY: | |||
| abort(401) | |||
| return view(*args, **kwargs) | |||
| return decorated | |||
| def enterprise_inner_api_only(view): | |||
| @wraps(view) | |||
| def decorated(*args, **kwargs): | |||
| @@ -123,7 +123,7 @@ class BillingService: | |||
| return BillingService._send_request("GET", "/education/verify", params=params) | |||
| @classmethod | |||
| def is_active(cls, account_id: str): | |||
| def status(cls, account_id: str): | |||
| params = {"account_id": account_id} | |||
| return BillingService._send_request("GET", "/education/status", params=params) | |||
| @@ -1,18 +0,0 @@ | |||
| from pydantic import BaseModel | |||
| from tasks.mail_enterprise_task import send_enterprise_email_task | |||
| class DifyMail(BaseModel): | |||
| to: list[str] | |||
| subject: str | |||
| body: str | |||
| substitutions: dict[str, str] = {} | |||
| class EnterpriseMailService: | |||
| @classmethod | |||
| def send_mail(cls, mail: DifyMail): | |||
| send_enterprise_email_task.delay( | |||
| to=mail.to, subject=mail.subject, body=mail.body, substitutions=mail.substitutions | |||
| ) | |||
| @@ -11,7 +11,7 @@ from libs.email_i18n import get_email_i18n_service | |||
| @shared_task(queue="mail") | |||
| def send_enterprise_email_task(to: list[str], subject: str, body: str, substitutions: Mapping[str, str]): | |||
| def send_inner_email_task(to: list[str], subject: str, body: str, substitutions: Mapping[str, str]): | |||
| if not mail.is_inited(): | |||
| return | |||