mcp-home/mail/server.py

"""MCP Mail Server — searches and reads IMAP backup emails."""

import os
import sys
import base64
import io
import contextlib
import imaplib
import mailbox
import gc
import time
import ctypes
from email.header import decode_header
from email.mime.text import MIMEText
from email.utils import formatdate, parsedate_to_datetime
from pathlib import Path
from typing import Annotated

from pydantic import Field
from mcp.server.fastmcp import FastMCP
from mcp.types import TextContent, ImageContent, EmbeddedResource, BlobResourceContents
from starlette.applications import Starlette
from starlette.routing import Mount

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from common import get_current_user, OAUTH_ROUTES, BearerAuthMiddleware
from pdfutil import pdf_to_content
from officeutil import office_to_content, is_office

from common import load_config as _lc
_cfg = _lc()
MAIL_ROOTS = _cfg["mail_roots"]

IMAP_ACCOUNTS = {
    "stefan": {
        "host": "imap.1blu.de",
        "user": "d370128_0-slohmaier",
        "password": _cfg["imap_passwords"]["stefan"],
        "drafts_folder": "Drafts",
        "from_addr": "stefan@slohmaier.de",
    },
    "kati": {
        "host": "imap.gmail.com",
        "user": "akpolke@gmail.com",
        "password": _cfg["imap_passwords"]["kati"],
        "drafts_folder": "[Google Mail]/Entw&APw-rfe",
        "from_addr": "akpolke@gmail.com",
    },
}

mcp = FastMCP("Mail", stateless_http=True,
              transport_security={"enable_dns_rebinding_protection": False})


try:
    _libc = ctypes.CDLL("libc.so.6")
except OSError:
    _libc = None


def _trim_memory():
    """gc + malloc_trim: Speicher nach grossen Operationen ans OS zurueckgeben."""
    gc.collect()
    if _libc is not None:
        try:
            _libc.malloc_trim(0)
        except Exception:
            pass


def _safe_decode(payload, charset):
    """Decode bytes robust gegen unbekannte/kaputte Charsets (z.B. 'x-unknown')."""
    if not isinstance(payload, bytes):
        return str(payload)
    for cs in (charset, "utf-8", "latin-1"):
        if not cs:
            continue
        try:
            return payload.decode(cs, errors="replace")
        except (LookupError, TypeError):
            continue
    # Letzter Fallback: latin-1 akzeptiert jedes Byte
    return payload.decode("latin-1", errors="replace")


def _decode_hdr(raw):
    if not raw:
        return ""
    parts = decode_header(str(raw))
    decoded = []
    for data, charset in parts:
        if isinstance(data, bytes):
            decoded.append(_safe_decode(data, charset))
        else:
            decoded.append(str(data))
    return " ".join(decoded)


def _get_body(msg):
    if msg.is_multipart():
        for part in msg.walk():
            if part.get_content_type() == "text/plain":
                payload = part.get_payload(decode=True)
                if payload:
                    return _safe_decode(payload, part.get_content_charset())
        for part in msg.walk():
            if part.get_content_type() == "text/html":
                payload = part.get_payload(decode=True)
                if payload:
                    return "[HTML] " + _safe_decode(payload, part.get_content_charset())
    else:
        payload = msg.get_payload(decode=True)
        if payload:
            return _safe_decode(payload, msg.get_content_charset())
    return ""


def _get_attachments(msg):
    """Return list of (index, filename, mimetype, size, part) for attachments."""
    attachments = []
    if not msg.is_multipart():
        return attachments
    idx = 0
    for part in msg.walk():
        if part.get_content_maintype() == "multipart":
            continue
        disp = str(part.get("Content-Disposition", ""))
        filename = part.get_filename()
        ctype = part.get_content_type()
        # Attachment = has filename, or explicit attachment disposition, or non-text inline
        is_attachment = (
            "attachment" in disp.lower()
            or (filename is not None)
            or (ctype not in ("text/plain", "text/html") and part.get_content_maintype() != "multipart")
        )
        if not is_attachment:
            continue
        payload = part.get_payload(decode=True)
        if payload is None:
            continue
        idx += 1
        name = _decode_hdr(filename) if filename else f"attachment_{idx}.{ctype.split('/')[-1]}"
        attachments.append({"index": idx, "filename": name, "mime": ctype, "size": len(payload), "payload": payload})
    return attachments


