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