Files
mytoolkit/bin/commands/mail.py
T
Zhengshou Lai ffaf7e58f2 fix(mail): improve attachment filename and reply body handling
- Use RFC 2047 encoded-word for non-ASCII attachment filenames instead of
  RFC 2231 filename* segments, which Coremail Web fails to display.
- Build multipart messages with MIMEMultipart when attachments are present
  to keep address headers and attachment filenames Coremail-compatible.
- Add --attach support to reply and forward commands.
- Quote the original message body in reply drafts for proper threading.
2026-06-23 10:24:23 +08:00

520 lines
16 KiB
Python

"""Mail client for SYSU Coremail (IMAP/SMTP)."""
import base64
import imaplib
import json
import os
import re
import smtplib
from datetime import datetime
from email import policy
from email.header import Header, decode_header
from email.message import EmailMessage
from email.mime.application import MIMEApplication
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.parser import BytesParser
from email.utils import parseaddr
from pathlib import Path
import click
from bin.utils import handle_errors
_HOME = Path(os.environ.get("MYTOOLKIT_HOME", Path.home() / ".mytoolkit"))
CONFIG_PATH = _HOME / "mail.json"
def load_config():
"""Load mail config."""
if not CONFIG_PATH.exists():
click.echo("Mail not configured. Run: mytoolkit mail config", err=True)
raise click.Abort()
with open(CONFIG_PATH, encoding="utf-8") as f:
return json.load(f)
def decode_password(pwd: str) -> str:
"""Decode base64-encoded password."""
try:
return base64.b64decode(pwd).decode("utf-8")
except Exception:
return pwd
def get_imap(cfg: dict):
"""Connect to IMAP server."""
mail = imaplib.IMAP4_SSL(cfg["imap_host"], cfg.get("imap_port", 993))
pwd = decode_password(cfg["password"])
mail.login(cfg["email"], pwd)
return mail
def get_smtp(cfg: dict):
"""Connect to SMTP server."""
smtp = smtplib.SMTP_SSL(cfg["smtp_host"], cfg.get("smtp_port", 465))
pwd = decode_password(cfg["password"])
smtp.login(cfg["email"], pwd)
return smtp
def decode_mime_str(s: str) -> str:
"""Decode MIME-encoded header string (e.g. =?utf-8?...?=) to plain text."""
if not s:
return ""
parts = decode_header(s)
out = []
for value, charset in parts:
if isinstance(value, bytes):
try:
out.append(value.decode(charset or "utf-8", errors="replace"))
except (LookupError, TypeError):
out.append(value.decode("utf-8", errors="replace"))
else:
out.append(value)
return "".join(out)
def uid_fetch(imap, uid, parts):
"""Fetch by real IMAP UID. Returns raw bytes or empty bytes if not found."""
typ, msg_data = imap.uid("FETCH", str(uid), parts)
if typ != "OK" or not msg_data or not msg_data[0]:
return b""
return msg_data[0][1]
def safe_filename(name: str) -> str:
"""Strip path separators from a filename to prevent traversal."""
name = decode_mime_str(name).strip()
name = name.replace("/", "_").replace("\\", "_")
if name in ("", ".", ".."):
return "attachment"
return name
def _encode_filename(filename: str) -> str:
"""Return a Content-Disposition filename value compatible with Coremail.
Coremail Web client does not reliably display RFC 2231 filename* segments.
We use RFC 2047 encoded-word for the plain filename parameter, which older
Chinese university webmail clients generally handle better.
"""
ascii_name = filename.encode("ascii", "ignore").decode()
if ascii_name == filename:
return f'filename="{filename}"'
encoded = Header(filename, charset="utf-8").encode()
return f'filename="{encoded}"'
def _quote_body(text: str) -> str:
"""Prefix each line of a body with '> ' for reply quoting."""
return "\n".join(f"> {line}" for line in text.splitlines())
def _build_reply_body(original, body: str) -> str:
"""Build a reply body that quotes the original message."""
orig_body = ""
b = original.get_body(preferencelist=("plain",))
if b:
orig_body = b.get_content() or ""
quoted = _quote_body(orig_body) if orig_body else ""
return (
f"{body}\n\n"
f"--- Original message ---\n"
f"From: {original['From']}\n"
f"Date: {original['Date']}\n"
f"Subject: {original['Subject']}\n\n"
f"{quoted}"
)
def _references(original) -> str:
"""Build References header for a reply."""
refs = []
if original["References"]:
refs.append(str(original["References"]))
if original["Message-ID"]:
refs.append(str(original["Message-ID"]))
return " ".join(refs)
def _build_address_header(raw_address: str) -> str:
"""Encode a 'Display Name <addr>' header for Coremail compatibility.
Python's compat32 policy folds long encoded words incorrectly when the
address contains both non-ASCII display name and angle brackets. We encode
the display name manually and preserve the raw address outside the encoded
word.
"""
if not raw_address:
return ""
display_name, addr_spec = parseaddr(raw_address)
if not display_name or display_name == addr_spec:
return addr_spec
encoded = base64.b64encode(display_name.encode("utf-8")).decode("ascii")
return f"=?utf-8?b?{encoded}?= <{addr_spec}>"
def _build_message(
cfg: dict,
headers: dict,
body: str,
html: bool,
attach_paths,
) -> bytes:
"""Build a MIME message and return its serialized bytes.
Uses MIMEMultipart whenever attachments are present so that non-ASCII
display names in address headers and non-ASCII attachment filenames are
both serialized in a Coremail-compatible way.
"""
attach_paths = [p for p in attach_paths if p]
if attach_paths:
msg = MIMEMultipart("mixed")
msg["From"] = cfg["email"]
for key, value in headers.items():
if key in ("To", "Cc") and value:
msg[key] = _build_address_header(value)
elif key == "Subject" and value:
msg[key] = Header(value, charset="utf-8").encode()
elif value:
msg[key] = value
if html:
text_part = MIMEText("Draft includes HTML version.", "plain", "utf-8")
html_part = MIMEText(body, "html", "utf-8")
msg.attach(text_part)
msg.attach(html_part)
else:
text_part = MIMEText(body, "plain", "utf-8")
msg.attach(text_part)
for path in attach_paths:
p = Path(path)
if not p.exists():
click.echo(f"Attachment not found: {path}", err=True)
continue
with open(p, "rb") as f:
data = f.read()
part = MIMEApplication(data)
part["Content-Disposition"] = f"attachment; {_encode_filename(p.name)}"
msg.attach(part)
return msg.as_bytes()
# No attachments: use EmailMessage with default policy.
msg = EmailMessage()
msg["From"] = cfg["email"]
for key, value in headers.items():
if key == "Subject" and value:
msg[key] = value
elif value:
msg[key] = value
if html:
msg.set_content("Draft includes HTML version.", subtype="plain")
msg.add_alternative(body, subtype="html")
else:
msg.set_content(body)
return msg.as_bytes()
@click.group()
def mail():
"""Mail client (IMAP/SMTP). No send command — use draft instead."""
pass
@mail.command()
@click.option("--imap-host", default="mail.sysu.edu.cn", help="IMAP server host")
@click.option("--imap-port", default=993, type=int, help="IMAP server port")
@click.option("--smtp-host", default="mail.sysu.edu.cn", help="SMTP server host")
@click.option("--smtp-port", default=465, type=int, help="SMTP server port")
@click.option("--email", prompt="Email address", help="Email address")
@click.option("--password", prompt="Password", hide_input=True, help="Email password")
@handle_errors
def config(imap_host, imap_port, smtp_host, smtp_port, email, password):
"""Configure mail account."""
CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
cfg = {
"imap_host": imap_host,
"imap_port": imap_port,
"smtp_host": smtp_host,
"smtp_port": smtp_port,
"email": email,
"password": base64.b64encode(password.encode()).decode(),
}
with open(CONFIG_PATH, "w", encoding="utf-8") as f:
json.dump(cfg, f, indent=2)
click.echo(f"Config saved to {CONFIG_PATH}")
@mail.command()
@click.option("--limit", "-n", default=20, type=int, help="Number of messages to show")
@click.option("--unseen", "-u", is_flag=True, help="Show only unseen messages")
@handle_errors
def inbox(limit, unseen):
"""List inbox messages. UID column is the real IMAP UID (stable across sessions)."""
cfg = load_config()
imap = get_imap(cfg)
imap.select("INBOX", readonly=True)
criterion = "UNSEEN" if unseen else "ALL"
_, data = imap.uid("SEARCH", None, criterion)
uids = data[0].split()
uids = uids[-limit:] if len(uids) > limit else uids
click.echo(f"{'UID':>10} {'Date':<20} {'From':<30} Subject")
click.echo("-" * 100)
for uid in reversed(uids):
raw_header = uid_fetch(imap, uid.decode(), "(RFC822.HEADER)")
header = BytesParser(policy=policy.default).parsebytes(raw_header)
subject = header["Subject"] or "(no subject)"
from_addr = header["From"] or ""
from_addr = re.sub(r"<.*>", "", from_addr).strip() or from_addr
from_addr = from_addr[:28]
date_str = ""
if header["Date"]:
try:
dt = datetime.strptime(header["Date"][:31], "%a, %d %b %Y %H:%M:%S")
date_str = dt.strftime("%Y-%m-%d %H:%M")
except ValueError:
date_str = str(header["Date"])[:19]
click.echo(f"{uid.decode():>10} {date_str:<20} {from_addr:<30} {subject[:50]}")
imap.close()
imap.logout()
@mail.command()
@click.argument("uid", type=int)
@click.option("--raw", is_flag=True, help="Show raw headers")
@handle_errors
def read(uid, raw):
"""Read a message by IMAP UID."""
cfg = load_config()
imap = get_imap(cfg)
imap.select("INBOX", readonly=True)
raw_msg = uid_fetch(imap, uid, "(RFC822)")
msg = BytesParser(policy=policy.default).parsebytes(raw_msg)
if raw:
click.echo(msg.as_string())
else:
click.echo(f"From: {msg['From']}")
click.echo(f"To: {msg['To']}")
click.echo(f"Date: {msg['Date']}")
click.echo(f"Subject: {msg['Subject']}")
attachments = [
safe_filename(p.get_filename() or "")
for p in msg.walk()
if p.get_content_disposition() == "attachment" and p.get_filename()
]
if attachments:
click.echo(f"Attachments: {', '.join(attachments)}")
click.echo("-" * 60)
body = msg.get_body(preferencelist=("plain",))
if body:
click.echo(body.get_content())
else:
click.echo("(no plain text body)")
imap.close()
imap.logout()
@mail.command()
@click.argument("uid", type=int)
@click.option("--out", "-o", "out_dir", default=".", help="Output directory (default: current dir)")
@click.option("--inline", is_flag=True, help="Also save inline parts (e.g. embedded images)")
@handle_errors
def download(uid, out_dir, inline):
"""Download attachments from a message by IMAP UID."""
cfg = load_config()
imap = get_imap(cfg)
imap.select("INBOX", readonly=True)
raw_msg = uid_fetch(imap, uid, "(RFC822)")
imap.close()
imap.logout()
if not raw_msg:
click.echo(f"Message UID {uid} not found.", err=True)
raise click.Abort()
msg = BytesParser(policy=policy.default).parsebytes(raw_msg)
out = Path(out_dir).expanduser()
out.mkdir(parents=True, exist_ok=True)
wanted = {"attachment"}
if inline:
wanted.add("inline")
saved = []
for part in msg.walk():
cd = part.get_content_disposition()
if cd not in wanted:
continue
fn = part.get_filename()
if not fn:
continue
fn = safe_filename(fn)
target = out / fn
# avoid overwriting existing files
i = 1
while target.exists():
stem, suffix = target.stem, target.suffix
target = out / f"{stem}-{i}{suffix}"
i += 1
payload = part.get_payload(decode=True)
if not isinstance(payload, (bytes, bytearray)):
continue
target.write_bytes(payload)
saved.append((target, len(payload), cd))
if not saved:
click.echo(f"No attachments in UID {uid}.")
return
for target, size, cd in saved:
marker = "[inline]" if cd == "inline" else " "
click.echo(f"{marker} {target} ({size} bytes)")
@mail.command()
@click.option("--to", required=True, help="Recipient address")
@click.option("--cc", help="CC addresses (comma-separated)")
@click.option("--subject", "-s", required=True, help="Subject")
@click.option("--body", "-b", required=True, help="Body text (\\n expanded)")
@click.option("--attach", multiple=True, help="Attachment path(s)")
@click.option("--html", is_flag=True, help="Body is HTML")
@handle_errors
def draft(to, cc, subject, body, attach, html):
"""Create a draft email (does NOT send). Saves to drafts folder."""
cfg = load_config()
body = body.replace("\\n", "\n").replace("\\t", "\t")
headers = {"To": to, "Subject": subject}
if cc:
headers["Cc"] = cc
raw = _build_message(cfg, headers, body, html, attach)
imap = get_imap(cfg)
imap.select("INBOX")
imap.append("Drafts", None, None, raw)
imap.close()
imap.logout()
click.echo("Draft saved to Drafts folder.")
click.echo(f" To: {to}")
click.echo(f" Subject: {subject}")
@mail.command()
@click.argument("uid", type=int)
@click.option("--body", "-b", required=True, help="Reply body")
@click.option("--attach", multiple=True, help="Attachment path(s)")
@click.option("--html", is_flag=True, help="Body is HTML")
@handle_errors
def reply(uid, body, attach, html):
"""Reply to a message by IMAP UID (draft mode — saves to Drafts)."""
cfg = load_config()
body = body.replace("\\n", "\n").replace("\\t", "\t")
imap = get_imap(cfg)
imap.select("INBOX", readonly=True)
raw_orig = uid_fetch(imap, uid, "(RFC822)")
if not raw_orig:
imap.close()
imap.logout()
click.echo(f"Message UID {uid} not found.", err=True)
raise click.Abort()
original = BytesParser(policy=policy.default).parsebytes(raw_orig)
subject = original["Subject"] or ""
if not str(subject).startswith("Re:"):
subject = f"Re: {subject}"
headers = {
"To": original["Reply-To"] or original["From"],
"Subject": subject,
"In-Reply-To": original["Message-ID"],
"References": _references(original),
}
content = _build_reply_body(original, body)
if html:
# For HTML replies we still include the plain quoted body as plain text
# alternative and an HTML rendering as primary alternative.
content = _build_reply_body(original, body)
raw = _build_message(cfg, headers, content.replace("\n", "<br>\n"), True, attach)
else:
raw = _build_message(cfg, headers, content, False, attach)
imap.append("Drafts", None, None, raw)
imap.close()
imap.logout()
click.echo(f"Reply draft saved for UID {uid}.")
@mail.command()
@click.argument("uid", type=int)
@click.option("--to", required=True, help="Forward to")
@click.option("--body", "-b", default="", help="Additional body text")
@click.option("--attach", multiple=True, help="Attachment path(s)")
@click.option("--html", is_flag=True, help="Body is HTML")
@handle_errors
def forward(uid, to, body, attach, html):
"""Forward a message by IMAP UID (draft mode — saves to Drafts)."""
cfg = load_config()
body = body.replace("\\n", "\n").replace("\\t", "\t")
imap = get_imap(cfg)
imap.select("INBOX", readonly=True)
raw_orig = uid_fetch(imap, uid, "(RFC822)")
if not raw_orig:
imap.close()
imap.logout()
click.echo(f"Message UID {uid} not found.", err=True)
raise click.Abort()
original = BytesParser(policy=policy.default).parsebytes(raw_orig)
orig_body = ""
b = original.get_body(preferencelist=("plain",))
if b:
orig_body = b.get_content()
subject = f"Fwd: {original['Subject']}"
content = (
f"{body}\n\n"
f"--- Forwarded message ---\n"
f"From: {original['From']}\n"
f"Date: {original['Date']}\n"
f"Subject: {original['Subject']}\n\n"
f"{orig_body}"
)
headers = {"To": to, "Subject": subject}
if html:
raw = _build_message(cfg, headers, content.replace("\n", "<br>\n"), True, attach)
else:
raw = _build_message(cfg, headers, content, False, attach)
imap.append("Drafts", None, None, raw)
imap.close()
imap.logout()
click.echo(f"Forward draft saved for UID {uid}.")