email parser document
This commit is contained in:
tsukasa
@@ -0,0 +1,156 @@
|
||||
import os
|
||||
import aiofiles
|
||||
import chardet
|
||||
import logging
|
||||
import string
|
||||
from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal
|
||||
from aios_kernel.storage import AIStorage
|
||||
|
||||
class KnowledgeEmailSource:
|
||||
def __init__(self, env: KnowledgePipelineEnvironment, config:dict):
|
||||
self.config = config
|
||||
|
||||
# @classmethod
|
||||
# def user_config_items(cls):
|
||||
# return [("address", "email address"),
|
||||
# ("password", "email password"),
|
||||
# ("imap_server", "imap server"),
|
||||
# ("imap_port", "imap port")
|
||||
# ]
|
||||
|
||||
async def run_once(self):
|
||||
# read config from toml file
|
||||
# and read from config config.local.toml if exists (config.local.toml is ignored by git)
|
||||
logging.debug(f"knowledge email source {self.id()} run once")
|
||||
filter = "ALL"
|
||||
self.client = self.email_client()
|
||||
await self.read_emails(imap_keyword=filter)
|
||||
|
||||
def email_client(self) -> imaplib.IMAP4_SSL:
|
||||
logging.info(f"read email config from {self.config.get('imap_server')}")
|
||||
client = imaplib.IMAP4_SSL(
|
||||
host=self.config.get('imap_server'),
|
||||
port=self.config.get('imap_port')
|
||||
)
|
||||
client.login(self.config.get('address'), self.config.get('password'))
|
||||
return client
|
||||
|
||||
async def read_emails(self, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"):
|
||||
journal_client = KnowledgeJournalClient()
|
||||
latest_journal = journal_client.latest_journal(self.id())
|
||||
latest_uid = 0 if latest_journal is None else int(latest_journal.item_id)
|
||||
self.client.select(folder)
|
||||
_, data = self.client.uid('search', None, imap_keyword)
|
||||
|
||||
# get email uid list
|
||||
email_list = data[0].split()
|
||||
logging.info(f"got {len(email_list)} emails")
|
||||
journal_client = KnowledgeJournalClient()
|
||||
for uid in email_list:
|
||||
_uid = int.from_bytes(uid)
|
||||
if _uid > latest_uid:
|
||||
email_dir = self.check_email_saved(uid)
|
||||
if email_dir is not None:
|
||||
logging.info(f"email uid {uid} already saved")
|
||||
else:
|
||||
email_dir = self.read_and_save_email(uid)
|
||||
logging.info(f"email uid {uid} saved")
|
||||
email_object = EmailObjectBuilder({}, email_dir).build()
|
||||
await KnowledgeBase().insert_object(email_object)
|
||||
journal_client.insert(KnowledgeJournal("email", self.id(), str(int.from_bytes(uid)), str(email_object.calculate_id())))
|
||||
|
||||
|
||||
def read_and_save_email(self, uid: str) -> str:
|
||||
message_parts = "(BODY.PEEK[])"
|
||||
_, email_data = self.client.uid('fetch', uid, message_parts)
|
||||
mail = mailparser.parse_from_bytes(email_data[0][1])
|
||||
logging.info(f"got email subject [{mail.subject}]")
|
||||
self.save_email(mail)
|
||||
return self.get_local_dir_name(mail)
|
||||
|
||||
def get_local_dir_name(self, mail: mailparser.MailParser) -> str:
|
||||
dir = f"{self.local_root()}/{self.config.get('address')}"
|
||||
name = f"{mail.subject}__{mail.date}"
|
||||
name = hashlib.md5(name.encode('utf-8')).hexdigest()
|
||||
return f"{dir}/{name}"
|
||||
|
||||
def check_email_saved(self, uid: str) -> str:
|
||||
message_parts = "(BODY[HEADER])"
|
||||
_, email_data = self.client.uid('fetch', uid, message_parts)
|
||||
mail = mailparser.parse_from_bytes(email_data[0][1])
|
||||
logging.info(f"[{uid}]check email subject [{mail.subject}]")
|
||||
dir = self.get_local_dir_name(mail)
|
||||
logging.info(f"check email saved {dir}")
|
||||
file = f"{dir}/email.txt"
|
||||
if os.path.exists(file):
|
||||
return dir
|
||||
return None
|
||||
|
||||
# save email attachment(images)
|
||||
def save_email_attachment(self, mail: mailparser.MailParser, email_dir: str):
|
||||
for attachment in mail.attachments:
|
||||
if attachment['mail_content_type'] in ['image/png', 'image/jpeg', 'image/gif']:
|
||||
print('current mail have image attachment')
|
||||
img_dir = f"{email_dir}/image"
|
||||
if not os.path.exists(img_dir):
|
||||
os.makedirs(img_dir)
|
||||
filename = attachment['filename']
|
||||
filefullname = f"{img_dir}/{filename}"
|
||||
image_data = attachment['payload']
|
||||
try:
|
||||
image_data = base64.b64decode(image_data)
|
||||
except base64.binascii.Error:
|
||||
image_data = image_data.encode()
|
||||
with open(filefullname, 'wb') as f:
|
||||
f.write(image_data)
|
||||
logging.info(f"save email image {filename} success")
|
||||
|
||||
# save email body images(html content)
|
||||
def save_body_images(self, html_content: str, email_dir: str):
|
||||
# get all image urls
|
||||
soup = BeautifulSoup(html_content, 'html.parser')
|
||||
img_tags = soup.find_all('img')
|
||||
img_urls = [img['src'] for img in img_tags if 'src' in img.attrs]
|
||||
logging.info(f'Found {len(img_urls)} images in email body')
|
||||
|
||||
name_count = 0
|
||||
|
||||
if not os.path.exists(email_dir):
|
||||
os.makedirs(email_dir)
|
||||
|
||||
for img_url in img_urls:
|
||||
# keep the original image filename(last of url)
|
||||
ext = img_url.split('/')[-1].split('.')[-1]
|
||||
img_filename = os.path.join(email_dir, f"{name_count}.{ext}")
|
||||
name_count += 1
|
||||
# download image
|
||||
response = requests.get(img_url, stream=True)
|
||||
if response.status_code == 200:
|
||||
with open(img_filename, 'wb') as img_file:
|
||||
for chunk in response.iter_content(1024):
|
||||
img_file.write(chunk)
|
||||
logging.info(f'Downloaded {img_url} to {img_filename}')
|
||||
else:
|
||||
logging.info(f'Failed to download {img_url}')
|
||||
|
||||
# save email content to local dir
|
||||
def save_email(self, mail: mailparser.MailParser):
|
||||
dir = f"{self.local_root()}/{self.config.get('address')}"
|
||||
if not os.path.exists(dir):
|
||||
os.makedirs(dir)
|
||||
email_dir = self.get_local_dir_name(mail)
|
||||
logging.info(f"save email to {email_dir}")
|
||||
if not os.path.exists(email_dir):
|
||||
os.makedirs(email_dir)
|
||||
with open(f"{email_dir}/email.txt", "w", encoding='utf-8') as f:
|
||||
# soup = BeautifulSoup(mail.body, 'html.parser')
|
||||
f.write(mail.body)
|
||||
with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f:
|
||||
mail_dict = json.loads(mail.mail_json)
|
||||
if 'body' in mail_dict:
|
||||
del mail_dict['body']
|
||||
json.dump(mail_dict, f, ensure_ascii=False, indent=4)
|
||||
logging.info(f"save email meta info {f.name}")
|
||||
|
||||
self.save_email_attachment(mail, email_dir)
|
||||
self.save_body_images(mail.body, f"{email_dir}/body_image")
|
||||
Reference in New Issue
Block a user