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