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