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                    try:
 81                        await self._sync()
 82                    except:
 83                        pass
 84                    else:
 85                        self._streaming_task = asyncio.create_task(
 86                            self._stream_user_updates()
 87                        )
 88                        self.set_status(NodeStatus.healthy)
 89                        logger.info("Connected to node %i", self.id)
 90            await asyncio.sleep(10)
 91
 92    async def _stream_user_updates(self):
 93        try:
 94            async with self._stub.SyncUsers.open() as stream:
 95                logger.debug("opened the stream")
 96                while True:
 97                    user_update = await self._updates_queue.get()
 98                    logger.debug("got something from queue")
 99                    user = user_update["user"]
100                    await stream.send_message(
101                        UserData(
102                            user=User(
103                                id=user.id,
104                                username=user.username,
105                                key=user.key,
106                            ),
107                            inbounds=[
108                                Inbound(tag=t) for t in user_update["inbounds"]
109                            ],
110                        )
111                    )
112        except (OSError, ConnectionError, GRPCError, StreamTerminatedError):
113            logger.info("node %i detached", self.id)
114            self.synced = False
115
116    async def update_user(self, user, inbounds: set[str] | None = None):
117        if inbounds is None:
118            inbounds = set()
119
120        await self._updates_queue.put({"user": user, "inbounds": inbounds})
121
122    async def _repopulate_users(self, users_data: list[dict]) -> None:
123        updates = [
124            UserData(
125                user=User(id=u["id"], username=u["username"], key=u["key"]),
126                inbounds=[Inbound(tag=t) for t in u["inbounds"]],
127            )
128            for u in users_data
129        ]
130        await self._stub.RepopulateUsers(UsersData(users_data=updates))
131
132    async def fetch_users_stats(self):
133        response = await self._stub.FetchUsersStats(Empty())
134        return response.users_stats
135
136    async def _fetch_inbounds(self) -> list:
137        response = await self._stub.FetchInbounds(Empty())
138        return response.inbounds
139
140    async def _sync(self):
141        inbounds = await self._fetch_inbounds()
142        self.store_inbounds(inbounds)
143        users = self.list_users()
144        await self._repopulate_users(users)
145        self.synced = True
146
147    async def get_logs(self, include_buffer=True):
148        async with self._stub.StreamXrayLogs.open() as stm:
149            await stm.send_message(
150                XrayLogsRequest(include_buffer=include_buffer)
151            )
152            while True:
153                response = await stm.recv_message()
154                yield response.line
155
156    async def restart_xray(self, config: str):
157        try:
158            await self._stub.RestartXray(XrayConfig(configuration=config))
159            await self._sync()
160        except:
161            self.synced = False
162            self.set_status(NodeStatus.unhealthy)
163            raise
164        else:
165            self.set_status(NodeStatus.healthy)
166
167    async def get_xray_config(self):
168        response = await self._stub.FetchXrayConfig(Empty())
169        return response.configuration