def _discover_accounts(user):
    root = MAIL_ROOTS.get(user)
    if not root or not os.path.isdir(root):
        return {}
    return {a: os.path.join(root, a) for a in sorted(os.listdir(root)) if os.path.isdir(os.path.join(root, a))}


def _discover_folders(acct_path):
    folders = []
    for entry in sorted(Path(acct_path).rglob("cur")):
        rel = str(entry.parent.relative_to(acct_path))
        if rel != "." and rel not in folders:
            folders.append(rel)
    if "INBOX" in folders:
        folders.remove("INBOX")
        folders.insert(0, "INBOX")
    return folders


def _open_folder(acct_path, folder_name):
    base = os.path.realpath(acct_path)
    path = os.path.realpath(os.path.join(base, folder_name))
    # Path-Traversal verhindern: muss innerhalb des Account-Verzeichnisses bleiben
    if path != base and not path.startswith(base + os.sep):
        return None
    return mailbox.Maildir(path, create=False) if os.path.isdir(path) else None


MAX_SCAN = 2000  # Obergrenze gesammelter Treffer (Speicher-Schutz bei breiten Queries)
DEADLINE_S = 120  # Wall-Clock-Limit: Server heilt sich selbst, nie wieder endloser Haenger


def _iter_messages(md):
    """Yield (key, msg); ueberspringt unparsbare Mails statt die Suche zu killen."""
    try:
        keys = list(md.keys())
    except Exception:
        return
    for key in keys:
        try:
            msg = md.get_message(key)
        except Exception:
            continue
        yield key, msg


def _date_sortkey(date_str):
    try:
        dt = parsedate_to_datetime(date_str)
        return dt.timestamp() if dt else 0.0
    except Exception:
        return 0.0


@mcp.tool()
def list_accounts() -> str:
    """List all email accounts with their folder count. Call this first to see available accounts."""
    user = get_current_user()
    if not user:
        return "Error: not authenticated"
    accounts = _discover_accounts(user)
    if not accounts:
        return "No mail accounts found"
    return "\n".join(f"{name}: {len(_discover_folders(path))} folders" for name, path in accounts.items())


@mcp.tool()
def list_folders(
    account: Annotated[str, Field(description="Account name from list_accounts, e.g. 'd370128_0-slohmaier' or 'gmail'")]
) -> str:
    """List all IMAP folders in a mail account (INBOX, Sent, Archive, etc.)."""
    user = get_current_user()
    if not user:
        return "Error: not authenticated"
    acct_path = _discover_accounts(user).get(account)
    if not acct_path:
        return f"Account not found: {account}. Use list_accounts to see available accounts."
    folders = _discover_folders(acct_path)
    return "\n".join(folders) if folders else "No folders"


