263 lines
9.3 KiB
Python
263 lines
9.3 KiB
Python
import logging
|
||
import os
|
||
import email
|
||
from email.policy import default
|
||
from sqlalchemy.orm import Session
|
||
from datetime import datetime
|
||
import re
|
||
|
||
from ..models.domain import Domain
|
||
from ..models.mailbox import Mailbox
|
||
from ..models.email import Email
|
||
from ..models.attachment import Attachment
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class MailStore:
|
||
"""邮件存储服务，负责保存和检索邮件"""
|
||
|
||
def __init__(self, db_session_factory, storage_path=None):
|
||
"""
|
||
初始化邮件存储服务
|
||
|
||
参数:
|
||
db_session_factory: 数据库会话工厂函数
|
||
storage_path: 附件存储路径
|
||
"""
|
||
self.db_session_factory = db_session_factory
|
||
self.storage_path = storage_path or os.path.join(os.getcwd(), 'email_data')
|
||
|
||
# 确保存储目录存在
|
||
if not os.path.exists(self.storage_path):
|
||
os.makedirs(self.storage_path)
|
||
|
||
async def save_email(self, sender, recipient, message, raw_data):
|
||
"""
|
||
保存一封电子邮件
|
||
|
||
参数:
|
||
sender: 发件人地址
|
||
recipient: 收件人地址
|
||
message: 解析后的邮件对象
|
||
raw_data: 原始邮件数据
|
||
|
||
返回:
|
||
成功返回邮件ID，失败返回None
|
||
"""
|
||
# 从收件人地址中提取用户名和域名
|
||
try:
|
||
address, domain_name = recipient.split('@', 1)
|
||
except ValueError:
|
||
logger.warning(f"无效的收件人地址格式: {recipient}")
|
||
return None
|
||
|
||
# 获取数据库会话
|
||
db = self.db_session_factory()
|
||
|
||
try:
|
||
# 检查域名是否存在且活跃
|
||
domain = db.query(Domain).filter_by(name=domain_name, active=True).first()
|
||
if not domain:
|
||
logger.warning(f"不支持的域名: {domain_name}")
|
||
return None
|
||
|
||
# 查找或创建邮箱
|
||
mailbox = db.query(Mailbox).filter_by(address=address, domain_id=domain.id).first()
|
||
if not mailbox:
|
||
# 自动创建新邮箱
|
||
mailbox = Mailbox(
|
||
address=address,
|
||
domain_id=domain.id,
|
||
active=True
|
||
)
|
||
db.add(mailbox)
|
||
db.flush() # 获取ID但不提交
|
||
logger.info(f"已为 {recipient} 自动创建邮箱")
|
||
|
||
# 提取邮件内容
|
||
subject = message.get('subject', '')
|
||
|
||
# 获取文本和HTML内容
|
||
body_text = None
|
||
body_html = None
|
||
attachments_data = []
|
||
|
||
if message.is_multipart():
|
||
for part in message.walk():
|
||
content_type = part.get_content_type()
|
||
content_disposition = part.get_content_disposition()
|
||
|
||
# 处理文本内容
|
||
if content_disposition is None or content_disposition == 'inline':
|
||
if content_type == 'text/plain' and not body_text:
|
||
body_text = part.get_content()
|
||
elif content_type == 'text/html' and not body_html:
|
||
body_html = part.get_content()
|
||
|
||
# 处理附件
|
||
elif content_disposition == 'attachment':
|
||
filename = part.get_filename()
|
||
if filename:
|
||
content = part.get_payload(decode=True)
|
||
if content:
|
||
attachments_data.append({
|
||
'filename': filename,
|
||
'content_type': content_type,
|
||
'data': content,
|
||
'size': len(content)
|
||
})
|
||
else:
|
||
# 非多部分邮件
|
||
content_type = message.get_content_type()
|
||
if content_type == 'text/plain':
|
||
body_text = message.get_content()
|
||
elif content_type == 'text/html':
|
||
body_html = message.get_content()
|
||
|
||
# 创建邮件记录
|
||
email_obj = Email(
|
||
mailbox_id=mailbox.id,
|
||
sender=sender,
|
||
recipients=recipient,
|
||
subject=subject,
|
||
body_text=body_text,
|
||
body_html=body_html,
|
||
headers={k: v for k, v in message.items()}
|
||
)
|
||
|
||
# 保存邮件
|
||
db.add(email_obj)
|
||
db.flush() # 获取ID但不提交
|
||
|
||
# 提取验证信息
|
||
email_obj.extract_verification_data()
|
||
|
||
# 保存附件
|
||
for attachment_data in attachments_data:
|
||
attachment = Attachment(
|
||
email_id=email_obj.id,
|
||
filename=attachment_data['filename'],
|
||
content_type=attachment_data['content_type'],
|
||
size=attachment_data['size']
|
||
)
|
||
|
||
db.add(attachment)
|
||
db.flush()
|
||
|
||
# 决定存储位置
|
||
if attachment_data['size'] > 1024 * 1024: # 大于1MB的存储到文件系统
|
||
attachments_dir = os.path.join(self.storage_path, 'attachments')
|
||
attachment.save_to_filesystem(attachment_data['data'], attachments_dir)
|
||
else:
|
||
# 小附件直接存储在数据库
|
||
attachment.content = attachment_data['data']
|
||
|
||
# 提交所有更改
|
||
db.commit()
|
||
logger.info(f"邮件已成功保存: {sender} -> {recipient}, ID: {email_obj.id}")
|
||
return email_obj.id
|
||
|
||
except Exception as e:
|
||
db.rollback()
|
||
logger.error(f"保存邮件时出错: {str(e)}")
|
||
return None
|
||
finally:
|
||
db.close()
|
||
|
||
def get_emails_for_mailbox(self, mailbox_id, limit=50, offset=0, unread_only=False):
|
||
"""获取指定邮箱的邮件列表"""
|
||
db = self.db_session_factory()
|
||
try:
|
||
query = db.query(Email).filter(Email.mailbox_id == mailbox_id)
|
||
|
||
if unread_only:
|
||
query = query.filter(Email.read == False)
|
||
|
||
total = query.count()
|
||
|
||
emails = query.order_by(Email.received_at.desc()) \
|
||
.limit(limit) \
|
||
.offset(offset) \
|
||
.all()
|
||
|
||
return {
|
||
'total': total,
|
||
'items': [email.to_dict() for email in emails]
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"获取邮件列表时出错: {str(e)}")
|
||
return {'total': 0, 'items': []}
|
||
finally:
|
||
db.close()
|
||
|
||
def get_email_by_id(self, email_id, mark_as_read=True):
|
||
"""获取指定ID的邮件详情"""
|
||
db = self.db_session_factory()
|
||
try:
|
||
email = db.query(Email).filter(Email.id == email_id).first()
|
||
|
||
if not email:
|
||
return None
|
||
|
||
if mark_as_read and not email.read:
|
||
email.read = True
|
||
email.last_read = datetime.utcnow()
|
||
db.commit()
|
||
|
||
# 获取附件信息
|
||
attachments = [attachment.to_dict() for attachment in email.attachments]
|
||
|
||
# 构建完整响应
|
||
result = email.to_dict()
|
||
result['body_text'] = email.body_text
|
||
result['body_html'] = email.body_html
|
||
result['attachments'] = attachments
|
||
|
||
return result
|
||
except Exception as e:
|
||
db.rollback()
|
||
logger.error(f"获取邮件详情时出错: {str(e)}")
|
||
return None
|
||
finally:
|
||
db.close()
|
||
|
||
def delete_email(self, email_id):
|
||
"""删除指定ID的邮件"""
|
||
db = self.db_session_factory()
|
||
try:
|
||
email = db.query(Email).filter(Email.id == email_id).first()
|
||
|
||
if not email:
|
||
return False
|
||
|
||
db.delete(email)
|
||
db.commit()
|
||
return True
|
||
except Exception as e:
|
||
db.rollback()
|
||
logger.error(f"删除邮件时出错: {str(e)}")
|
||
return False
|
||
finally:
|
||
db.close()
|
||
|
||
def get_attachment_content(self, attachment_id):
|
||
"""获取附件内容"""
|
||
db = self.db_session_factory()
|
||
try:
|
||
attachment = db.query(Attachment).filter(Attachment.id == attachment_id).first()
|
||
|
||
if not attachment:
|
||
return None
|
||
|
||
content = attachment.get_content()
|
||
|
||
return {
|
||
'content': content,
|
||
'filename': attachment.filename,
|
||
'content_type': attachment.content_type
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"获取附件内容时出错: {str(e)}")
|
||
return None
|
||
finally:
|
||
db.close() |