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