@mcp.tool()
def search_mail(
    query: Annotated[str, Field(description="Search term — matches subject, from, to, and body. Example: 'Rechnung', 'Amazon', 'meeting'")],
    account: Annotated[str, Field(description="Filter by account name. Leave empty to search all accounts.")] = "",
    folder: Annotated[str, Field(description="Filter by folder name, e.g. 'INBOX', 'Sent'. Leave empty for all folders.")] = "",
    limit: Annotated[int, Field(description="Maximum number of results to return")] = 20,
) -> str:
    """Search emails by keyword. Returns date, sender, recipient, subject, and a location key for reading the full email with read_mail."""
    user = get_current_user()
    if not user:
        return "Error: not authenticated"
    query_lower = query.lower()
    matches = []
    truncated = False
    t0 = time.monotonic()
    for acct_name, acct_path in _discover_accounts(user).items():
        if account and account not in acct_name:
            continue
        for fld in _discover_folders(acct_path):
            if folder and folder.lower() not in fld.lower():
                continue
            md = _open_folder(acct_path, fld)
            if not md:
                continue
            qbytes = query_lower.encode("utf-8", "ignore")
            try:
                keys = list(md.keys())
            except Exception:
                keys = []
            for key in keys:
                if time.monotonic() - t0 > DEADLINE_S:
                    truncated = True
                    break
                # Schneller Roh-Byte-Vorfilter: nur Mails voll parsen, die den
                # Suchbegriff im Rohtext enthalten (grep-Tempo statt 13k Voll-Parses).
                try:
                    raw = md.get_bytes(key)
                except Exception:
                    continue
                if qbytes not in raw.lower():
                    continue
                try:
                    msg = md.get_message(key)
                    subj = _decode_hdr(msg.get("Subject", ""))
                    frm = _decode_hdr(msg.get("From", ""))
                    to = _decode_hdr(msg.get("To", ""))
                    date_str = msg.get("Date", "")
                    # Bestaetigen ueber dekodierte Header ODER Body (Roh-Treffer verifizieren)
                    if query_lower not in f"{subj} {frm} {to}".lower():
                        if query_lower not in _get_body(msg).lower():
                            continue
                    matches.append((date_str, frm, to, subj, acct_name, fld, key))
                except Exception:
                    continue
                if len(matches) >= MAX_SCAN:
                    truncated = True
                    break
            if truncated:
                break
        if truncated:
            break
    matches.sort(key=lambda m: _date_sortkey(m[0]), reverse=True)
    shown = matches[:limit]
    _trim_memory()
    if not shown:
        return "No results found"
    lines = [f"[{d}] {f} -> {t}\n  Subject: {s}\n  Account: {a}, Folder: {fl}, Key: {k}"
             for (d, f, t, s, a, fl, k) in shown]
    more = len(matches) - len(shown)
    if more > 0:
        if truncated:
            lines.append(f"... mindestens {more} weitere Treffer (Scan-Limit erreicht; Suche eingrenzen).")
        else:
            lines.append(f"... {more} weitere Treffer (Suche eingrenzen oder limit erhoehen).")
    return "\n\n".join(lines)


@mcp.tool()
def read_mail(
    account: Annotated[str, Field(description="Account name from search results, e.g. 'd370128_0-slohmaier'")],
    folder: Annotated[str, Field(description="Folder name from search results, e.g. 'INBOX'")],
    key: Annotated[str, Field(description="Message key from search results")],
) -> str:
    """Read the full content of a single email. Use the account, folder, and key values from search_mail results."""
    user = get_current_user()
    if not user:
        return "Error: not authenticated"
    acct_path = _discover_accounts(user).get(account)
    if not acct_path:
        return f"Account not found: {account}"
    md = _open_folder(acct_path, folder)
    if not md:
        return f"Folder not found: {folder}"
    msg = md.get(key)
    if not msg:
        return f"Message not found: {key}"
    out = (
        f"From: {_decode_hdr(msg.get('From', ''))}\n"
        f"To: {_decode_hdr(msg.get('To', ''))}\n"
        f"Cc: {_decode_hdr(msg.get('Cc', ''))}\n"
        f"Date: {msg.get('Date', '')}\n"
        f"Subject: {_decode_hdr(msg.get('Subject', ''))}\n\n"
        f"{_get_body(msg)[:50000]}"
    )
    attachments = _get_attachments(msg)
    if attachments:
        out += "\n\nAnhaenge (nutze read_attachment mit dem Index):\n"
        for a in attachments:
            out += f"  {a['index']}. {a['filename']} ({a['mime']}, {a['size']:,} bytes)\n"
    return out


