main
py 587 lines 21.8 KB
Raw
1 # import asyncio
2 # import email
3 # import os
4 # import re
5 # import uuid
6 # from dataclasses import dataclass
7 # from email.header import decode_header
8 # from email.message import Message as EmailMessage
9 # from fnmatch import fnmatch
10 # from typing import Any, Dict, List, Optional, Tuple
11
12 # import html2text
13 # from bs4 import BeautifulSoup
14 # from imapclient import IMAPClient
15
16 # from helpers import files
17 # from helpers.errors import RepairableException, format_error
18 # from helpers.print_style import PrintStyle
19
20
21 # @dataclass
22 # class Message:
23 # """Email message representation with sender, subject, body, and attachments."""
24 # sender: str
25 # subject: str
26 # body: str
27 # attachments: List[str]
28
29
30 # class EmailClient:
31 # """
32 # Async email client for reading messages from IMAP and Exchange servers.
33
34 # """
35
36 # def __init__(
37 # self,
38 # account_type: str = "imap",
39 # server: str = "",
40 # port: int = 993,
41 # username: str = "",
42 # password: str = "",
43 # options: Optional[Dict[str, Any]] = None,
44 # ):
45 # """
46 # Initialize email client with connection parameters.
47
48 # Args:
49 # account_type: Type of account - "imap" or "exchange"
50 # server: Mail server address (e.g., "imap.gmail.com")
51 # port: Server port (default 993 for IMAP SSL)
52 # username: Email account username
53 # password: Email account password
54 # options: Optional configuration dict with keys:
55 # - ssl: Use SSL/TLS (default: True)
56 # - timeout: Connection timeout in seconds (default: 30)
57 # """
58 # self.account_type = account_type.lower()
59 # self.server = server
60 # self.port = port
61 # self.username = username
62 # self.password = password
63 # self.options = options or {}
64
65 # # Default options
66 # self.ssl = self.options.get("ssl", True)
67 # self.timeout = self.options.get("timeout", 30)
68
69 # self.client: Optional[IMAPClient] = None
70 # self.exchange_account = None
71
72 # async def connect(self) -> None:
73 # """Establish connection to email server."""
74 # try:
75 # if self.account_type == "imap":
76 # await self._connect_imap()
77 # elif self.account_type == "exchange":
78 # await self._connect_exchange()
79 # else:
80 # raise RepairableException(
81 # f"Unsupported account type: {self.account_type}. "
82 # "Supported types: 'imap', 'exchange'"
83 # )
84 # except Exception as e:
85 # err = format_error(e)
86 # PrintStyle.error(f"Failed to connect to email server: {err}")
87 # raise RepairableException(f"Email connection failed: {err}") from e
88
89 # async def _connect_imap(self) -> None:
90 # """Establish IMAP connection."""
91 # loop = asyncio.get_event_loop()
92
93 # def _sync_connect():
94 # client = IMAPClient(self.server, port=self.port, ssl=self.ssl, timeout=self.timeout)
95 # # Increase line length limit to handle large emails (default is 10000)
96 # # This fixes "line too long" errors for emails with large headers or embedded content
97 # client._imap._maxline = 100000
98 # client.login(self.username, self.password)
99 # return client
100
101 # self.client = await loop.run_in_executor(None, _sync_connect)
102 # PrintStyle.standard(f"Connected to IMAP server: {self.server}")
103
104 # async def _connect_exchange(self) -> None:
105 # """Establish Exchange connection."""
106 # try:
107 # from exchangelib import Account, Configuration, Credentials, DELEGATE
108
109 # loop = asyncio.get_event_loop()
110
111 # def _sync_connect():
112 # creds = Credentials(username=self.username, password=self.password)
113 # config = Configuration(server=self.server, credentials=creds)
114 # return Account(
115 # primary_smtp_address=self.username,
116 # config=config,
117 # autodiscover=False,
118 # access_type=DELEGATE
119 # )
120
121 # self.exchange_account = await loop.run_in_executor(None, _sync_connect)
122 # PrintStyle.standard(f"Connected to Exchange server: {self.server}")
123 # except ImportError as e:
124 # raise RepairableException(
125 # "exchangelib not installed. Install with: pip install exchangelib>=5.4.3"
126 # ) from e
127
128 # async def disconnect(self) -> None:
129 # """Clean up connection."""
130 # try:
131 # if self.client:
132 # loop = asyncio.get_event_loop()
133 # await loop.run_in_executor(None, self.client.logout)
134 # self.client = None
135 # PrintStyle.standard("Disconnected from IMAP server")
136 # elif self.exchange_account:
137 # self.exchange_account = None
138 # PrintStyle.standard("Disconnected from Exchange server")
139 # except Exception as e:
140 # PrintStyle.error(f"Error during disconnect: {format_error(e)}")
141
142 # async def read_messages(
143 # self,
144 # download_folder: str,
145 # filter: Optional[Dict[str, Any]] = None,
146 # ) -> List[Message]:
147 # """
148 # Read messages based on filter criteria.
149
150 # Args:
151 # download_folder: Folder to save attachments (relative to /a0/)
152 # filter: Filter criteria dict with keys:
153 # - unread: Boolean to filter unread messages (default: True)
154 # - sender: Sender pattern with wildcards (e.g., "*@company.com")
155 # - subject: Subject pattern with wildcards (e.g., "*invoice*")
156 # - since_date: Optional datetime for date filtering
157
158 # Returns:
159 # List of Message objects with attachments saved to download_folder
160 # """
161 # filter = filter or {}
162
163 # if self.account_type == "imap":
164 # return await self._fetch_imap_messages(download_folder, filter)
165 # elif self.account_type == "exchange":
166 # return await self._fetch_exchange_messages(download_folder, filter)
167 # else:
168 # raise RepairableException(f"Unsupported account type: {self.account_type}")
169
170 # async def _fetch_imap_messages(
171 # self,
172 # download_folder: str,
173 # filter: Dict[str, Any],
174 # ) -> List[Message]:
175 # """Fetch messages from IMAP server."""
176 # if not self.client:
177 # raise RepairableException("IMAP client not connected. Call connect() first.")
178
179 # loop = asyncio.get_event_loop()
180 # messages: List[Message] = []
181
182 # def _sync_fetch():
183 # # Select inbox
184 # self.client.select_folder("INBOX")
185
186 # # Build search criteria
187 # search_criteria = []
188 # if filter.get("unread", True):
189 # search_criteria.append("UNSEEN")
190
191 # if filter.get("since_date"):
192 # since_date = filter["since_date"]
193 # search_criteria.append(["SINCE", since_date])
194
195 # # Search for messages
196 # if not search_criteria:
197 # search_criteria = ["ALL"]
198
199 # message_ids = self.client.search(search_criteria)
200 # return message_ids
201
202 # message_ids = await loop.run_in_executor(None, _sync_fetch)
203
204 # if not message_ids:
205 # PrintStyle.hint("No messages found matching criteria")
206 # return messages
207
208 # PrintStyle.standard(f"Found {len(message_ids)} messages")
209
210 # # Fetch and process messages
211 # for msg_id in message_ids:
212 # try:
213 # msg = await self._fetch_and_parse_imap_message(msg_id, download_folder, filter)
214 # if msg:
215 # messages.append(msg)
216 # except Exception as e:
217 # PrintStyle.error(f"Error processing message {msg_id}: {format_error(e)}")
218 # continue
219
220 # return messages
221
222 # async def _fetch_and_parse_imap_message(
223 # self,
224 # msg_id: int,
225 # download_folder: str,
226 # filter: Dict[str, Any],
227 # ) -> Optional[Message]:
228 # """Fetch and parse a single IMAP message with retry logic for large messages."""
229 # loop = asyncio.get_event_loop()
230
231 # def _sync_fetch():
232 # try:
233 # # Try standard RFC822 fetch first
234 # return self.client.fetch([msg_id], ["RFC822"])[msg_id]
235 # except Exception as e:
236 # error_msg = str(e).lower()
237 # # If "line too long" error, try fetching in parts
238 # if "line too long" in error_msg or "fetch_failed" in error_msg:
239 # PrintStyle.warning(f"Message {msg_id} too large for standard fetch, trying alternative method")
240 # # Fetch headers and body separately to avoid line length issues
241 # try:
242 # envelope = self.client.fetch([msg_id], ["BODY.PEEK[]"])[msg_id]
243 # return envelope
244 # except Exception as e2:
245 # PrintStyle.error(f"Alternative fetch also failed for message {msg_id}: {format_error(e2)}")
246 # raise
247 # raise
248
249 # try:
250 # raw_msg = await loop.run_in_executor(None, _sync_fetch)
251
252 # # Extract email data from response
253 # if b"RFC822" in raw_msg:
254 # email_data = raw_msg[b"RFC822"]
255 # elif b"BODY[]" in raw_msg:
256 # email_data = raw_msg[b"BODY[]"]
257 # else:
258 # PrintStyle.error(f"Unexpected response format for message {msg_id}")
259 # return None
260
261 # email_msg = email.message_from_bytes(email_data)
262
263 # # Apply sender filter
264 # sender = self._decode_header(email_msg.get("From", ""))
265 # if filter.get("sender") and not fnmatch(sender, filter["sender"]):
266 # return None
267
268 # # Apply subject filter
269 # subject = self._decode_header(email_msg.get("Subject", ""))
270 # if filter.get("subject") and not fnmatch(subject, filter["subject"]):
271 # return None
272
273 # # Parse message
274 # return await self._parse_message(email_msg, download_folder)
275
276 # except Exception as e:
277 # PrintStyle.error(f"Failed to fetch/parse message {msg_id}: {format_error(e)}")
278 # return None
279
280 # async def _fetch_exchange_messages(
281 # self,
282 # download_folder: str,
283 # filter: Dict[str, Any],
284 # ) -> List[Message]:
285 # """Fetch messages from Exchange server."""
286 # if not self.exchange_account:
287 # raise RepairableException("Exchange account not connected. Call connect() first.")
288
289 # from exchangelib import Q
290
291 # loop = asyncio.get_event_loop()
292 # messages: List[Message] = []
293
294 # def _sync_fetch():
295 # # Build query
296 # query = None
297 # if filter.get("unread", True):
298 # query = Q(is_read=False)
299
300 # if filter.get("sender"):
301 # sender_pattern = filter["sender"].replace("*", "")
302 # sender_q = Q(sender__contains=sender_pattern)
303 # query = query & sender_q if query else sender_q
304
305 # if filter.get("subject"):
306 # subject_pattern = filter["subject"].replace("*", "")
307 # subject_q = Q(subject__contains=subject_pattern)
308 # query = query & subject_q if query else subject_q
309
310 # # Fetch messages from inbox
311 # inbox = self.exchange_account.inbox
312 # items = inbox.filter(query) if query else inbox.all()
313 # return list(items)
314
315 # exchange_messages = await loop.run_in_executor(None, _sync_fetch)
316
317 # PrintStyle.standard(f"Found {len(exchange_messages)} Exchange messages")
318
319 # # Process messages
320 # for ex_msg in exchange_messages:
321 # try:
322 # msg = await self._parse_exchange_message(ex_msg, download_folder)
323 # if msg:
324 # messages.append(msg)
325 # except Exception as e:
326 # PrintStyle.error(f"Error processing Exchange message: {format_error(e)}")
327 # continue
328
329 # return messages
330
331 # async def _parse_exchange_message(
332 # self,
333 # ex_msg,
334 # download_folder: str,
335 # ) -> Message:
336 # """Parse an Exchange message."""
337 # loop = asyncio.get_event_loop()
338
339 # def _get_body():
340 # return str(ex_msg.text_body or ex_msg.body or "")
341
342 # body = await loop.run_in_executor(None, _get_body)
343
344 # # Process HTML if present
345 # if ex_msg.body and str(ex_msg.body).strip().startswith("<"):
346 # body = self._html_to_text(str(ex_msg.body))
347
348 # # Save attachments
349 # attachment_paths = []
350 # if ex_msg.attachments:
351 # for attachment in ex_msg.attachments:
352 # if hasattr(attachment, "content"):
353 # path = await self._save_attachment_bytes(
354 # attachment.name,
355 # attachment.content,
356 # download_folder
357 # )
358 # attachment_paths.append(path)
359
360 # return Message(
361 # sender=str(ex_msg.sender.email_address) if ex_msg.sender else "",
362 # subject=str(ex_msg.subject or ""),
363 # body=body,
364 # attachments=attachment_paths
365 # )
366
367 # async def _parse_message(
368 # self,
369 # email_msg: EmailMessage,
370 # download_folder: str,
371 # ) -> Message:
372 # """
373 # Parse email message and extract content with inline attachments.
374
375 # Processes multipart messages, converts HTML to text, and maintains
376 # positional context for inline attachments.
377 # """
378 # sender = self._decode_header(email_msg.get("From", ""))
379 # subject = self._decode_header(email_msg.get("Subject", ""))
380
381 # # Extract body and attachments
382 # body = ""
383 # attachment_paths: List[str] = []
384 # cid_map: Dict[str, str] = {} # Map Content-ID to file paths
385 # body_parts: List[str] = [] # Track parts in order
386
387 # if email_msg.is_multipart():
388 # # Process parts in order to maintain attachment positions
389 # for part in email_msg.walk():
390 # content_type = part.get_content_type()
391 # content_disposition = str(part.get("Content-Disposition", ""))
392
393 # # Skip multipart containers
394 # if part.get_content_maintype() == "multipart":
395 # continue
396
397 # # Handle attachments
398 # if "attachment" in content_disposition or part.get("Content-ID"):
399 # filename = part.get_filename()
400 # if filename:
401 # filename = self._decode_header(filename)
402 # content = part.get_payload(decode=True)
403 # if content:
404 # path = await self._save_attachment_bytes(
405 # filename, content, download_folder
406 # )
407 # attachment_paths.append(path)
408
409 # # Map Content-ID for inline images
410 # cid = part.get("Content-ID")
411 # if cid:
412 # cid = cid.strip("<>")
413 # cid_map[cid] = path
414
415 # # Add positional marker for non-cid attachments
416 # # (cid attachments are positioned via HTML references)
417 # if not cid and body_parts:
418 # body_parts.append(f"\n[file://{path}]\n")
419
420 # # Handle body text
421 # elif content_type == "text/plain":
422 # if not body: # Use first text/plain as primary body
423 # charset = part.get_content_charset() or "utf-8"
424 # body = part.get_payload(decode=True).decode(charset, errors="ignore")
425 # body_parts.append(body)
426
427 # elif content_type == "text/html":
428 # if not body: # Use first text/html as primary body if no text/plain
429 # charset = part.get_content_charset() or "utf-8"
430 # html_content = part.get_payload(decode=True).decode(charset, errors="ignore")
431 # body = self._html_to_text(html_content, cid_map)
432 # body_parts.append(body)
433
434 # # Combine body parts if we built them up
435 # if len(body_parts) > 1:
436 # body = "".join(body_parts)
437 # else:
438 # # Single part message
439 # content_type = email_msg.get_content_type()
440 # charset = email_msg.get_content_charset() or "utf-8"
441 # content = email_msg.get_payload(decode=True)
442 # if content:
443 # if content_type == "text/html":
444 # body = self._html_to_text(content.decode(charset, errors="ignore"), cid_map)
445 # else:
446 # body = content.decode(charset, errors="ignore")
447
448 # return Message(
449 # sender=sender,
450 # subject=subject,
451 # body=body,
452 # attachments=attachment_paths
453 # )
454
455 # def _html_to_text(self, html_content: str, cid_map: Optional[Dict[str, str]] = None) -> str:
456 # """
457 # Convert HTML to plain text with inline attachment references.
458
459 # Replaces inline images with [file:///a0/...] markers to maintain
460 # positional context.
461 # """
462 # cid_map = cid_map or {}
463
464 # # Replace cid: references with file paths before conversion
465 # if cid_map:
466 # soup = BeautifulSoup(html_content, "html.parser")
467 # for img in soup.find_all("img"):
468 # src = img.get("src", "")
469 # if src.startswith("cid:"):
470 # cid = src[4:] # Remove "cid:" prefix
471 # if cid in cid_map:
472 # # Replace with file path marker
473 # file_marker = f"[file://{cid_map[cid]}]"
474 # img.replace_with(soup.new_string(file_marker))
475 # html_content = str(soup)
476
477 # # Convert HTML to text
478 # h = html2text.HTML2Text()
479 # h.ignore_links = False
480 # h.ignore_images = False
481 # h.ignore_emphasis = False
482 # h.body_width = 0 # Don't wrap lines
483
484 # text = h.handle(html_content)
485
486 # # Clean up extra whitespace
487 # text = re.sub(r"\n{3,}", "\n\n", text) # Max 2 consecutive newlines
488 # text = text.strip()
489
490 # return text
491
492 # async def _save_attachment_bytes(
493 # self,
494 # filename: str,
495 # content: bytes,
496 # download_folder: str,
497 # ) -> str:
498 # """
499 # Save attachment to disk and return absolute path.
500
501 # Uses Agent Zero's file helpers for path management.
502 # """
503 # # Sanitize filename
504 # filename = files.safe_file_name(filename)
505
506 # # Generate unique filename if needed
507 # unique_id = uuid.uuid4().hex[:8]
508 # name, ext = os.path.splitext(filename)
509 # unique_filename = f"{name}_{unique_id}{ext}"
510
511 # # Build relative path and save
512 # relative_path = os.path.join(download_folder, unique_filename)
513 # files.write_file_bin(relative_path, content)
514
515 # # Return absolute path
516 # abs_path = files.get_abs_path(relative_path)
517 # return abs_path
518
519 # def _decode_header(self, header: str) -> str:
520 # """Decode email header handling various encodings."""
521 # if not header:
522 # return ""
523
524 # decoded_parts = []
525 # for part, encoding in decode_header(header):
526 # if isinstance(part, bytes):
527 # decoded_parts.append(part.decode(encoding or "utf-8", errors="ignore"))
528 # else:
529 # decoded_parts.append(str(part))
530
531 # return " ".join(decoded_parts)
532
533
534 # async def read_messages(
535 # account_type: str = "imap",
536 # server: str = "",
537 # port: int = 993,
538 # username: str = "",
539 # password: str = "",
540 # download_folder: str = "usr/email",
541 # options: Optional[Dict[str, Any]] = None,
542 # filter: Optional[Dict[str, Any]] = None,
543 # ) -> List[Message]:
544 # """
545 # Convenience wrapper for reading email messages.
546
547 # Automatically handles connection and disconnection.
548
549 # Args:
550 # account_type: "imap" or "exchange"
551 # server: Mail server address
552 # port: Server port (default 993 for IMAP SSL)
553 # username: Email username
554 # password: Email password
555 # download_folder: Folder to save attachments (relative to /a0/)
556 # options: Optional configuration dict
557 # filter: Filter criteria dict
558
559 # Returns:
560 # List of Message objects
561
562 # Example:
563 # from helpers.email_client import read_messages
564 # messages = await read_messages(
565 # server="imap.gmail.com",
566 # port=993,
567 # username=secrets.get("EMAIL_USER"),
568 # password=secrets.get("EMAIL_PASSWORD"),
569 # download_folder="tmp/email/inbox",
570 # filter={"unread": True, "sender": "*@company.com"}
571 # )
572 # """
573 # client = EmailClient(
574 # account_type=account_type,
575 # server=server,
576 # port=port,
577 # username=username,
578 # password=password,
579 # options=options,
580 # )
581
582 # try:
583 # await client.connect()
584 # messages = await client.read_messages(download_folder, filter)
585 # return messages
586 # finally:
587 # await client.disconnect()