gitscope
/
← marzneshin / app/marznode/grpclib.py / blame
Files Commits 446 Branches 1 Tags 43 Markers 4 README License
https://dev.kronoz.win/r/marzneshin.git
app/marznode/grpclib.py · v0.3.1
Source
1 f6a32ea Dawsh 2y ago import asyncio
2 f6a32ea Dawsh 2y ago import atexit
3 f6a32ea Dawsh 2y ago import logging
4 f6a32ea Dawsh 2y ago import ssl
5 f6a32ea Dawsh 2y ago import tempfile
6 f6a32ea Dawsh 2y ago
7 f6a32ea Dawsh 2y ago from grpclib import GRPCError
8 f6a32ea Dawsh 2y ago from grpclib.client import Channel
9 f6a32ea Dawsh 2y ago from grpclib.exceptions import StreamTerminatedError
10 f6a32ea Dawsh 2y ago
11 f6a32ea Dawsh 2y ago from .base import MarzNodeBase
12 f6a32ea Dawsh 2y ago from .database import MarzNodeDB
13 f6a32ea Dawsh 2y ago from .marznode_grpc import MarzServiceStub
14 01cdf2e Dawsh 2y ago from .marznode_pb2 import (
15 01cdf2e Dawsh 2y ago UserData,
16 01cdf2e Dawsh 2y ago UsersData,
17 01cdf2e Dawsh 2y ago Empty,
18 01cdf2e Dawsh 2y ago User,
19 01cdf2e Dawsh 2y ago Inbound,
20 01cdf2e Dawsh 2y ago XrayLogsRequest,
21 01cdf2e Dawsh 2y ago XrayConfig,
22 01cdf2e Dawsh 2y ago )
23 f6a32ea Dawsh 2y ago from ..models.node import NodeStatus
24 f6a32ea Dawsh 2y ago
25 f6a32ea Dawsh 2y ago logger = logging.getLogger(__name__)
26 f6a32ea Dawsh 2y ago
27 f6a32ea Dawsh 2y ago
28 f6a32ea Dawsh 2y ago def string_to_temp_file(content: str):
29 01cdf2e Dawsh 2y ago file = tempfile.NamedTemporaryFile(mode="w+t")
30 f6a32ea Dawsh 2y ago file.write(content)
31 f6a32ea Dawsh 2y ago file.flush()
32 f6a32ea Dawsh 2y ago return file
33 f6a32ea Dawsh 2y ago
34 f6a32ea Dawsh 2y ago
35 f6a32ea Dawsh 2y ago class MarzNodeGRPCLIB(MarzNodeBase, MarzNodeDB):
36 01cdf2e Dawsh 2y ago def __init__(
37 01cdf2e Dawsh 2y ago self,
38 01cdf2e Dawsh 2y ago node_id: int,
39 01cdf2e Dawsh 2y ago address: str,
40 01cdf2e Dawsh 2y ago port: int,
41 01cdf2e Dawsh 2y ago ssl_key: str,
42 01cdf2e Dawsh 2y ago ssl_cert: str,
43 01cdf2e Dawsh 2y ago usage_coefficient: int = 1,
44 01cdf2e Dawsh 2y ago ):
45 f6a32ea Dawsh 2y ago self.id = node_id
46 f6a32ea Dawsh 2y ago self._address = address
47 f6a32ea Dawsh 2y ago self._port = port
48 f6a32ea Dawsh 2y ago
49 f6a32ea Dawsh 2y ago self._key_file = string_to_temp_file(ssl_key)
50 f6a32ea Dawsh 2y ago self._cert_file = string_to_temp_file(ssl_cert)
51 f6a32ea Dawsh 2y ago
52 f6a32ea Dawsh 2y ago ctx = ssl.create_default_context()
53 f6a32ea Dawsh 2y ago ctx.load_cert_chain(self._cert_file.name, self._key_file.name)
54 f6a32ea Dawsh 2y ago ctx.check_hostname = False
55 f6a32ea Dawsh 2y ago ctx.verify_mode = ssl.CERT_NONE
56 f6a32ea Dawsh 2y ago
57 f6a32ea Dawsh 2y ago self._channel = Channel(self._address, self._port, ssl=ctx)
58 f6a32ea Dawsh 2y ago self._stub = MarzServiceStub(self._channel)
59 f6a32ea Dawsh 2y ago asyncio.create_task(self._monitor_channel())
60 f6a32ea Dawsh 2y ago self._streaming_task = None
61 f6a32ea Dawsh 2y ago
62 f6a32ea Dawsh 2y ago self._updates_queue = asyncio.Queue(1)
63 f6a32ea Dawsh 2y ago self.synced = False
64 f6a32ea Dawsh 2y ago self.usage_coefficient = usage_coefficient
65 f6a32ea Dawsh 2y ago atexit.register(self._channel.close)
66 f6a32ea Dawsh 2y ago
67 f6a32ea Dawsh 2y ago async def _monitor_channel(self):
68 6504ac9 Dawsh 2y ago while state := self._channel._state:
69 f6a32ea Dawsh 2y ago logger.debug("node %i channel state: %s", self.id, state.value)
70 f6a32ea Dawsh 2y ago try:
71 f6a32ea Dawsh 2y ago await asyncio.wait_for(self._channel.__connect__(), timeout=2)
72 f6a32ea Dawsh 2y ago except Exception:
73 512d183 Mrdn 2y ago logger.debug("timeout for node, id: %i", self.id)
74 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.unhealthy, "timeout")
75 f6a32ea Dawsh 2y ago self.synced = False
76 f6a32ea Dawsh 2y ago if self._streaming_task:
77 f6a32ea Dawsh 2y ago self._streaming_task.cancel()
78 f6a32ea Dawsh 2y ago else:
79 f6a32ea Dawsh 2y ago if not self.synced:
80 4a2faf0 Dawsh 2y ago try:
81 4a2faf0 Dawsh 2y ago await self._sync()
82 4a2faf0 Dawsh 2y ago except:
83 4a2faf0 Dawsh 2y ago pass
84 4a2faf0 Dawsh 2y ago else:
85 4a2faf0 Dawsh 2y ago self._streaming_task = asyncio.create_task(
86 4a2faf0 Dawsh 2y ago self._stream_user_updates()
87 4a2faf0 Dawsh 2y ago )
88 4a2faf0 Dawsh 2y ago self.set_status(NodeStatus.healthy)
89 4a2faf0 Dawsh 2y ago logger.info("Connected to node %i", self.id)
90 512d183 Mrdn 2y ago await asyncio.sleep(10)
91 f6a32ea Dawsh 2y ago
92 f6a32ea Dawsh 2y ago async def _stream_user_updates(self):
93 f6a32ea Dawsh 2y ago try:
94 f6a32ea Dawsh 2y ago async with self._stub.SyncUsers.open() as stream:
95 4a2faf0 Dawsh 2y ago logger.debug("opened the stream")
96 f6a32ea Dawsh 2y ago while True:
97 f6a32ea Dawsh 2y ago user_update = await self._updates_queue.get()
98 4a2faf0 Dawsh 2y ago logger.debug("got something from queue")
99 f6a32ea Dawsh 2y ago user = user_update["user"]
100 f6a32ea Dawsh 2y ago await stream.send_message(
101 01cdf2e Dawsh 2y ago UserData(
102 f1ab16c Dawsh 2y ago user=User(
103 f1ab16c Dawsh 2y ago id=user.id,
104 f1ab16c Dawsh 2y ago username=user.username,
105 f1ab16c Dawsh 2y ago key=user.key,
106 f1ab16c Dawsh 2y ago ),
107 f1ab16c Dawsh 2y ago inbounds=[
108 f1ab16c Dawsh 2y ago Inbound(tag=t) for t in user_update["inbounds"]
109 f1ab16c Dawsh 2y ago ],
110 01cdf2e Dawsh 2y ago )
111 01cdf2e Dawsh 2y ago )
112 f6a32ea Dawsh 2y ago except (OSError, ConnectionError, GRPCError, StreamTerminatedError):
113 f6a32ea Dawsh 2y ago logger.info("node %i detached", self.id)
114 f6a32ea Dawsh 2y ago self.synced = False
115 f6a32ea Dawsh 2y ago
116 f6a32ea Dawsh 2y ago async def update_user(self, user, inbounds: set[str] | None = None):
117 f6a32ea Dawsh 2y ago if inbounds is None:
118 f6a32ea Dawsh 2y ago inbounds = set()
119 f6a32ea Dawsh 2y ago
120 4a2faf0 Dawsh 2y ago await self._updates_queue.put({"user": user, "inbounds": inbounds})
121 f6a32ea Dawsh 2y ago
122 f6a32ea Dawsh 2y ago async def _repopulate_users(self, users_data: list[dict]) -> None:
123 01cdf2e Dawsh 2y ago updates = [
124 01cdf2e Dawsh 2y ago UserData(
125 01cdf2e Dawsh 2y ago user=User(id=u["id"], username=u["username"], key=u["key"]),
126 01cdf2e Dawsh 2y ago inbounds=[Inbound(tag=t) for t in u["inbounds"]],
127 01cdf2e Dawsh 2y ago )
128 01cdf2e Dawsh 2y ago for u in users_data
129 01cdf2e Dawsh 2y ago ]
130 f6a32ea Dawsh 2y ago await self._stub.RepopulateUsers(UsersData(users_data=updates))
131 f6a32ea Dawsh 2y ago
132 f6a32ea Dawsh 2y ago async def fetch_users_stats(self):
133 f6a32ea Dawsh 2y ago response = await self._stub.FetchUsersStats(Empty())
134 f6a32ea Dawsh 2y ago return response.users_stats
135 f6a32ea Dawsh 2y ago
136 f6a32ea Dawsh 2y ago async def _fetch_inbounds(self) -> list:
137 f6a32ea Dawsh 2y ago response = await self._stub.FetchInbounds(Empty())
138 f6a32ea Dawsh 2y ago return response.inbounds
139 f6a32ea Dawsh 2y ago
140 f6a32ea Dawsh 2y ago async def _sync(self):
141 f6a32ea Dawsh 2y ago inbounds = await self._fetch_inbounds()
142 f6a32ea Dawsh 2y ago self.store_inbounds(inbounds)
143 f6a32ea Dawsh 2y ago users = self.list_users()
144 f6a32ea Dawsh 2y ago await self._repopulate_users(users)
145 f6a32ea Dawsh 2y ago self.synced = True
146 512d183 Mrdn 2y ago
147 512d183 Mrdn 2y ago async def get_logs(self, include_buffer=True):
148 512d183 Mrdn 2y ago async with self._stub.StreamXrayLogs.open() as stm:
149 f1ab16c Dawsh 2y ago await stm.send_message(
150 f1ab16c Dawsh 2y ago XrayLogsRequest(include_buffer=include_buffer)
151 f1ab16c Dawsh 2y ago )
152 512d183 Mrdn 2y ago while True:
153 512d183 Mrdn 2y ago response = await stm.recv_message()
154 512d183 Mrdn 2y ago yield response.line
155 512d183 Mrdn 2y ago
156 512d183 Mrdn 2y ago async def restart_xray(self, config: str):
157 758c903 Dawsh 2y ago try:
158 758c903 Dawsh 2y ago await self._stub.RestartXray(XrayConfig(configuration=config))
159 758c903 Dawsh 2y ago await self._sync()
160 758c903 Dawsh 2y ago except:
161 758c903 Dawsh 2y ago self.synced = False
162 758c903 Dawsh 2y ago self.set_status(NodeStatus.unhealthy)
163 758c903 Dawsh 2y ago raise
164 758c903 Dawsh 2y ago else:
165 758c903 Dawsh 2y ago self.set_status(NodeStatus.healthy)
166 512d183 Mrdn 2y ago
167 512d183 Mrdn 2y ago async def get_xray_config(self):
168 512d183 Mrdn 2y ago response = await self._stub.FetchXrayConfig(Empty())
169 512d183 Mrdn 2y ago return response.configuration