80 lines
3.3 KiB
Python
80 lines
3.3 KiB
Python
|
from django.core.management.base import BaseCommand, CommandError
|
|||
|
|
from user_management.models import GmailCredential, User
|
|||
|
|
from user_management.gmail_integration import GmailIntegration
|
|||
|
|
import logging
|
|||
|
|
import pickle
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
class Command(BaseCommand):
|
|||
|
|
help = '更新Gmail凭证中的邮箱信息'
|
|||
|
|
|
|||
|
|
def add_arguments(self, parser):
|
|||
|
|
parser.add_argument('--email', type=str, help='指定用户邮箱')
|
|||
|
|
parser.add_argument('--gmail', type=str, help='指定要设置的Gmail邮箱')
|
|||
|
|
parser.add_argument('--all', action='store_true', help='更新所有凭证')
|
|||
|
|
|
|||
|
|
def handle(self, *args, **options):
|
|||
|
|
email = options.get('email')
|
|||
|
|
gmail = options.get('gmail')
|
|||
|
|
update_all = options.get('all')
|
|||
|
|
|
|||
|
|
if update_all:
|
|||
|
|
# 更新所有凭证
|
|||
|
|
credentials = GmailCredential.objects.filter(is_active=True)
|
|||
|
|
self.stdout.write(f"找到 {credentials.count()} 个活跃的Gmail凭证")
|
|||
|
|
|
|||
|
|
for credential in credentials:
|
|||
|
|
self._update_credential(credential, gmail)
|
|||
|
|
elif email:
|
|||
|
|
# 更新指定用户的凭证
|
|||
|
|
try:
|
|||
|
|
user = User.objects.get(email=email)
|
|||
|
|
credentials = GmailCredential.objects.filter(user=user, is_active=True)
|
|||
|
|
|
|||
|
|
if not credentials.exists():
|
|||
|
|
raise CommandError(f"未找到用户 {email} 的Gmail凭证")
|
|||
|
|
|
|||
|
|
for credential in credentials:
|
|||
|
|
self._update_credential(credential, gmail)
|
|||
|
|
except User.DoesNotExist:
|
|||
|
|
raise CommandError(f"未找到用户 {email}")
|
|||
|
|
else:
|
|||
|
|
self.stdout.write("请提供--email参数或--all参数")
|
|||
|
|
|
|||
|
|
def _update_credential(self, credential, gmail=None):
|
|||
|
|
"""更新单个凭证"""
|
|||
|
|
user = credential.user
|
|||
|
|
self.stdout.write(f"正在更新用户 {user.email} 的Gmail凭证...")
|
|||
|
|
|
|||
|
|
if gmail:
|
|||
|
|
# 如果指定了Gmail邮箱，直接使用
|
|||
|
|
credential.gmail_email = gmail
|
|||
|
|
credential.save()
|
|||
|
|
self.stdout.write(self.style.SUCCESS(f"已手动设置Gmail邮箱为: {gmail}"))
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 尝试使用API获取Gmail邮箱
|
|||
|
|
try:
|
|||
|
|
# 从凭证数据中恢复服务
|
|||
|
|
creds = pickle.loads(credential.credentials)
|
|||
|
|
|
|||
|
|
if creds and not creds.invalid:
|
|||
|
|
# 创建Gmail集成实例
|
|||
|
|
integration = GmailIntegration(user=user)
|
|||
|
|
integration.credentials = creds
|
|||
|
|
|
|||
|
|
# 尝试调用API获取用户资料
|
|||
|
|
profile = integration.gmail_service.users().getProfile(userId='me').execute()
|
|||
|
|
gmail_email = profile.get('emailAddress')
|
|||
|
|
|
|||
|
|
if gmail_email:
|
|||
|
|
credential.gmail_email = gmail_email
|
|||
|
|
credential.save()
|
|||
|
|
self.stdout.write(self.style.SUCCESS(f"已更新Gmail邮箱为: {gmail_email}"))
|
|||
|
|
else:
|
|||
|
|
self.stdout.write(self.style.WARNING("无法从API获取Gmail邮箱"))
|
|||
|
|
else:
|
|||
|
|
self.stdout.write(self.style.ERROR("凭证无效，请重新授权"))
|
|||
|
|
except Exception as e:
|
|||
|
|
self.stdout.write(self.style.ERROR(f"更新失败: {str(e)}"))
|