add progress ui to fetch

This commit is contained in:
Tim Bendt
2025-05-06 11:29:52 -06:00
parent f4903495e3
commit a23de0b8d2

View File

@@ -10,8 +10,13 @@ from dateutil import parser
from dateutil.tz import UTC
from email.message import EmailMessage
from email.utils import format_datetime
from rich import print
from rich.progress import Progress, track, SpinnerColumn, MofNCompleteColumn
from rich.panel import Panel
import time
import html2text
import asyncio
import argparse
# Filepath for caching timestamp
cache_timestamp_file = 'cache_timestamp.json'
@@ -32,7 +37,14 @@ def save_sync_timestamp():
with open(sync_timestamp_file, 'w') as f:
json.dump({'last_sync': time.time()}, f)
def synchronize_maildir(maildir_path, headers):
# Add argument parsing for dry-run mode
arg_parser = argparse.ArgumentParser(description="Fetch and synchronize emails.")
arg_parser.add_argument("--dry-run", action="store_true", help="Run in dry-run mode without making changes.", default=True)
args = arg_parser.parse_args()
dry_run = args.dry_run
async def mark_read_async(maildir_path, headers, progress, task_id):
last_sync = load_last_sync_timestamp()
current_time = time.time()
@@ -43,36 +55,83 @@ def synchronize_maildir(maildir_path, headers):
cur_files = set(glob.glob(os.path.join(cur_dir, '*.eml')))
moved_to_cur = [os.path.basename(f) for f in cur_files - new_files]
progress.update(task_id, total=len(moved_to_cur))
for filename in moved_to_cur:
message_id = filename.split('.')[0] # Extract the Message-ID from the filename
print(f"Marking message as read: {message_id}")
if not dry_run:
response = requests.patch(
f'https://graph.microsoft.com/v1.0/me/messages/{message_id}',
headers=headers,
json={'isRead': True}
)
if response.status_code != 200:
print(f"Failed to mark message as read: {message_id}, {response.status_code}, {response.text}")
progress.console.print(Panel(f"Failed to mark message as read: {message_id}, {response.status_code}, {response.text}", padding=2, border_style="red"))
else:
progress.console.print(f"[DRY-RUN] Would mark message as read: {message_id}")
progress.advance(task_id)
# Find messages moved to ".Trash/cur" and delete them on the server
trash_dir = os.path.join(maildir_path, '.Trash', 'cur')
trash_files = set(glob.glob(os.path.join(trash_dir, '*.eml')))
for filepath in trash_files:
message_id = os.path.basename(filepath).split('.')[0] # Extract the Message-ID from the filename
print(f"Moving message to trash: {message_id}")
response = requests.delete(
f'https://graph.microsoft.com/v1.0/me/messages/{message_id}',
headers=headers
)
if response.status_code == 204: # 204 No Content indicates success
os.remove(filepath) # Remove the file from local trash
# Save the current sync timestamp
if not dry_run:
save_sync_timestamp()
else:
progress.console.print("[DRY-RUN] Would save sync timestamp.")
async def fetch_mail_async(maildir_path, attachments_dir, headers, progress, task_id):
mail_url = 'https://graph.microsoft.com/v1.0/me/mailFolders/inbox/messages?$top=100&$orderby=receivedDateTime asc&$select=id,subject,from,toRecipients,ccRecipients,receivedDateTime,isRead,body,attachments'
messages = []
# Fetch the total count of messages in the inbox
inbox_url = 'https://graph.microsoft.com/v1.0/me/mailFolders/inbox'
response = requests.get(inbox_url, headers=headers)
if response.status_code != 200:
raise Exception(f"Failed to fetch inbox details: {response.status_code} {response.text}")
total_messages = response.json().get('totalItemCount', 0)
progress.update(task_id, total=total_messages)
while mail_url:
response = requests.get(mail_url, headers=headers)
if response.status_code != 200:
raise Exception(f"Failed to fetch mail: {response.status_code} {response.text}")
response_data = response.json()
messages.extend(response_data.get('value', []))
progress.advance(task_id, (len(response_data.get('value', [])) / 2))
# Get the next page URL from @odata.nextLink
mail_url = response_data.get('@odata.nextLink')
inbox_msg_ids = set(message['id'] for message in messages)
progress.update(task_id, completed=(len(messages) / 2))
new_dir = os.path.join(maildir_path, 'new')
cur_dir = os.path.join(maildir_path, 'cur')
new_files = set(glob.glob(os.path.join(new_dir, '*.eml')))
cur_files = set(glob.glob(os.path.join(cur_dir, '*.eml')))
for filename in Set.union(cur_files, new_files):
message_id = filename.split('.')[0].split('/')[-1] # Extract the Message-ID from the filename
if (message_id not in inbox_msg_ids):
if not dry_run:
progress.console.print(f"Deleting {filename} from inbox")
os.remove(filename)
else:
progress.console.print(f"[DRY-RUN] Would delete {filename} from inbox")
for message in messages:
progress.console.print(f"Processing message: {message.get('subject', 'No Subject')}", end='\r')
save_email_to_maildir(maildir_path, message, attachments_dir, progress)
progress.update(task_id, advance=0.5)
progress.update(task_id, completed=len(messages))
progress.console.print(f"\nFinished saving {len(messages)} messages.")
async def archive_mail_async(maildir_path, headers, progress, task_id):
# Find messages moved to ".Archives/**/*" and move them to the "Archive" folder on the server
archive_dir = os.path.join(maildir_path, '.Archives')
archive_files = glob.glob(os.path.join(archive_dir, '**', '*.eml'), recursive=True)
progress.update(task_id, total=len(archive_files))
# Fetch the list of folders to find the "Archive" folder ID
print("Fetching server folders to locate 'Archive' folder...")
progress.console.print("Fetching server folders to locate 'Archive' folder...")
folder_response = requests.get('https://graph.microsoft.com/v1.0/me/mailFolders', headers=headers)
if folder_response.status_code != 200:
raise Exception(f"Failed to fetch mail folders: {folder_response.status_code}, {folder_response.text}")
@@ -89,28 +148,124 @@ def synchronize_maildir(maildir_path, headers):
for filepath in archive_files:
message_id = os.path.basename(filepath).split('.')[0] # Extract the Message-ID from the filename
print(f"Moving message to 'Archive' folder: {message_id}")
progress.console.print(f"Moving message to 'Archive' folder: {message_id}")
if not dry_run:
response = requests.post(
f'https://graph.microsoft.com/v1.0/me/messages/{message_id}/move',
headers=headers,
json={'destinationId': archive_folder_id}
)
if response.status_code != 201: # 201 Created indicates success
print(f"Failed to move message to 'Archive': {message_id}, {response.status_code}, {response.text}")
progress.console.print(f"Failed to move message to 'Archive': {message_id}, {response.status_code}, {response.text}")
if response.status_code == 404:
os.remove(filepath) # Remove the file from local archive if not found on server
# Save the current sync timestamp
save_sync_timestamp()
# Load cached timestamp if it exists
if os.path.exists(cache_timestamp_file):
with open(cache_timestamp_file, 'r') as f:
cache_timestamp = json.load(f)
else:
cache_timestamp = {}
progress.console.print(f"[DRY-RUN] Would move message to 'Archive' folder: {message_id}")
progress.advance(task_id)
return
async def delete_mail_async(maildir_path, headers, progress, task_id):
# Find messages moved to ".Trash/cur" and delete them on the server
trash_dir = os.path.join(maildir_path, '.Trash', 'cur')
trash_files = set(glob.glob(os.path.join(trash_dir, '*.eml')))
progress.update(task_id, total=len(trash_files))
for filepath in trash_files:
message_id = os.path.basename(filepath).split('.')[0] # Extract the Message-ID from the filename
if not dry_run:
progress.console.print(f"Moving message to trash: {message_id}")
response = requests.delete(
f'https://graph.microsoft.com/v1.0/me/messages/{message_id}',
headers=headers
)
if response.status_code == 204: # 204 No Content indicates success
os.remove(filepath) # Remove the file from local trash
else:
progress.console.print(f"[DRY-RUN] Would delete message: {message_id}")
progress.advance(task_id)
async def fetch_calendar_async(headers, progress, task_id):
total_event_url = 'https://graph.microsoft.com/v1.0/me/events?$count=true'
total = requests.get(total_event_url, headers=headers)
if (total.status_code != 200):
raise Exception(f"Failed to fetch count: {response.status_code} {response.text}")
total_events = total.json().get('@odata.count', 0)
progress.update(task_id, total=total_events)
calendar_url = 'https://graph.microsoft.com/v1.0/me/events?$top=100&$orderby=start/dateTime asc'
events = []
while calendar_url:
response = requests.get(calendar_url, headers=headers)
if response.status_code != 200:
raise Exception(f"Failed to fetch calendar: {response.status_code} {response.text}")
response_data = response.json()
events.extend(response_data.get('value', []))
progress.advance(task_id, len(response_data.get('value', [])))
# Get the next page URL from @odata.nextLink
calendar_url = response_data.get('@odata.nextLink')
# Call the synchronization function before fetching mail
async def download_calendar_events(headers, progress, task_id):
# Fetch the total count of events in the calendar
total_event_url = 'https://graph.microsoft.com/v1.0/me/events?$count=true'
total = requests.get(total_event_url, headers=headers)
if (response.status_code != 200):
raise Exception(f"Failed to fetch count: {response.status_code} {response.text}")
total_events = total.json().get('@odata.count', 0)
progress.update(task_id, total=total_events)
print(f"Total events in calendar: {total_events}")
# Fetch events with pagination and expand recurring events
events_url = 'https://graph.microsoft.com/v1.0/me/events?$top=100&$expand=instances'
events = []
progress.console.print("Fetching Calendar events...")
while events_url:
response = requests.get(events_url, headers=headers)
response_data = response.json()
events.extend(response_data.get('value', []))
# print(f"Fetched {len(events)} events so far...", end='\r')
events_url = response_data.get('@odata.nextLink')
progress.advance(task_id, len(response_data.get('value', [])))
# Save events to a file in iCalendar format
output_file = f'output_ics/outlook_events_latest.ics'
if not dry_run:
os.makedirs(os.path.dirname(output_file), exist_ok=True)
progress.console.print(f"Saving events to {output_file}...")
with open(output_file, 'w') as f:
f.write("BEGIN:VCALENDAR\nVERSION:2.0\n")
for event in events:
if 'start' in event and 'end' in event:
start = parser.isoparse(event['start']['dateTime'])
end = parser.isoparse(event['end']['dateTime'])
f.write(f"BEGIN:VEVENT\nSUMMARY:{event['subject']}\n")
f.write(f"DTSTART:{start.strftime('%Y%m%dT%H%M%S')}\n")
f.write(f"DTEND:{end.strftime('%Y%m%dT%H%M%S')}\n")
if 'recurrence' in event and event['recurrence']: # Check if 'recurrence' exists and is not None
for rule in event['recurrence']:
if rule.startswith('RRULE'):
rule_parts = rule.split(';')
new_rule_parts = []
for part in rule_parts:
if part.startswith('UNTIL='):
until_value = part.split('=')[1]
until_date = parser.isoparse(until_value)
if start.tzinfo is not None and until_date.tzinfo is None:
until_date = until_date.replace(tzinfo=UTC)
new_rule_parts.append(f"UNTIL={until_date.strftime('%Y%m%dT%H%M%SZ')}")
else:
new_rule_parts.append(part)
rule = ';'.join(new_rule_parts)
f.write(f"{rule}\n")
f.write("END:VEVENT\n")
f.write("END:VCALENDAR\n")
progress.console.print(f"Saved events to {output_file}")
else:
progress.console.print(f"[DRY-RUN] Would save events to {output_file}")
# Function to check if the cache is still valid
def is_cache_valid():
@@ -119,15 +274,13 @@ def is_cache_valid():
cache_expiry_time = cache_timestamp['timestamp'] + cache_timestamp['max_age']
return current_time < cache_expiry_time
return False
# Function to create Maildir structure
def create_maildir_structure(base_path):
os.makedirs(os.path.join(base_path, 'cur'), exist_ok=True)
os.makedirs(os.path.join(base_path, 'new'), exist_ok=True)
os.makedirs(os.path.join(base_path, 'tmp'), exist_ok=True)
def save_email_to_maildir(maildir_path, email_data, attachments_dir):
def save_email_to_maildir(maildir_path, email_data, attachments_dir, progress):
# Create a new EmailMessage object
msg = EmailMessage()
@@ -159,15 +312,19 @@ def save_email_to_maildir(maildir_path, email_data, attachments_dir):
msg.set_content(body_markdown)
# Download attachments
progress.console.print(f"Downloading attachments for message: {msg['Message-ID']}")
for attachment in email_data.get('attachments', []):
attachment_id = attachment.get('id')
attachment_name = attachment.get('name', 'unknown')
attachment_content = attachment.get('contentBytes')
if attachment_content:
attachment_path = os.path.join(attachments_dir, attachment_name)
if not dry_run:
with open(attachment_path, 'wb') as f:
f.write(attachment_content.encode('utf-8'))
msg.add_attachment(attachment_content.encode('utf-8'), filename=attachment_name)
else:
progress.console.print(f"[DRY-RUN] Would save attachment to {attachment_path}")
# Determine the directory based on isRead
target_dir = 'cur' if email_data.get('isRead', False) else 'new'
@@ -177,13 +334,31 @@ def save_email_to_maildir(maildir_path, email_data, attachments_dir):
# Check if the file already exists in any subfolder
for root, _, files in os.walk(maildir_path):
if email_filename in files:
print(f"Message {msg['Message-ID']} already exists in {root}. Skipping save.")
progress.console.print(f"Message {msg['Message-ID']} already exists in {root}. Skipping save.")
return
# Save the email to the Maildir
if not dry_run:
with open(email_filepath, 'w') as f:
f.write(msg.as_string())
print(f"Saved message {msg['Message-ID']} to {email_filepath}")
progress.console.print(f"Saved message {msg['Message-ID']}")
else:
progress.console.print(f"[DRY-RUN] Would save message {msg['Message-ID']}")
async def main():
# Load cached timestamp if it exists
if os.path.exists(cache_timestamp_file):
with open(cache_timestamp_file, 'r') as f:
cache_timestamp = json.load(f)
else:
cache_timestamp = {}
# Save emails to Maildir
maildir_path = os.getenv('MAILDIR_PATH', os.path.expanduser('~/Mail')) + "/corteva"
attachments_dir = os.path.join(maildir_path, 'attachments')
os.makedirs(attachments_dir, exist_ok=True)
create_maildir_structure(maildir_path)
# Read Azure app credentials from environment variables
client_id = os.getenv('AZURE_CLIENT_ID')
@@ -192,8 +367,6 @@ tenant_id = os.getenv('AZURE_TENANT_ID')
if not client_id or not tenant_id:
raise ValueError("Please set the AZURE_CLIENT_ID and AZURE_TENANT_ID environment variables.")
# Token cache
cache = msal.SerializableTokenCache()
cache_file = 'token_cache.bin'
@@ -223,7 +396,7 @@ else:
flow = app.initiate_device_flow(scopes=scopes)
if 'user_code' not in flow:
raise Exception("Failed to create device flow")
print(flow['message'])
print(Panel(flow['message'], border_style="magenta", padding=2, title="MSAL Login Flow Link"))
token_response = app.acquire_token_by_device_flow(flow)
if 'access_token' not in token_response:
@@ -240,124 +413,27 @@ accounts = app.get_accounts()
if not accounts:
raise Exception("No accounts found")
# Call the synchronization function before fetching mail
print("Synchronizing maildir with server...")
synchronize_maildir(maildir_path=os.getenv('MAILDIR_PATH', os.path.expanduser('~/Mail')) + "/corteva", headers=headers)
print("Synchronization complete.")
mail_url = 'https://graph.microsoft.com/v1.0/me/mailFolders/inbox/messages?$top=100&$orderby=receivedDateTime asc&$select=id,subject,from,toRecipients,ccRecipients,receivedDateTime,isRead,body,attachments'
messages = []
print("Fetching mail...")
# Fetch the total count of messages in the inbox
inbox_url = 'https://graph.microsoft.com/v1.0/me/mailFolders/inbox'
response = requests.get(inbox_url, headers=headers)
if response.status_code != 200:
raise Exception(f"Failed to fetch inbox details: {response.status_code} {response.text}")
total_messages = response.json().get('totalItemCount', 0)
print(f"Total messages in inbox: {total_messages}")
while mail_url:
if is_cache_valid():
print("Using cached messages...")
break # No need to fetch further, cache is still valid
response = requests.get(mail_url, headers=headers)
if response.status_code != 200:
raise Exception(f"Failed to fetch mail: {response.status_code} {response.text}")
# Parse the Cache-Control header to get the max-age value
cache_control = response.headers.get('Cache-Control', '')
max_age = 0
if 'max-age=' in cache_control:
max_age = int(cache_control.split('max-age=')[1].split(',')[0])
# Update the cache timestamp and max-age
cache_timestamp['timestamp'] = time.time()
cache_timestamp['max_age'] = max_age
with open(cache_timestamp_file, 'w') as f:
json.dump(cache_timestamp, f)
# Process the response
response_data = response.json()
messages.extend(response_data.get('value', [])) # Add the current page of messages to the list
# Calculate and display progress percentage
progress = (len(messages) / total_messages) * 100 if total_messages > 0 else 0
print(f"Fetched {len(messages)} of {total_messages} messages ({progress:.2f}%)", end='\r')
# Get the next page URL from @odata.nextLink
mail_url = response_data.get('@odata.nextLink')
print("\nFinished fetching mail. Now saving them to maildir.")
# Save emails to Maildir
maildir_path = os.getenv('MAILDIR_PATH', os.path.expanduser('~/Mail')) + "/corteva"
attachments_dir = os.path.join(maildir_path, 'attachments')
os.makedirs(attachments_dir, exist_ok=True)
create_maildir_structure(maildir_path)
for message in messages:
print(f"Processing message: {message.get('subject', 'No Subject')}", end='\r')
save_email_to_maildir(maildir_path, message, attachments_dir)
progress = Progress(
SpinnerColumn(),
MofNCompleteColumn(),
*Progress.get_default_columns()
)
with progress:
task_fetch = progress.add_task("[green]Syncing Inbox...", total=0)
task_calendar = progress.add_task("[cyan]Fetching calendar...", total=0)
task_read = progress.add_task("[blue]Marking as read...", total=0)
task_archive = progress.add_task("[yellow]Archiving mail...", total=0)
task_delete = progress.add_task("[red]Deleting mail...", total=0)
print(f"\nFinished saving {len(messages)} messages.")
await asyncio.gather(
mark_read_async(maildir_path, headers, progress, task_read),
archive_mail_async(maildir_path, headers, progress, task_archive),
delete_mail_async(maildir_path, headers, progress, task_delete),
fetch_mail_async(maildir_path, attachments_dir, headers, progress, task_fetch),
fetch_calendar_async(headers, progress, task_calendar)
)
inbox_msg_ids = set(message['id'] for message in messages)
new_dir = os.path.join(maildir_path, 'new')
cur_dir = os.path.join(maildir_path, 'cur')
new_files = set(glob.glob(os.path.join(new_dir, '*.eml')))
cur_files = set(glob.glob(os.path.join(cur_dir, '*.eml')))
for filename in Set.union(cur_files, new_files):
message_id = filename.split('.')[0].split('/')[-1] # Extract the Message-ID from the filename
if (message_id not in inbox_msg_ids):
print(f"Deleting {filename} from inbox")
os.remove(filename)
# Fetch events with pagination and expand recurring events
events_url = 'https://graph.microsoft.com/v1.0/me/events?$top=100&$expand=instances'
events = []
print("Fetching Calendar events...")
while events_url:
response = requests.get(events_url, headers=headers)
response_data = response.json()
events.extend(response_data.get('value', []))
print(f"Fetched {len(events)} events so far...", end='\r')
events_url = response_data.get('@odata.nextLink')
# Save events to a file in iCalendar format
output_file = f'output_ics/outlook_events_latest.ics'
print(f"Saving events to {output_file}...")
with open(output_file, 'w') as f:
f.write("BEGIN:VCALENDAR\nVERSION:2.0\n")
for event in events:
if 'start' in event and 'end' in event:
start = parser.isoparse(event['start']['dateTime'])
end = parser.isoparse(event['end']['dateTime'])
f.write(f"BEGIN:VEVENT\nSUMMARY:{event['subject']}\n")
f.write(f"DTSTART:{start.strftime('%Y%m%dT%H%M%S')}\n")
f.write(f"DTEND:{end.strftime('%Y%m%dT%H%M%S')}\n")
if 'recurrence' in event and event['recurrence']: # Check if 'recurrence' exists and is not None
for rule in event['recurrence']:
if rule.startswith('RRULE'):
rule_parts = rule.split(';')
new_rule_parts = []
for part in rule_parts:
if part.startswith('UNTIL='):
until_value = part.split('=')[1]
until_date = parser.isoparse(until_value)
if start.tzinfo is not None and until_date.tzinfo is None:
until_date = until_date.replace(tzinfo=UTC)
new_rule_parts.append(f"UNTIL={until_date.strftime('%Y%m%dT%H%M%SZ')}")
else:
new_rule_parts.append(part)
rule = ';'.join(new_rule_parts)
f.write(f"{rule}\n")
f.write("END:VEVENT\n")
f.write("END:VCALENDAR\n")
if __name__ == "__main__":
asyncio.run(main())