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 self._monitor_task = 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 stop(self):
71 self._channel.close()
72 self._monitor_task.cancel()
73
74 async def _monitor_channel(self):
75 while state := self._channel._state:
76 logger.debug("node %i channel state: %s", self.id, state.value)
77 try:
78 await asyncio.wait_for(self._channel.__connect__(), timeout=2)
79 except Exception:
80 logger.debug("timeout for node, id: %i", self.id)
81 self.set_status(NodeStatus.unhealthy, "timeout")
82 self.synced = False
83 if self._streaming_task:
84 self._streaming_task.cancel()
85 else:
86 if not self.synced:
87 try:
88 await self._sync()
89 except:
90 pass
91 else:
92 self._streaming_task = asyncio.create_task(
93 self._stream_user_updates()
94 )
95 self.set_status(NodeStatus.healthy)
96 logger.info("Connected to node %i", self.id)
97 await asyncio.sleep(10)
98
99 async def _stream_user_updates(self):
100 try:
101 async with self._stub.SyncUsers.open() as stream:
102 logger.debug("opened the stream")
103 while True:
104 user_update = await self._updates_queue.get()
105 logger.debug("got something from queue")
106 user = user_update["user"]
107 await stream.send_message(
108 UserData(
109 user=User(
110 id=user.id,
111 username=user.username,
112 key=user.key,
113 ),
114 inbounds=[
115 Inbound(tag=t) for t in user_update["inbounds"]
116 ],
117 )
118 )
119 except (OSError, ConnectionError, GRPCError, StreamTerminatedError):
120 logger.info("node %i detached", self.id)
121 self.synced = False
122
123 async def update_user(self, user, inbounds: set[str] | None = None):
124 if inbounds is None:
125 inbounds = set()
126
127 await self._updates_queue.put({"user": user, "inbounds": inbounds})
128
129 async def _repopulate_users(self, users_data: list[dict]) -> None:
130 updates = [
131 UserData(
132 user=User(id=u["id"], username=u["username"], key=u["key"]),
133 inbounds=[Inbound(tag=t) for t in u["inbounds"]],
134 )
135 for u in users_data
136 ]
137 await self._stub.RepopulateUsers(UsersData(users_data=updates))
138
139 async def fetch_users_stats(self):
140 response = await self._stub.FetchUsersStats(Empty())
141 return response.users_stats
142
143 async def _fetch_backends(self) -> list:
144 response = await self._stub.FetchBackends(Empty())
145 return response.backends
146
147 async def _sync(self):
148 backends = await self._fetch_backends()
149 self.store_backends(backends)
150 users = self.list_users()
151 await self._repopulate_users(users)
152 self.synced = True
153
154 async def get_logs(self, name: str = "xray", include_buffer=True):
155 async with self._stub.StreamBackendLogs.open() as stm:
156 await stm.send_message(
157 BackendLogsRequest(
158 backend_name=name, include_buffer=include_buffer
159 )
160 )
161 while True:
162 response = await stm.recv_message()
163 yield response.line
164
165 async def restart_backend(
166 self, name: str, config: str, config_format: int
167 ):
168 try:
169 await self._stub.RestartBackend(
170 RestartBackendRequest(
171 backend_name=name,
172 config=BackendConfig(
173 configuration=config, config_format=config_format
174 ),
175 )
176 )
177 await self._sync()
178 except:
179 self.synced = False
180 self.set_status(NodeStatus.unhealthy)
181 raise
182 else:
183 self.set_status(NodeStatus.healthy)
184
185 async def get_backend_config(self, name: str):
186 response: BackendConfig = await self._stub.FetchBackendConfig(
187 Backend(name=name)
188 )
189 return response.configuration, response.config_format
190
191 async def get_backend_stats(self, name: str):
192 response: BackendStats = await self._stub.GetBackendStats(
193 Backend(name=name)
194 )
195 return response