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.2.0
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 512d183 Mrdn 2y ago from .marznode_pb2 import UserData, UsersData, Empty, User, Inbound, XrayLogsRequest, XrayConfig
15 f6a32ea Dawsh 2y ago from ..models.node import NodeStatus
16 f6a32ea Dawsh 2y ago
17 f6a32ea Dawsh 2y ago logger = logging.getLogger(__name__)
18 f6a32ea Dawsh 2y ago
19 f6a32ea Dawsh 2y ago
20 f6a32ea Dawsh 2y ago def string_to_temp_file(content: str):
21 f6a32ea Dawsh 2y ago file = tempfile.NamedTemporaryFile(mode='w+t')
22 f6a32ea Dawsh 2y ago file.write(content)
23 f6a32ea Dawsh 2y ago file.flush()
24 f6a32ea Dawsh 2y ago return file
25 f6a32ea Dawsh 2y ago
26 f6a32ea Dawsh 2y ago
27 f6a32ea Dawsh 2y ago class MarzNodeGRPCLIB(MarzNodeBase, MarzNodeDB):
28 f6a32ea Dawsh 2y ago def __init__(self, node_id: int, address: str, port: int, ssl_key: str, ssl_cert: str, usage_coefficient: int = 1):
29 f6a32ea Dawsh 2y ago self.id = node_id
30 f6a32ea Dawsh 2y ago self._address = address
31 f6a32ea Dawsh 2y ago self._port = port
32 f6a32ea Dawsh 2y ago
33 f6a32ea Dawsh 2y ago self._key_file = string_to_temp_file(ssl_key)
34 f6a32ea Dawsh 2y ago self._cert_file = string_to_temp_file(ssl_cert)
35 f6a32ea Dawsh 2y ago
36 f6a32ea Dawsh 2y ago ctx = ssl.create_default_context()
37 f6a32ea Dawsh 2y ago ctx.load_cert_chain(self._cert_file.name, self._key_file.name)
38 f6a32ea Dawsh 2y ago ctx.check_hostname = False
39 f6a32ea Dawsh 2y ago ctx.verify_mode = ssl.CERT_NONE
40 f6a32ea Dawsh 2y ago
41 f6a32ea Dawsh 2y ago self._channel = Channel(self._address, self._port, ssl=ctx)
42 f6a32ea Dawsh 2y ago self._stub = MarzServiceStub(self._channel)
43 f6a32ea Dawsh 2y ago asyncio.create_task(self._monitor_channel())
44 f6a32ea Dawsh 2y ago self._streaming_task = None
45 f6a32ea Dawsh 2y ago
46 f6a32ea Dawsh 2y ago self._updates_queue = asyncio.Queue(1)
47 f6a32ea Dawsh 2y ago self.synced = False
48 f6a32ea Dawsh 2y ago self.usage_coefficient = usage_coefficient
49 f6a32ea Dawsh 2y ago atexit.register(self._channel.close)
50 f6a32ea Dawsh 2y ago
51 f6a32ea Dawsh 2y ago async def _monitor_channel(self):
52 6504ac9 Dawsh 2y ago while state := self._channel._state:
53 f6a32ea Dawsh 2y ago logger.debug("node %i channel state: %s", self.id, state.value)
54 f6a32ea Dawsh 2y ago try:
55 f6a32ea Dawsh 2y ago await asyncio.wait_for(self._channel.__connect__(), timeout=2)
56 f6a32ea Dawsh 2y ago except Exception:
57 512d183 Mrdn 2y ago logger.debug("timeout for node, id: %i", self.id)
58 f6a32ea Dawsh 2y ago self.set_status(NodeStatus.unhealthy, "timeout")
59 f6a32ea Dawsh 2y ago self.synced = False
60 f6a32ea Dawsh 2y ago if self._streaming_task:
61 f6a32ea Dawsh 2y ago self._streaming_task.cancel()
62 f6a32ea Dawsh 2y ago else:
63 f6a32ea Dawsh 2y ago if not self.synced:
64 512d183 Mrdn 2y ago await self._sync()
65 512d183 Mrdn 2y ago self._streaming_task = asyncio.create_task(self._stream_user_updates())
66 512d183 Mrdn 2y ago self.set_status(NodeStatus.healthy)
67 512d183 Mrdn 2y ago logger.info("Connected to node %i", self.id)
68 512d183 Mrdn 2y ago await asyncio.sleep(10)
69 f6a32ea Dawsh 2y ago
70 f6a32ea Dawsh 2y ago async def _stream_user_updates(self):
71 f6a32ea Dawsh 2y ago try:
72 f6a32ea Dawsh 2y ago async with self._stub.SyncUsers.open() as stream:
73 f6a32ea Dawsh 2y ago logger.info("opened the stream")
74 f6a32ea Dawsh 2y ago while True:
75 f6a32ea Dawsh 2y ago user_update = await self._updates_queue.get()
76 f6a32ea Dawsh 2y ago logger.info("got something from queue")
77 f6a32ea Dawsh 2y ago user = user_update["user"]
78 f6a32ea Dawsh 2y ago await stream.send_message(
79 f6a32ea Dawsh 2y ago UserData(user=User(id=user.id, username=user.username, key=user.key),
80 f6a32ea Dawsh 2y ago inbounds=[Inbound(tag=t) for t in user_update["inbounds"]]))
81 f6a32ea Dawsh 2y ago except (OSError, ConnectionError, GRPCError, StreamTerminatedError):
82 f6a32ea Dawsh 2y ago logger.info("node %i detached", self.id)
83 f6a32ea Dawsh 2y ago self.synced = False
84 f6a32ea Dawsh 2y ago
85 f6a32ea Dawsh 2y ago async def update_user(self, user, inbounds: set[str] | None = None):
86 f6a32ea Dawsh 2y ago if inbounds is None:
87 f6a32ea Dawsh 2y ago inbounds = set()
88 f6a32ea Dawsh 2y ago
89 f6a32ea Dawsh 2y ago if self.synced:
90 f6a32ea Dawsh 2y ago await self._updates_queue.put({"user": user, "inbounds": inbounds})
91 f6a32ea Dawsh 2y ago
92 f6a32ea Dawsh 2y ago async def _repopulate_users(self, users_data: list[dict]) -> None:
93 f6a32ea Dawsh 2y ago updates = [UserData(user=User(id=u["id"], username=u["username"], key=u["key"]),
94 f6a32ea Dawsh 2y ago inbounds=[Inbound(tag=t) for t in u["inbounds"]]) for u in users_data]
95 f6a32ea Dawsh 2y ago await self._stub.RepopulateUsers(UsersData(users_data=updates))
96 f6a32ea Dawsh 2y ago
97 f6a32ea Dawsh 2y ago async def fetch_users_stats(self):
98 f6a32ea Dawsh 2y ago response = await self._stub.FetchUsersStats(Empty())
99 f6a32ea Dawsh 2y ago return response.users_stats
100 f6a32ea Dawsh 2y ago
101 f6a32ea Dawsh 2y ago async def _fetch_inbounds(self) -> list:
102 f6a32ea Dawsh 2y ago response = await self._stub.FetchInbounds(Empty())
103 f6a32ea Dawsh 2y ago return response.inbounds
104 f6a32ea Dawsh 2y ago
105 f6a32ea Dawsh 2y ago async def _sync(self):
106 f6a32ea Dawsh 2y ago inbounds = await self._fetch_inbounds()
107 f6a32ea Dawsh 2y ago self.store_inbounds(inbounds)
108 f6a32ea Dawsh 2y ago users = self.list_users()
109 f6a32ea Dawsh 2y ago await self._repopulate_users(users)
110 f6a32ea Dawsh 2y ago self.synced = True
111 512d183 Mrdn 2y ago
112 512d183 Mrdn 2y ago async def get_logs(self, include_buffer=True):
113 512d183 Mrdn 2y ago async with self._stub.StreamXrayLogs.open() as stm:
114 512d183 Mrdn 2y ago await stm.send_message(XrayLogsRequest(include_buffer=include_buffer))
115 512d183 Mrdn 2y ago while True:
116 512d183 Mrdn 2y ago response = await stm.recv_message()
117 512d183 Mrdn 2y ago yield response.line
118 512d183 Mrdn 2y ago
119 512d183 Mrdn 2y ago async def restart_xray(self, config: str):
120 512d183 Mrdn 2y ago await self._stub.RestartXray(XrayConfig(configuration=config))
121 512d183 Mrdn 2y ago await self._sync()
122 512d183 Mrdn 2y ago
123 512d183 Mrdn 2y ago async def get_xray_config(self):
124 512d183 Mrdn 2y ago response = await self._stub.FetchXrayConfig(Empty())
125 512d183 Mrdn 2y ago return response.configuration