1import asyncio
 2from collections import defaultdict
 3from typing import TYPE_CHECKING
 4
 5from app import marznode
 6from .grpcio import MarzNodeGRPCIO
 7from .grpclib import MarzNodeGRPCLIB
 8from ..models.node import NodeConnectionBackend
 9
10if TYPE_CHECKING:
11    from app.db import User as DBUser
12
13
14def update_user(
15    user: "DBUser", old_inbounds: set | None = None, remove: bool = False
16):
17    """updates a user on all related nodes"""
18    if old_inbounds is None:
19        old_inbounds = set()
20
21    node_inbounds = defaultdict(list)
22    if remove:
23        for inb in user.inbounds:
24            node_inbounds[inb.node_id]
25    else:
26        for inb in user.inbounds:
27            node_inbounds[inb.node_id].append(inb.tag)
28
29    for inb in old_inbounds:
30        node_inbounds[inb[0]]
31
32    for node_id, tags in node_inbounds.items():
33        if marznode.nodes.get(node_id):
34            asyncio.ensure_future(
35                marznode.nodes[node_id].update_user(user=user, inbounds=tags)
36            )
37
38
39async def remove_user(user: "DBUser"):
40    node_ids = set(inb.node_id for inb in user.inbounds)
41
42    for node_id in node_ids:
43        if marznode.nodes.get(node_id):
44            asyncio.ensure_future(
45                marznode.nodes[node_id].update_user(user=user, inbounds=[])
46            )
47
48
49async def remove_node(node_id: int):
50    if node_id in marznode.nodes:
51        del marznode.nodes[node_id]
52
53
54async def add_node(db_node, certificate):
55    await remove_node(db_node.id)
56    if db_node.connection_backend == NodeConnectionBackend.grpcio:
57        node = MarzNodeGRPCIO(
58            db_node.id,
59            db_node.address,
60            db_node.port,
61            usage_coefficient=db_node.usage_coefficient,
62        )
63    else:
64        node = MarzNodeGRPCLIB(
65            db_node.id,
66            db_node.address,
67            db_node.port,
68            certificate.key,
69            certificate.certificate,
70            usage_coefficient=db_node.usage_coefficient,
71        )
72    marznode.nodes[db_node.id] = node
73
74
75__all__ = ["update_user", "add_node", "remove_node"]