SenSu/services/proxy_service.py

#!/usr/bin/env python3
"""SenSu ProxyService — 反向代理，将任意 URL 映射到框架路径"""
import asyncio, logging, aiohttp
from aiohttp import web, ClientSession, WSMsgType
from typing import Dict, Optional
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

@dataclass
class ProxyTarget:
    path: str
    target_url: str
    description: str = ""
    is_external: bool = False
    status: str = "active"
    strip_prefix: bool = True

class ProxyService:
    def __init__(self, service_manager=None):
        self.sm = service_manager
        self.proxies: Dict[str, ProxyTarget] = {}
        self._session: Optional[ClientSession] = None

    async def start(self):
        self._session = ClientSession()
        logger.info("ProxyService 已就绪")

    async def register_proxy(self, path: str, target_url: str,
                             description: str = "", is_external: bool = False,
                             strip_prefix: bool = True) -> bool:
        path = "/" + path.strip("/")
        if path in self.proxies:
            logger.warning(f"代理路径已存在: {path}")
            return False
        if not target_url.endswith("/"):
            target_url += "/"
        self.proxies[path] = ProxyTarget(path=path, target_url=target_url,
                                         description=description, is_external=is_external,
                                         strip_prefix=strip_prefix)
        logger.info(f"代理已注册: {path} → {target_url}")
        return True

    def unregister_proxy(self, path: str):
        path = "/" + path.strip("/")
        if path in self.proxies:
            del self.proxies[path]
            logger.info(f"代理已注销: {path}")

    def list_proxies(self):
        return [{"path": p.path, "target": p.target_url, "status": p.status,
                 "description": p.description, "external": p.is_external}
                for p in self.proxies.values()]

    def setup_routes(self, app: web.Application):
        """注册代理路由到 aiohttp app"""
        async def proxy_handler(request):
            path = request.path
            # Find matching proxy (longest prefix match)
            proxy = None
            for p in sorted(self.proxies.keys(), key=len, reverse=True):
                if path.startswith(p) or path == p:
                    proxy = self.proxies[p]
                    break
            if not proxy:
                # Check bare path
                lookup = "/" + path.strip("/")
                proxy = self.proxies.get(lookup)

            if not proxy:
                return web.json_response({"error": "no proxy for path"}, status=404)

            # Build target URL
            remaining = path[len(proxy.path):] if proxy.strip_prefix else path
            target = proxy.target_url.rstrip("/") + "/" + remaining.lstrip("/")

            try:
                # Forward request
                headers = {k: v for k, v in request.headers.items()
                          if k.lower() not in ("host", "content-length")}
                headers["X-Forwarded-For"] = request.remote
                headers["X-Proxy-By"] = "SenSu"

                async with self._session.request(
                    request.method, target, headers=headers,
                    data=await request.read(), timeout=30
                ) as resp:
                    body = await resp.read()
                    proxy_resp = web.Response(body=body, status=resp.status)
                    for k, v in resp.headers.items():
                        if k.lower() not in ("transfer-encoding", "content-encoding"):
                            proxy_resp.headers[k] = v
                    return proxy_resp
            except asyncio.TimeoutError:
                return web.json_response({"error": "proxy timeout"}, status=504)
            except Exception as e:
                logger.error(f"代理错误 {path}: {e}")
                return web.json_response({"error": str(e)}, status=502)

        # WebSocket proxy
        async def ws_proxy_handler(request):
            path = request.path
            proxy = None
            for p in sorted(self.proxies.keys(), key=len, reverse=True):
                if path.startswith(p):
                    proxy = self.proxies[p]
                    break
            if not proxy:
                return web.json_response({"error": "no ws proxy"}, status=404)

            target = proxy.target_url.rstrip("/") + "/" + path[len(proxy.path):].lstrip("/")
            if target.startswith("http"):
                target = target.replace("http://", "ws://").replace("https://", "wss://")

            ws_client = web.WebSocketResponse()
            await ws_client.prepare(request)
            try:
                async with self._session.ws_connect(target) as ws_target:
                    async def forward(src, dst):
                        async for msg in src:
                            if msg.type == WSMsgType.TEXT:
                                await dst.send_str(msg.data)
                            elif msg.type == WSMsgType.BINARY:
                                await dst.send_bytes(msg.data)
                            elif msg.type in (WSMsgType.CLOSE, WSMsgType.ERROR):
                                break

                    await asyncio.gather(
                        forward(ws_client, ws_target),
                        forward(ws_target, ws_client),
                    )
            except Exception as e:
                logger.debug(f"WS proxy error: {e}")
            return ws_client

        app.router.add_route("*", "/proxy/{tail:.*}", proxy_handler)
        # Register individual proxy routes
        for path in self.proxies:
            app.router.add_route("*", f"{path}/{{tail:.*}}", proxy_handler)

        # WebSocket proxy
        app.router.add_route("GET", "/wsproxy/{tail:.*}", ws_proxy_handler)
        logger.info(f"代理路由已注册 ({len(self.proxies)} targets)")

    async def check_health(self, path: str) -> dict:
        proxy = self.proxies.get("/" + path.strip("/"))
        if not proxy:
            return {"ok": False, "error": "not found"}
        try:
            async with self._session.get(proxy.target_url, timeout=5) as resp:
                return {"ok": True, "status": resp.status, "target": proxy.target_url}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    async def shutdown(self):
        if self._session:
            await self._session.close()
        logger.info("ProxyService 已关闭")