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 BackendConfig,
21 BackendLogsRequest,
22 Backend,
23 RestartBackendRequest,
24 BackendStats,
25)
26from ..models.node import NodeStatus
27
28logger = logging.getLogger(__name__)
29
30
31def string_to_temp_file(content: str):
32 file = tempfile.NamedTemporaryFile(mode="w+t")
33 file.write(content)
34 file.flush()
35 return file
36
37
38class MarzNodeGRPCLIB(MarzNodeBase, MarzNodeDB):
39 def __init__(
40 self,
41 node_id: int,
42 address: str,
43 port: int,
44 ssl_key: str,
45 ssl_cert: str,
46 usage_coefficient: int = 1,
47 ):
48 self.id = node_id
49 self._address = address
50 self._port = port
51
52 self._key_file = string_to_temp_file(ssl_key)
53 self._cert_file = string_to_temp_file(ssl_cert)
54
55 ctx = ssl.create_default_context()
56 ctx.load_cert_chain(self._cert_file.name, self._key_file.name)
57 ctx.check_hostname = False
58 ctx.verify_mode = ssl.CERT_NONE
59
60 self._channel = Channel(self._address, self._port, ssl=ctx)
61 self._stub = MarzServiceStub(self._channel)
62 asyncio.create_task(self._monitor_channel())
63 self._streaming_task = None
64
65 self._updates_queue = asyncio.Queue(1)
66 self.synced = False
67 self.usage_coefficient = usage_coefficient
68 atexit.register(self._channel.close)
69
70 async def _monitor_channel(self):
71 while state := self._channel._state:
72 logger.debug("node %i channel state: %s", self.id, state.value)
73 try:
74 await asyncio.wait_for(self._channel.__connect__(), timeout=2)
75 except Exception:
76 logger.debug("timeout for node, id: %i", self.id)
77 self.set_status(NodeStatus.unhealthy, "timeout")
78 self.synced = False
79 if self._streaming_task:
80 self._streaming_task.cancel()
81 else:
82 if not self.synced:
83 try:
84 await self._sync()
85 except:
86 pass
87 else:
88 self._streaming_task = asyncio.create_task(
89 self._stream_user_updates()
90 )
91 self.set_status(NodeStatus.healthy)
92 logger.info("Connected to node %i", self.id)
93 await asyncio.sleep(10)
94
95 async def _stream_user_updates(self):
96 try:
97 async with self._stub.SyncUsers.open() as stream:
98 logger.debug("opened the stream")
99 while True:
100 user_update = await self._updates_queue.get()
101 logger.debug("got something from queue")
102 user = user_update["user"]
103 await stream.send_message(
104 UserData(
105 user=User(
106 id=user.id,
107 username=user.username,
108 key=user.key,
109 ),
110 inbounds=[
111 Inbound(tag=t) for t in user_update["inbounds"]
112 ],
113 )
114 )
115 except (OSError, ConnectionError, GRPCError, StreamTerminatedError):
116 logger.info("node %i detached", self.id)
117 self.synced = False
118
119 async def update_user(self, user, inbounds: set[str] | None = None):
120 if inbounds is None:
121 inbounds = set()
122
123 await self._updates_queue.put({"user": user, "inbounds": inbounds})
124
125 async def _repopulate_users(self, users_data: list[dict]) -> None:
126 updates = [
127 UserData(
128 user=User(id=u["id"], username=u["username"], key=u["key"]),
129 inbounds=[Inbound(tag=t) for t in u["inbounds"]],
130 )
131 for u in users_data
132 ]
133 await self._stub.RepopulateUsers(UsersData(users_data=updates))
134
135 async def fetch_users_stats(self):
136 response = await self._stub.FetchUsersStats(Empty())
137 return response.users_stats
138
139 async def _fetch_backends(self) -> list:
140 response = await self._stub.FetchBackends(Empty())
141 return response.backends
142
143 async def _sync(self):
144 backends = await self._fetch_backends()
145 self.store_backends(backends)
146 users = self.list_users()
147 await self._repopulate_users(users)
148 self.synced = True
149
150 async def get_logs(self, name: str = "xray", include_buffer=True):
151 async with self._stub.StreamBackendLogs.open() as stm:
152 await stm.send_message(
153 BackendLogsRequest(
154 backend_name=name, include_buffer=include_buffer
155 )
156 )
157 while True:
158 response = await stm.recv_message()
159 yield response.line
160
161 async def restart_backend(
162 self, name: str, config: str, config_format: int
163 ):
164 try:
165 await self._stub.RestartBackend(
166 RestartBackendRequest(
167 backend_name=name,
168 config=BackendConfig(
169 configuration=config, config_format=config_format
170 ),
171 )
172 )
173 await self._sync()
174 except:
175 self.synced = False
176 self.set_status(NodeStatus.unhealthy)
177 raise
178 else:
179 self.set_status(NodeStatus.healthy)
180
181 async def get_backend_config(self, name: str):
182 response: BackendConfig = await self._stub.FetchBackendConfig(
183 Backend(name=name)
184 )
185 return response.configuration, response.config_format
186
187 async def get_backend_stats(self, name: str):
188 response: BackendStats = await self._stub.GetBackendStats(
189 Backend(name=name)
190 )
191 return response