define bas environment

2023-12-01 14:29:10 +08:00
parent 66e407add5
commit b7b968d5f7
18 changed files with 1700 additions and 1023 deletions
@@ -22,7 +22,7 @@ class LocalEmail:
                if latest_journal.is_finish():
                    yield None
                    continue
-                parsed = str(latest_journal.get_object_id())
+                parsed = latest_journal.get_input()
            
            mail_id = self.mail_storage.next_mail_id(parsed)
            if mail_id is None:
@@ -7,17 +7,20 @@ import datetime
 from bs4 import BeautifulSoup
 import sqlite3
 import html2text
+from urllib.parse import urlparse
 from aios import *

+
+
 class Mail:
    def __init__(self, **kwargs) -> None:
-        self.from_addr = kwargs.get("From")
-        self.to_addr = kwargs.get("To")
-        self.subject = kwargs.get("Subject")
-        self.date = kwargs.get("Date")
-        self.bcc = kwargs.get("BCC")
-        self.cc = kwargs.get("CC")
-        self.reply_to = None
+        self.from_addr = kwargs.get("from")
+        self.to_addr = kwargs.get("to")
+        self.subject = kwargs.get("subject")
+        self.date = kwargs.get("date")
+        self.bcc = kwargs.get("bcc")
+        self.cc = kwargs.get("cc")
+        self.reply_to = kwargs.get("reply_to")
        self.id: str = None
        self.content: str = None

@@ -192,20 +195,36 @@ class MailStorage:
                    self.conn.commit()
            await asyncio.sleep(10)
    
-    def download(self, uid, mail: mailparser.MailParser):
+    def download(self, uid, parser: mailparser.MailParser, 
+                 save_image=True, 
+                 from_field="From",
+                 to_field="To",
+                 subject_field="Subject",
+                 date_field="Date",
+                 reply_to_field="In-Reply-To",
+                 cc_field="CC",
+                 bcc_field="BCC"):
        mail_dir = self.mail_dir(uid)
-        os.makedirs(dir)
+        if not os.path.exists(mail_dir):
+            os.makedirs(mail_dir)

-        meta = json.loads(mail.mail_json)
-        mail = Mail(**meta)
-        reply_to = meta.get("In-Reply-To")
+        src_meta = json.loads(parser.mail_json)
+        meta = {}
+        meta["from"] = src_meta.get(from_field)
+        meta["to"] = src_meta.get(to_field)
+        meta["subject"] = src_meta.get(subject_field)
+        meta["date"] = src_meta.get(date_field)
+        meta["bcc"] = src_meta.get(bcc_field)
+        meta["cc"] = src_meta.get(cc_field)
+        reply_to = src_meta.get(reply_to_field)
        if reply_to:
-            mail.reply_to = self.uid_to_object_id(reply_to)
+            meta["reply_to"] = self.uid_to_object_id(reply_to)
+        mail = Mail(**meta)       
        
        h = html2text.HTML2Text()
        h.ignore_links = True
        h.ignore_images = True
-        mail_content = h.handle(mail.body)    
+        mail_content = h.handle(parser.body)    
        mail.content = mail_content
       
        mail.calculate_id()
@@ -216,41 +235,52 @@ class MailStorage:
        with open(f"{mail_dir}/mail.txt", "w", encoding='utf-8') as f:
            f.write(mail_content)
        
-        for attachment in mail.attachments:
-            if attachment['mail_content_type'] in ['image/png', 'image/jpeg', 'image/gif']:
-                filename = attachment['filename']
-                filefullname = f"{mail_dir}/{filename}"
-                image_data = attachment['payload']
+        if save_image:
+            for attachment in parser.attachments:
+                if attachment['mail_content_type'] in ['image/png', 'image/jpg', 'image/jpeg', 'image/gif', 'image/svg']:
+                    filename = attachment['filename']
+                    filefullname = f"{mail_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")
+
+            # get all image urls
+            soup = BeautifulSoup(parser.body, '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
+            
+            for img_url in img_urls:
+                # keep the original image filename(last of url)
+                url_result = urlparse(img_url)
+                if url_result.scheme not in ['http', 'https']:
+                    continue
+                ext = url_result.path.split('/')[-1].split('.')[-1]
+                if ext in ['png', 'jpg', 'jpeg', 'gif', 'svg']:
+                    img_filename = os.path.join(mail_dir, f"{name_count}.{ext}")
+                else :
+                    img_filename = os.path.join(mail_dir, f"{name_count}")
+                name_count += 1
+                # download image 
                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")
-
-        # get all image urls
-        soup = BeautifulSoup(mail.body, '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
-        
-        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(mail_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}')
+                    response = requests.get(img_url, stream=True)
+                except requests.exceptions.RequestException as e:
+                    logging.error(f'Failed to download {img_url}: {e}')
+                    continue
+                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.error(f'Failed to download {img_url}')

        cursor = self.conn.cursor()
        cursor.execute(
@@ -260,5 +290,8 @@ class MailStorage:
        """,
            (uid, mail.id, mail.date, mail.from_addr),
        )
+        self.conn.commit()
+
+        return mail.id

    
@@ -1,9 +1,13 @@
 import os
 import logging
 import json
+import string
 import imaplib
 import mailparser
-from aios import *
+
+from knowledge import *
+from aios_kernel.storage import AIStorage
+from .mail import Mail, MailStorage


 class EmailSpider:
@@ -16,14 +20,22 @@ class EmailSpider:
            port=self.config.get('imap_port')
        )
        self.client.login(self.config.get('address'), self.config.get('password'))
-        self.mail_local_root = os.path.join(self.env.pipeline_path, self.config.get("address"))
-        os.makedirs(self.mail_local_root)
+        self.client.select("INBOX")
+        local_path = string.Template(config["path"]).substitute(myai_dir=AIStorage.get_instance().get_myai_dir())
+        local_path = os.path.join(local_path, self.config.get('address'))
+        self.mail_storage = MailStorage(local_path)
+       

    async def next(self):
        while True:
-            _, data = self.client.uid('search', None, "ALL")
+            try:
+                _, data = self.client.uid('search', None, "ALL")
+            except Exception as e:
+                self.env.get_logger().error(f"email spider error: {e}")
+                yield (None, None)
+                continue
            uid_list = data[0].split()
-            if uid_list.len() == 0:
+            if len(uid_list) == 0:
                yield (None, None)
                continue
            
@@ -43,9 +55,16 @@ class EmailSpider:
                _uid = int.from_bytes(uid)
                if _uid > from_uid:
                    message_parts = "(BODY.PEEK[])"
-                    _, email_data = self.client.uid('fetch', uid, message_parts)
-                    mail = mailparser.parse_from_bytes(email_data[0][1])
-                    self.save_email(_uid, mail)
+                    try:
+                        _, email_data = self.client.uid('fetch', uid, message_parts)
+                        mail = mailparser.parse_from_bytes(email_data[0][1])
+                        id = self.mail_storage.download(_uid, mail)
+                    except Exception as e:
+                        self.env.get_logger().error(f"email spider error: {e}")
+                        yield (None, None)
+                        break
+                    yield (ObjectID.from_base58(id), str(_uid))
+                   

            yield (None, None)