From dab8c386e8c7158c06eb71f5203f9dc5d299a975 Mon Sep 17 00:00:00 2001
From: huangzhenpc <huangzhenpc@git586vip.cn>
Date: Thu, 27 Feb 2025 10:08:58 +0800
Subject: [PATCH] testapiaddyuming

---
 app/__init__.py | 23 ++++++++++++++++++++++-
 app/utils.py    | 30 +++++++++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/app/__init__.py b/app/__init__.py
index 79b1dcb..763a010 100644
--- a/app/__init__.py
+++ b/app/__init__.py
@@ -1,6 +1,6 @@
 from flask import Flask, request, jsonify
 from .config import Config
-from .utils import get_latest_emails, get_latest_email_with_code
+from .utils import get_latest_emails, get_latest_email_with_code, add_allowed_domain, remove_allowed_domain, get_allowed_domains
 
 def create_app():
     app = Flask(__name__)
@@ -25,4 +25,25 @@ def create_app():
             return jsonify(email_data)
         return jsonify({'error': 'No emails found for this recipient'}), 404
 
+    @app.route('/allowed_domains/add', methods=['POST'])
+    def add_domain():
+        domain = request.json.get('domain')
+        if not domain:
+            return jsonify({'error': 'Domain is required'}), 400
+        add_allowed_domain(domain)
+        return jsonify({'message': 'Domain added successfully'}), 201
+
+    @app.route('/allowed_domains/remove', methods=['POST'])
+    def remove_domain():
+        domain = request.json.get('domain')
+        if not domain:
+            return jsonify({'error': 'Domain is required'}), 400
+        remove_allowed_domain(domain)
+        return jsonify({'message': 'Domain removed successfully'}), 200
+
+    @app.route('/allowed_domains/list', methods=['GET'])
+    def list_domains():
+        domains = get_allowed_domains()
+        return jsonify(domains), 200
+
     return app 
\ No newline at end of file
diff --git a/app/utils.py b/app/utils.py
index 2fee2a7..7e7196b 100644
--- a/app/utils.py
+++ b/app/utils.py
@@ -258,4 +258,32 @@ def extract_code_from_body(body):
     """从邮件正文中提取验证码"""
     import re
     match = re.search(r'\b(\d{6})\b', body)
-    return match.group(1) if match else None 
\ No newline at end of file
+    return match.group(1) if match else None
+
+
+def add_allowed_domain(domain):
+    """添加允许的域名"""
+    try:
+        redis_client.sadd('allowed_domains', domain)
+        logger.info(f'Added allowed domain: {domain}')
+    except Exception as e:
+        logger.error(f'Error adding allowed domain: {e}')
+
+
+def remove_allowed_domain(domain):
+    """删除允许的域名"""
+    try:
+        redis_client.srem('allowed_domains', domain)
+        logger.info(f'Removed allowed domain: {domain}')
+    except Exception as e:
+        logger.error(f'Error removing allowed domain: {e}')
+
+
+def get_allowed_domains():
+    """获取当前允许的域名列表"""
+    try:
+        domains = redis_client.smembers('allowed_domains')
+        return [domain.decode() for domain in domains]
+    except Exception as e:
+        logger.error(f'Error fetching allowed domains: {e}')
+        return [] 
\ No newline at end of file