Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>tags/1.8.0
| from datetime import datetime | |||||
| import pytz | import pytz | ||||
| from flask import request | from flask import request | ||||
| from flask_login import current_user | from flask_login import current_user | ||||
| class EducationApi(Resource): | class EducationApi(Resource): | ||||
| status_fields = { | status_fields = { | ||||
| "result": fields.Boolean, | "result": fields.Boolean, | ||||
| "is_student": fields.Boolean, | |||||
| "expire_at": TimestampField, | |||||
| "allow_refresh": fields.Boolean, | |||||
| } | } | ||||
| @setup_required | @setup_required | ||||
| def get(self): | def get(self): | ||||
| account = current_user | 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): | class EducationAutoCompleteApi(Resource): |
| from flask_restful import ( | |||||
| Resource, # type: ignore | |||||
| reqparse, | |||||
| ) | |||||
| from flask_restful import Resource, reqparse | |||||
| from controllers.console.wraps import setup_required | from controllers.console.wraps import setup_required | ||||
| from controllers.inner_api import api | 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): | 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 | 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(EnterpriseMail, "/enterprise/mail") | ||||
| api.add_resource(BillingMail, "/billing/mail") |
| from models.model import EndUser | 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): | def enterprise_inner_api_only(view): | ||||
| @wraps(view) | @wraps(view) | ||||
| def decorated(*args, **kwargs): | def decorated(*args, **kwargs): |
| return BillingService._send_request("GET", "/education/verify", params=params) | return BillingService._send_request("GET", "/education/verify", params=params) | ||||
| @classmethod | @classmethod | ||||
| def is_active(cls, account_id: str): | |||||
| def status(cls, account_id: str): | |||||
| params = {"account_id": account_id} | params = {"account_id": account_id} | ||||
| return BillingService._send_request("GET", "/education/status", params=params) | return BillingService._send_request("GET", "/education/status", params=params) | ||||
| 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 | |||||
| ) |
| @shared_task(queue="mail") | @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(): | if not mail.is_inited(): | ||||
| return | return | ||||