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