Update parse_email.py

parse_email.py

@@ -1,27 +1,53 @@
 import email
-import re
+from email import policy
 from bs4 import BeautifulSoup
+import re
 
 def parse_email(file_path):
     with open(file_path, "rb") as f:
-        msg = email.message_from_binary_file(f)
+        msg = email.message_from_binary_file(f, policy=policy.default)
 
+    # --- 1. Extract headers ---
     headers = dict(msg.items())
-    
-    # Extract body (handle plain + HTML)
+
+    # --- 2. Extract body (text + html) ---
     body = ""
     if msg.is_multipart():
         for part in msg.walk():
-            if part.get_content_type() == "text/plain":
-                body += part.get_payload(decode=True).decode(errors="ignore")
-            elif part.get_content_type() == "text/html":
-                html = part.get_payload(decode=True).decode(errors="ignore")
-                soup = BeautifulSoup(html, "html.parser")
-                body += soup.get_text()
+            content_type = part.get_content_type()
+            if content_type == "text/plain":
+                try:
+                    body += part.get_content()
+                except:
+                    pass
+            elif content_type == "text/html":
+                try:
+                    html_body = part.get_content()
+                    soup = BeautifulSoup(html_body, "html.parser")
+                    body += soup.get_text(" ", strip=True)
+                except:
+                    pass
     else:
-        body = msg.get_payload(decode=True).decode(errors="ignore")
+        try:
+            body = msg.get_content()
+        except:
+            body = ""
+
+    # --- 3. Extract URLs ---
+    urls = set()
+
+    # From plain text
+    urls.update(re.findall(r"https?://[^\s]+", body))
 
-    # Extract URLs
-    urls = re.findall(r'(https?://\S+)', body)
+    # From HTML parts
+    for part in msg.walk():
+        if part.get_content_type() == "text/html":
+            try:
+                html_body = part.get_content()
+                soup = BeautifulSoup(html_body, "html.parser")
+                for link in soup.find_all("a", href=True):
+                    urls.add(link["href"])
+            except:
+                pass
 
-    return headers, body, urls
+    return headers, body, list(urls)