@mcp.tool()
def read_attachment(
    account: Annotated[str, Field(description="Account name from search results")],
    folder: Annotated[str, Field(description="Folder name from search results")],
    key: Annotated[str, Field(description="Message key from search results")],
    attachment_index: Annotated[int, Field(description="Attachment number from the read_mail attachment list (1-based)")],
) -> list[TextContent | ImageContent | EmbeddedResource]:
    """Read an email attachment. Images shown inline, PDFs and Office documents (docx, xlsx, pptx) as extracted text, text directly, other formats as binary. Get the index from read_mail."""
    _trim_memory()
    user = get_current_user()
    if not user:
        return [TextContent(type="text", text="Error: not authenticated")]
    acct_path = _discover_accounts(user).get(account)
    if not acct_path:
        return [TextContent(type="text", text=f"Account not found: {account}")]
    md = _open_folder(acct_path, folder)
    if not md:
        return [TextContent(type="text", text=f"Folder not found: {folder}")]
    msg = md.get(key)
    if not msg:
        return [TextContent(type="text", text=f"Message not found: {key}")]
    attachments = _get_attachments(msg)
    att = next((a for a in attachments if a["index"] == attachment_index), None)
    if not att:
        return [TextContent(type="text", text=f"Anhang {attachment_index} nicht gefunden. {len(attachments)} Anhaenge vorhanden.")]
    mime = att["mime"]
    payload = att["payload"]
    if mime in ("image/jpeg", "image/png", "image/gif", "image/webp") and len(payload) < 10_000_000:
        return [ImageContent(type="image", data=base64.b64encode(payload).decode(), mimeType=mime)]
    if mime.startswith("text/"):
        return [TextContent(type="text", text=payload.decode("utf-8", errors="replace")[:100000])]
    if mime == "application/pdf" or att["filename"].lower().endswith(".pdf"):
        return pdf_to_content(payload, att["filename"], mime, uri=f"mail://attachment/{att['filename']}")
    if is_office(mime, att["filename"]):
        return office_to_content(payload, att["filename"], mime, uri=f"mail://attachment/{att['filename']}")
    return [EmbeddedResource(type="resource", resource=BlobResourceContents(
        uri=f"mail://attachment/{att['filename']}",
        blob=base64.b64encode(payload).decode(), mimeType=mime))]


@mcp.tool()
def create_draft(
    to: Annotated[str, Field(description="Recipient email address, e.g. 'max@example.com'")],
    subject: Annotated[str, Field(description="Email subject line")],
    body: Annotated[str, Field(description="Email body text (plain text)")],
    cc: Annotated[str, Field(description="CC recipients, comma-separated")] = "",
    account: Annotated[str, Field(description="Which account to create the draft in. Leave empty for default account.")] = "",
) -> str:
    """Create an email draft in the Drafts folder. The draft can then be reviewed and sent from Roundcube or phone. Does NOT send the email."""
    user = get_current_user()
    if not user: return "Error: not authenticated"
    acct = IMAP_ACCOUNTS.get(user)
    if not acct: return f"Kein IMAP-Account fuer {user} konfiguriert"

    msg = MIMEText(body, "plain", "utf-8")
    msg["From"] = acct["from_addr"]
    msg["To"] = to
    if cc:
        msg["Cc"] = cc
    msg["Subject"] = subject
    msg["Date"] = formatdate(localtime=True)
    msg["X-Mailer"] = "MCP Mail Server"

    try:
        imap = imaplib.IMAP4_SSL(acct["host"])
        imap.login(acct["user"], acct["password"])
        status, _ = imap.append(acct["drafts_folder"], "\\Draft", None, msg.as_bytes())
        imap.logout()
        if status == "OK":
            return f"Entwurf erstellt: An {to}, Betreff \"{subject}\". Oeffne Roundcube oder die Mail-App um ihn zu senden."
        return f"Fehler beim Erstellen des Entwurfs: {status}"
    except Exception as e:
        return f"IMAP-Fehler: {e}"


def create_app():
    from contextlib import asynccontextmanager
    mcp_app = mcp.streamable_http_app()

    @asynccontextmanager
    async def lifespan(app):
        async with contextlib.AsyncExitStack() as stack:
            await stack.enter_async_context(mcp_app.router.lifespan_context(mcp_app))
            yield

    routes = list(OAUTH_ROUTES) + [Mount("/", app=mcp_app)]
    app = Starlette(routes=routes, lifespan=lifespan)
    app.add_middleware(BearerAuthMiddleware)
    return app


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(create_app(), host="127.0.0.1", port=5100)