1import asyncio
  2import atexit
  3import logging
  4import ssl
  5import tempfile
  6
  7from grpclib import GRPCError
  8from grpclib.client import Channel
  9from grpclib.exceptions import StreamTerminatedError
 10
 11from .base import MarzNodeBase
 12from .database import MarzNodeDB
 13from .marznode_grpc import MarzServiceStub
 14from .marznode_pb2 import UserData, UsersData, Empty, User, Inbound, XrayLogsRequest, XrayConfig
 15from ..models.node import NodeStatus
 16
 17logger = logging.getLogger(__name__)
 18
 19
 20def string_to_temp_file(content: str):
 21    file = tempfile.NamedTemporaryFile(mode='w+t')
 22    file.write(content)
 23    file.flush()
 24    return file
 25
 26
 27class MarzNodeGRPCLIB(MarzNodeBase, MarzNodeDB):
 28    def __init__(self, node_id: int, address: str, port: int, ssl_key: str, ssl_cert: str, usage_coefficient: int = 1):
 29        self.id = node_id
 30        self._address = address
 31        self._port = port
 32
 33        self._key_file = string_to_temp_file(ssl_key)
 34        self._cert_file = string_to_temp_file(ssl_cert)
 35
 36        ctx = ssl.create_default_context()
 37        ctx.load_cert_chain(self._cert_file.name, self._key_file.name)
 38        ctx.check_hostname = False
 39        ctx.verify_mode = ssl.CERT_NONE
 40
 41        self._channel = Channel(self._address, self._port, ssl=ctx)
 42        self._stub = MarzServiceStub(self._channel)
 43        asyncio.create_task(self._monitor_channel())
 44        self._streaming_task = None
 45
 46        self._updates_queue = asyncio.Queue(1)
 47        self.synced = False
 48        self.usage_coefficient = usage_coefficient
 49        atexit.register(self._channel.close)
 50
 51    async def _monitor_channel(self):
 52        while state := self._channel._state:
 53            logger.debug("node %i channel state: %s", self.id, state.value)
 54            try:
 55                await asyncio.wait_for(self._channel.__connect__(), timeout=2)
 56            except Exception:
 57                logger.debug("timeout for node, id: %i", self.id)
 58                self.set_status(NodeStatus.unhealthy, "timeout")
 59                self.synced = False
 60                if self._streaming_task:
 61                    self._streaming_task.cancel()
 62            else:
 63                if not self.synced:
 64                    await self._sync()
 65                    self._streaming_task = asyncio.create_task(self._stream_user_updates())
 66                    self.set_status(NodeStatus.healthy)
 67                    logger.info("Connected to node %i", self.id)
 68            await asyncio.sleep(10)
 69
 70    async def _stream_user_updates(self):
 71        try:
 72            async with self._stub.SyncUsers.open() as stream:
 73                logger.info("opened the stream")
 74                while True:
 75                    user_update = await self._updates_queue.get()
 76                    logger.info("got something from queue")
 77                    user = user_update["user"]
 78                    await stream.send_message(
 79                        UserData(user=User(id=user.id, username=user.username, key=user.key),
 80                                 inbounds=[Inbound(tag=t) for t in user_update["inbounds"]]))
 81        except (OSError, ConnectionError, GRPCError, StreamTerminatedError):
 82            logger.info("node %i detached", self.id)
 83            self.synced = False
 84
 85    async def update_user(self, user, inbounds: set[str] | None = None):
 86        if inbounds is None:
 87            inbounds = set()
 88
 89        if self.synced:
 90            await self._updates_queue.put({"user": user, "inbounds": inbounds})
 91
 92    async def _repopulate_users(self, users_data: list[dict]) -> None:
 93        updates = [UserData(user=User(id=u["id"], username=u["username"], key=u["key"]),
 94                            inbounds=[Inbound(tag=t) for t in u["inbounds"]]) for u in users_data]
 95        await self._stub.RepopulateUsers(UsersData(users_data=updates))
 96
 97    async def fetch_users_stats(self):
 98        response = await self._stub.FetchUsersStats(Empty())
 99        return response.users_stats
100
101    async def _fetch_inbounds(self) -> list:
102        response = await self._stub.FetchInbounds(Empty())
103        return response.inbounds
104
105    async def _sync(self):
106        inbounds = await self._fetch_inbounds()
107        self.store_inbounds(inbounds)
108        users = self.list_users()
109        await self._repopulate_users(users)
110        self.synced = True
111
112    async def get_logs(self, include_buffer=True):
113        async with self._stub.StreamXrayLogs.open() as stm:
114            await stm.send_message(XrayLogsRequest(include_buffer=include_buffer))
115            while True:
116                response = await stm.recv_message()
117                yield response.line
118
119    async def restart_xray(self, config: str):
120        await self._stub.RestartXray(XrayConfig(configuration=config))
121        await self._sync()
122
123    async def get_xray_config(self):
124        response = await self._stub.FetchXrayConfig(Empty())
125        return response.configuration