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