Add src/server.py
This commit is contained in:
+338
@@ -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")))
|
||||
Reference in New Issue
Block a user