communex.module.server

  1"""
  2Server for Commune modules.
  3"""
  4
  5import inspect
  6import random
  7from functools import partial
  8from typing import Any
  9
 10import fastapi
 11from fastapi import APIRouter
 12from pydantic import BaseModel
 13from substrateinterface import Keypair
 14
 15from communex.key import check_ss58_address
 16from communex.module import _signer as signer
 17from communex.module._rate_limiters.limiters import (
 18    IpLimiterParams,
 19    StakeLimiterParams,
 20)
 21from communex.module.module import EndpointDefinition, Module, endpoint
 22from communex.module.routers.module_routers import (
 23    InputHandlerVerifier,
 24    IpLimiterVerifier,
 25    ListVerifier,
 26    StakeLimiterVerifier,
 27    build_route_class,
 28)
 29from communex.types import Ss58Address
 30from communex.util.memo import TTLDict
 31
 32
 33class ModuleServer:
 34    def __init__(
 35        self,
 36        module: Module,
 37        key: Keypair,
 38        max_request_staleness: int = 120,
 39        whitelist: list[Ss58Address] | None = None,
 40        blacklist: list[Ss58Address] | None = None,
 41        subnets_whitelist: list[int] | None = None,
 42        lower_ttl: int = 600,
 43        upper_ttl: int = 700,
 44        limiter: StakeLimiterParams | IpLimiterParams = StakeLimiterParams(),
 45        ip_blacklist: list[str] | None = None,
 46        use_testnet: bool = False,
 47    ) -> None:
 48        self._module = module
 49        self._app = fastapi.FastAPI()
 50        self._subnets_whitelist = subnets_whitelist
 51        self.key = key
 52        self.max_request_staleness = max_request_staleness
 53        ttl = random.randint(lower_ttl, upper_ttl)
 54        self._blockchain_cache = TTLDict[str, list[Ss58Address]](ttl)
 55        self._ip_blacklist = ip_blacklist
 56
 57        # to keep reference to add_to_blacklist and add_to_whitelist
 58        whitelist_: list[Ss58Address] = [] if whitelist is None else whitelist
 59        blacklist_: list[Ss58Address] = [] if blacklist is None else blacklist
 60
 61        self._whitelist = whitelist_
 62        self._blacklist = blacklist_
 63
 64        self._build_routers(use_testnet, limiter)
 65
 66    def _build_routers(
 67        self, use_testnet: bool, limiter: StakeLimiterParams | IpLimiterParams
 68    ):
 69        input_handler = InputHandlerVerifier(
 70            self._subnets_whitelist,
 71            check_ss58_address(self.key.ss58_address),
 72            self.max_request_staleness,
 73            self._blockchain_cache,
 74            self.key,
 75            use_testnet,
 76        )
 77        check_lists = ListVerifier(
 78            self._blacklist, self._whitelist, self._ip_blacklist
 79        )
 80        if isinstance(limiter, StakeLimiterParams):
 81            limiter_verifier = StakeLimiterVerifier(
 82                self._subnets_whitelist, limiter
 83            )
 84        else:
 85            limiter_verifier = IpLimiterVerifier(limiter)
 86
 87        # order of verifiers is extremely important
 88        verifiers = [check_lists, input_handler, limiter_verifier]
 89        route_class = build_route_class(verifiers)
 90        self._router = APIRouter(route_class=route_class)
 91        self.register_endpoints(self._router)
 92        self._app.include_router(self._router)
 93
 94    def get_fastapi_app(self):
 95        return self._app
 96
 97    def register_endpoints(self, router: APIRouter):
 98        endpoints = self._module.get_endpoints()
 99        for name, endpoint_def in endpoints.items():
100
101            class Body(BaseModel):
102                params: endpoint_def.params_model  # type: ignore
103
104            async def async_handler(
105                end_def: EndpointDefinition[Any, ...], body: Body
106            ):
107                return await end_def.fn(
108                    self._module,
109                    **body.params.model_dump(),  # type: ignore
110                )
111
112            def handler(end_def: EndpointDefinition[Any, ...], body: Body):
113                return end_def.fn(self._module, **body.params.model_dump())  # type: ignore
114
115            if inspect.iscoroutinefunction(endpoint_def.fn):
116                defined_handler = partial(async_handler, endpoint_def)
117            else:
118                defined_handler = partial(handler, endpoint_def)
119            router.post(f"/method/{name}")(defined_handler)
120
121    def add_to_blacklist(self, ss58_address: Ss58Address):
122        self._blacklist.append(ss58_address)
123
124    def add_to_whitelist(self, ss58_address: Ss58Address):
125        self._whitelist.append(ss58_address)
126
127
128def main():
129    class Amod(Module):
130        @endpoint
131        def do_the_thing(self, awesomness: int = 42):
132            if awesomness > 60:
133                msg = f"You're super awesome: {awesomness} awesomness"
134            else:
135                msg = f"You're not that awesome: {awesomness} awesomness"
136            return {"msg": msg}
137
138    a_mod = Amod()
139    keypair = Keypair.create_from_mnemonic(signer.TESTING_MNEMONIC)
140    server = ModuleServer(
141        a_mod,
142        keypair,
143        subnets_whitelist=[0],
144        blacklist=None,
145    )
146    app = server.get_fastapi_app()
147
148    import uvicorn
149
150    uvicorn.run(app, host="127.0.0.1", port=8000)  # type: ignore
151
152
153if __name__ == "__main__":
154    main()