magicwang1111-ComfyUI-Account-Manager/utils/internal_api.py

import base64
import hashlib
import hmac
import json
import os
import secrets
import time
from pathlib import Path

import aiohttp
from aiohttp import web


class InternalSigner:
    HEADER_TIMESTAMP = "X-Account-Manager-Timestamp"
    HEADER_SIGNATURE = "X-Account-Manager-Signature"

    def __init__(self, secret_file: str, max_age_seconds: int = 30):
        self.secret_file = Path(secret_file)
        self.max_age_seconds = max(5, int(max_age_seconds))
        self.secret = self._load_or_create_secret()

    def _load_or_create_secret(self) -> bytes:
        if self.secret_file.exists():
            return self.secret_file.read_text(encoding="utf-8").strip().encode("utf-8")
        self.secret_file.parent.mkdir(parents=True, exist_ok=True)
        value = secrets.token_hex(64)
        flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
        try:
            descriptor = os.open(self.secret_file, flags, 0o600)
            with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
                handle.write(value)
        except FileExistsError:
            return self.secret_file.read_text(encoding="utf-8").strip().encode("utf-8")
        return value.encode("utf-8")

    @staticmethod
    def _canonical(method: str, path: str, timestamp: str, body: bytes) -> bytes:
        digest = hashlib.sha256(body).hexdigest()
        return f"{method.upper()}\n{path}\n{timestamp}\n{digest}".encode("utf-8")

    def headers(self, method: str, path: str, body: bytes) -> dict[str, str]:
        timestamp = str(int(time.time()))
        signature = hmac.new(
            self.secret,
            self._canonical(method, path, timestamp, body),
            hashlib.sha256,
        ).hexdigest()
        return {
            self.HEADER_TIMESTAMP: timestamp,
            self.HEADER_SIGNATURE: signature,
        }

    def verify(
        self,
        method: str,
        path: str,
        body: bytes,
        timestamp: str,
        signature: str,
    ) -> bool:
        try:
            numeric_timestamp = int(timestamp)
        except (TypeError, ValueError):
            return False
        if abs(int(time.time()) - numeric_timestamp) > self.max_age_seconds:
            return False
        expected = hmac.new(
            self.secret,
            self._canonical(method, path, timestamp, body),
            hashlib.sha256,
        ).hexdigest()
        return hmac.compare_digest(expected, signature or "")

    async def verify_request(self, request: web.Request, body: bytes) -> bool:
        return self.verify(
            request.method,
            request.path,
            body,
            request.headers.get(self.HEADER_TIMESTAMP, ""),
            request.headers.get(self.HEADER_SIGNATURE, ""),
        )


class EventRelay:
    PATH = "/account-manager/internal/event"

    def __init__(self, server, scheduler, signer: InternalSigner, instance_port: int):
        self.server = server
        self.scheduler = scheduler
        self.signer = signer
        self.instance_port = int(instance_port)
        self._original_send_json = server.send_json
        self._original_send_bytes = server.send_bytes

    def install(self) -> None:
        self.server.send_json = self.send_json
        self.server.send_bytes = self.send_bytes

    async def _post(self, ingress_port: int, payload: dict) -> bool:
        body = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode(
            "utf-8"
        )
        headers = self.signer.headers("POST", self.PATH, body)
        headers["Content-Type"] = "application/json"
        url = f"http://127.0.0.1:{int(ingress_port)}{self.PATH}"
        try:
            timeout = aiohttp.ClientTimeout(total=10)
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.post(url, data=body, headers=headers) as response:
                    return 200 <= response.status < 300
        except (aiohttp.ClientError, TimeoutError, OSError):
            return False

    async def send_json(self, event, data, sid=None):
        if sid is None or sid in self.server.sockets:
            return await self._original_send_json(event, data, sid)
        ingress_port = self.scheduler.client_ingress(str(sid))
        if not ingress_port or ingress_port == self.instance_port:
            return await self._original_send_json(event, data, sid)
        relayed = await self._post(
            ingress_port,
            {
                "kind": "json",
                "event": event,
                "data": data,
                "sid": str(sid),
                "source_port": self.instance_port,
            },
        )
        if not relayed:
            return await self._original_send_json(event, data, sid)
        return None

    async def send_bytes(self, event, data, sid=None):
        if sid is None or sid in self.server.sockets:
            return await self._original_send_bytes(event, data, sid)
        ingress_port = self.scheduler.client_ingress(str(sid))
        if not ingress_port or ingress_port == self.instance_port:
            return await self._original_send_bytes(event, data, sid)
        encoded = self.server.encode_bytes(event, data)
        relayed = await self._post(
            ingress_port,
            {
                "kind": "bytes",
                "data": base64.b64encode(encoded).decode("ascii"),
                "sid": str(sid),
                "source_port": self.instance_port,
            },
        )
        if not relayed:
            return await self._original_send_bytes(event, data, sid)
        return None

    async def handle_event(self, request: web.Request) -> web.Response:
        body = await request.read()
        if not await self.signer.verify_request(request, body):
            return web.json_response({"error": "Invalid internal signature"}, status=403)
        try:
            payload = json.loads(body)
            sid = str(payload["sid"])
            if payload["kind"] == "json":
                await self._original_send_json(payload["event"], payload.get("data"), sid)
            elif payload["kind"] == "bytes":
                message = base64.b64decode(payload["data"], validate=True)
                socket = self.server.sockets.get(sid)
                if socket is not None:
                    await socket.send_bytes(message)
            else:
                raise ValueError("Unsupported event kind")
        except (KeyError, TypeError, ValueError, json.JSONDecodeError):
            return web.json_response({"error": "Invalid internal event"}, status=400)
        return web.json_response({"ok": True})