diff --git a/src/server.py b/src/server.py new file mode 100644 index 0000000..962c3a0 --- /dev/null +++ b/src/server.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 +__version__ = "0.1.0" +import asyncio +import logging +import os +import smtplib +import ssl +import time +from contextlib import asynccontextmanager +from datetime import date +from email import policy +from email.mime.multipart import MIMEMultipart +from email.mime.text import MIMEText +from email.parser import BytesParser +from typing import Optional +import hmac +import httpx +import uvicorn +import yaml +from imapclient import IMAPClient +from fastmcp import FastMCP, Context +from fastmcp.server.auth import TokenVerifier, AccessToken +from starlette.middleware import Middleware +from starlette.responses import JSONResponse, Response +from starlette.types import ASGIApp, Receive, Scope, Send + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("poke-mail") +logging.getLogger("httpx").setLevel(logging.WARNING) + +class ApiKeyAuth(TokenVerifier): + """Validates incoming requests against a static API key (MCP_API_KEY).""" + def __init__(self, api_key: str): + super().__init__() + self._api_key = api_key + async def verify_token(self, token: str) -> AccessToken | None: + if hmac.compare_digest(token, self._api_key): + return AccessToken(token=token, client_id="owner", scopes=["all"]) + return None + +class DropNonMCPRoutes: + """Return empty 404 for any path outside /mcp — reveals nothing to scanners.""" + def __init__(self, app: ASGIApp): + self.app = app + async def __call__(self, scope: Scope, receive: Receive, send: Send): + if scope["type"] == "http" and not scope["path"].startswith("/mcp"): + response = Response(status_code=404) + await response(scope, receive, send) + return + await self.app(scope, receive, send) + +class RateLimitMiddleware: + MAX_TRACKED_IPS = 1024 + def __init__(self, app: ASGIApp): + self.app = app + self.get_rpm = int(os.environ.get("RATE_LIMIT_GET_RPM", "30")) + self.post_rpm = int(os.environ.get("RATE_LIMIT_POST_RPM", "120")) + self.window = 60 + self._hits: dict[str, list[float]] = {} + self._last_cleanup = time.monotonic() + def _client_ip(self, scope: Scope) -> str: + for header_name, header_val in scope.get("headers", []): + if header_name == b"x-forwarded-for": + parts = header_val.decode().split(",") + return parts[-1].strip() + client = scope.get("client") + return client[0] if client else "unknown" + def _cleanup_stale(self, now: float) -> None: + if now - self._last_cleanup > 300: + by_recency = sorted(self._hits, key=lambda ip: self._hits[ip][-1]) + for ip in by_recency[: len(self._hits) - self.MAX_TRACKED_IPS]: + del self._hits[ip] + self._last_cleanup = now + def _is_limited(self, bucket: str, rpm: int) -> tuple[bool, int]: + now = time.monotonic() + self._cleanup_stale(now) + timestamps = self._hits.get(bucket, []) + cutoff = now - self.window + timestamps = [t for t in timestamps if t > cutoff] + self._hits[bucket] = timestamps + if len(timestamps) >= rpm: + oldest = timestamps[0] + retry_after = int(oldest + self.window - now) + 1 + return True, max(retry_after, 1) + timestamps.append(now) + return False, 0 + async def __call__(self, scope: Scope, receive: Receive, send: Send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + ip = self._client_ip(scope) + method = scope.get("method", "GET") + bucket, rpm = (f"{ip}:post", self.post_rpm) if method == "POST" else (f"{ip}:get", self.get_rpm) + limited, retry_after = self._is_limited(bucket, rpm) + if limited: + response = JSONResponse({"error": "rate_limited", "retry_after": retry_after}, status_code=429, headers={"Retry-After": str(retry_after)}) + await response(scope, receive, send) + return + await self.app(scope, receive, send) + +def load_config() -> dict: + path = os.environ.get("CONFIG_PATH", "config.yml") + try: + with open(path) as f: + return yaml.safe_load(f) or {} + except FileNotFoundError: + logger.warning("Config file %s not found, using env vars", path) + return {} + +def parse_accounts(config: dict) -> list[dict]: + accounts = config.get("accounts", []) + if not accounts: + imap_host = os.environ.get("IMAP_HOST") + if not imap_host: + raise RuntimeError("No accounts configured. Set POKE_MAIL_ACCOUNTS env var or create config.yml") + accounts = [{ + "id": "default", + "imap_host": imap_host, + "imap_port": int(os.environ.get("IMAP_PORT", "993")), + "imap_username": os.environ["IMAP_USERNAME"], + "imap_password": os.environ["IMAP_PASSWORD"], + "smtp_host": os.environ.get("SMTP_HOST", imap_host), + "smtp_port": int(os.environ.get("SMTP_PORT", "587")), + "smtp_username": os.environ.get("SMTP_USERNAME", os.environ["IMAP_USERNAME"]), + "smtp_password": os.environ.get("SMTP_PASSWORD", os.environ["IMAP_PASSWORD"]), + "from_address": os.environ.get("FROM_ADDRESS", os.environ.get("SMTP_USERNAME", os.environ["IMAP_USERNAME"])), + "watch_folders": ["INBOX"], + "mark_as_read": os.environ.get("MARK_AS_READ", "false").lower() == "true", + }] + global_allow_send = config.get("allow_send", True) + global_mark_as_read = config.get("mark_as_read", False) + required = ("imap_host", "imap_username", "imap_password") + for i, acc in enumerate(accounts): + acc.setdefault("id", f"account-{i}") + acc.setdefault("imap_port", 993) + acc.setdefault("watch_folders", ["INBOX"]) + acc.setdefault("smtp_host", acc.get("imap_host")) + acc.setdefault("smtp_port", 587) + acc.setdefault("smtp_username", acc.get("imap_username")) + acc.setdefault("smtp_password", acc.get("imap_password")) + acc.setdefault("from_address", acc.get("smtp_username")) + acc.setdefault("allow_send", global_allow_send) + acc.setdefault("mark_as_read", global_mark_as_read) + for field in required: + if field not in acc: + raise RuntimeError(f"Account '{acc['id']}' missing required field: {field}") + return accounts + +def resolve_account(accounts: list[dict], account_id: str) -> dict: + if not account_id or not account_id.strip(): + raise ValueError(f"account_id is required. Available accounts: {[a['id'] for a in accounts]}.") + needle = account_id.strip().lower() + for acc in accounts: + if acc["id"].lower() == needle: + return acc + matches = [acc for acc in accounts if needle in {(acc.get("from_address") or "").lower(), (acc.get("imap_username") or "").lower(), (acc.get("smtp_username") or "").lower()}] + if len(matches) == 1: + return matches[0] + if len(matches) > 1: + raise ValueError(f"account_id '{account_id}' matches multiple accounts: {[a['id'] for a in matches]}. Pass the unique 'id' instead.") + raise ValueError(f"Unknown account_id '{account_id}'. Available: {[a['id'] for a in accounts]}.") + +_IMAP_CLIENT_ID = {"name": "poke-mail", "version": "1.0.0", "vendor": "Poke Interactions"} + +def get_imap_client(account: dict) -> IMAPClient: + port = account["imap_port"] + use_ssl = port == 993 + client = IMAPClient(account["imap_host"], port=port, ssl=use_ssl) + if not use_ssl: + client.starttls() + client.login(account["imap_username"], account["imap_password"]) + try: + if client.has_capability("ID"): + client.id_(_IMAP_CLIENT_ID) + except Exception as e: + logger.debug("IMAP ID command failed for %s: %s", account.get("id"), e) + return client + +def parse_email_message(raw: bytes) -> dict: + msg = BytesParser(policy=policy.default).parsebytes(raw) + body_text = "" + body_html = "" + attachments = [] + if msg.is_multipart(): + for part in msg.walk(): + ct = part.get_content_type() + cd = str(part.get("Content-Disposition", "")) + if "attachment" in cd: + try: + content = part.get_content() + size = len(content) if hasattr(content, "__len__") else 0 + except Exception: + size = 0 + attachments.append({"filename": part.get_filename() or "unnamed", "content_type": ct, "size": size}) + elif ct == "text/plain" and not body_text: + try: + body_text = part.get_content() + except Exception: + body_text = part.get_payload(decode=True).decode(errors="replace") + elif ct == "text/html" and not body_html: + try: + body_html = part.get_content() + except Exception: + body_html = part.get_payload(decode=True).decode(errors="replace") + else: + ct = msg.get_content_type() + try: + content = msg.get_content() + except Exception: + payload = msg.get_payload(decode=True) + content = payload.decode(errors="replace") if payload else "" + if ct == "text/html": + body_html = content + else: + body_text = content + def parse_addresses(header): + if not header: + return [] + return [addr.strip() for addr in str(header).split(",") if addr.strip()] + return { + "from": str(msg["from"] or ""), + "to": parse_addresses(msg["to"]), + "cc": parse_addresses(msg["cc"]), + "subject": str(msg["subject"] or ""), + "date": str(msg["date"] or ""), + "body_text": body_text, + "body_html": body_html, + "headers": {k: str(v) for k, v in msg.items()}, + "attachments": attachments, + } + +def build_search_criteria(from_addr: Optional[str] = None, to_addr: Optional[str] = None, subject: Optional[str] = None, since: Optional[str] = None, before: Optional[str] = None) -> list: + criteria = [] + if from_addr: + criteria.extend(["FROM", from_addr]) + if to_addr: + criteria.extend(["TO", to_addr]) + if subject: + criteria.extend(["SUBJECT", subject]) + if since: + criteria.extend(["SINCE", date.fromisoformat(since)]) + if before: + criteria.extend(["BEFORE", date.fromisoformat(before)]) + if not criteria: + criteria = ["ALL"] + return criteria + +def detect_archive_folder(client: IMAPClient) -> str: + folders = client.list_folders() + for flags, _delim, name in folders: + if b"\\Archive" in flags: + return name + if name in ("[Gmail]/All Mail", "Archive"): + return name + return "Archive" + +def detect_drafts_folder(client: IMAPClient) -> str: + folders = client.list_folders() + for flags, _delim, name in folders: + if b"\\Drafts" in flags: + return name + for name in ("Drafts", "[Gmail]/Drafts", "INBOX.Drafts"): + if client.folder_exists(name): + return name + return "Drafts" + +async def forward_to_poke(email_data: dict, account: dict, webhook_url: str, api_key: str) -> bool: + payload = { + "account_id": account["id"], + "from_address": account["from_address"], + "from": email_data["from"], + "to": email_data["to"], + "subject": email_data["subject"], + "date": email_data["date"], + "body_text": email_data["body_text"], + "body_html": email_data["body_html"], + "headers": email_data.get("headers", {}), + "attachments": email_data.get("attachments", []), + } + async with httpx.AsyncClient() as client: + headers = {"Authorization": f"Bearer {api_key}"} + try: + resp = await client.post(webhook_url, json=payload, headers=headers, timeout=10) + resp.raise_for_status() + return True + except Exception as e: + logger.error("Failed to forward email to Poke: %s", e) + return False + +async def watch_account(account: dict, config: dict, webhook_url: str, api_key: str): + client = get_imap_client(account) + folder = config.get("watch_folder", "INBOX") + logger.info("Watching folder %s for account %s", folder, account["id"]) + client.select_folder(folder) + while True: + try: + client.idle() + await asyncio.sleep(0.1) + messages = client.search(["UNSEEN"]) + for uid in messages: + raw = client.fetch([uid], ["RFC822"])[uid][b"RFC822"] + email_data = parse_email_message(raw) + await forward_to_poke(email_data, account, webhook_url, api_key) + if account.get("mark_as_read", False): + client.add_flags(uid, [b"\Seen"]) + except Exception as e: + logger.error("Error watching account %s: %s", account["id"], e) + await asyncio.sleep(5) + finally: + client.idle_done() + +@asynccontextmanager +async def lifespan(app: FastMCP): + config = load_config() + webhook_url = config.get("webhook_url") + api_key = os.environ.get("MCP_API_KEY") + accounts = parse_accounts(config) + tasks = [] + for account in accounts: + tasks.append(watch_account(account, config, webhook_url, api_key)) + runner = asyncio.create_task(asyncio.gather(*tasks)) + yield + runner.cancel() + await runner + +def create_app(): + api_key = os.environ.get("MCP_API_KEY", "") + middleware = [ + Middleware(RateLimitMiddleware), + Middleware(DropNonMCPRoutes), + Middleware(ApiKeyAuth, api_key=api_key), + ] + return FastMCP(lifespan=lifespan, middleware=middleware) + +app = create_app() + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "3000")))