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)}")) |