rpc.py - dev - kotkov/eco

rpc.py
125 строк4 KB
Arseny Kotkov
fix: agent start up hotfix
27 май 2025, 06:47
27 май 2025, 06:47de888a7
""" JSON-RPC 2.0 via HTTP """
import functools
import json
from http.server import HTTPServer, BaseHTTPRequestHandler
from threading import Thread

import requests
from requests import post
from log import ecolog


class Message:
    def __init__(self, method: str, params: list):
        self._method = method
        self._params = params

    def encode(self):
        return json.dumps({
            'method': self._method,
            'params': self._params
        }).encode()


def rpc_method(desc):
    # TODO: return number of parameters, description
    def wrapper(func):
        func._is_rpc = True
        return func
    return wrapper


class RPCMeta(type):
    """ Adds @public methods to self.public_methods before class init.  """

    def __new__(mcs, name, bases, attrs):
        cls = super().__new__(mcs, name, bases, attrs)
        cls.public_methods = []
        for name, attr in attrs.items():
            if hasattr(attr, '_is_rpc'):
                cls.public_methods.append(name)
        return cls


class RPC:
    def __init__(self, cls):
        self._ip: str = ""
        self._port: int = 0
        self._cls = cls
        self._thread = None
        self._httpd = None

    @property
    def ip(self):
        return self._ip

    @ip.setter
    def ip(self, x):
        self._ip = x

    @property
    def port(self):
        return self._port

    @port.setter
    def port(self, x):
        self._port = x

    def serve(self, socket: (str, int), blocking=False) -> None:
        # sys.stdout = open('log.txt', 'a')
        self._ip, self._port = socket
        ecolog(f"RPC serving at {self.ip}:{self.port}")
        RES_SUFFIX = '_handler'
        _cls = self._cls

        class RequestHandler(BaseHTTPRequestHandler):
            cls = _cls

            def do_POST(self):
                data = self._get_data()
                self._send_headers()

                # Ищем ссылку на метод, который будем вызывать
                method = getattr(self.cls, data['method'])
                # Проверяем, доступен ли метод для RPC
                if self.is_public(method):
                    # Выполняем RPC, получаем ответ
                    response = method(*data['params'])
                    """
                    # Ищем метод-обработчик
                    response_method = getattr(
                        self.cls,
                        f'{data["method"]}{RES_SUFFIX}'  # e.g. add_handler
                    )
                    # Снова выполняем RPC, вызывая
                    response_method(response)
                    """
                    self.wfile.write(response.encode())
                else:
                    ecolog(f"Method {method.__name__} is private.")

            def _get_data(self):
                content_length = int(self.headers['Content-Length'])
                post_data = self.rfile.read(content_length)
                return json.loads(post_data)

            def _send_headers(self):
                self.send_response(200)
                self.send_header('Content-type', 'application/json')
                self.end_headers()

            def is_public(self, method) -> bool:
                return method.__name__ in self.cls.public_methods

        self._httpd = HTTPServer(socket, RequestHandler)
        self._thread = Thread(target=self._httpd.serve_forever)
        self._thread.daemon = True
        self._thread.start()
        if blocking:
            self._thread.join()

    def call(self, ip: str, port: int, message: Message):
        try:
            return post(f'http://{ip}:{port}', data=message.encode()).text
        except requests.exceptions.ConnectionError as e:
            return json.dumps({'error': str(e)}).encode()