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 UserStatus
10
11if TYPE_CHECKING:
12    from app.db import User as DBUser
13
14
15async def update_user(user: "DBUser", old_inbounds: set | None = None):
16    """updates a user on all nodes
17    even though this isn't efficient it is extremely precise"""
18    if old_inbounds is None:
19        old_inbounds = set()
20
21    node_inbounds = defaultdict(list)
22    if user.status in (UserStatus.on_hold, UserStatus.active):
23        for inb in user.inbounds:
24            node_inbounds[inb.node_id].append(inb.tag)
25    else:
26        for inb in user.inbounds:
27            node_inbounds[inb.node_id]
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.create_task(
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            await marznode.nodes[node_id].update_user(user=user, inbounds=[])
45
46
47async def remove_node(node_id: int):
48    if node_id in marznode.nodes:
49        del marznode.nodes[node_id]
50
51
52async def add_node(db_node, certificate):
53    await remove_node(db_node.id)
54    if db_node.connection_backend == NodeConnectionBackend.grpcio:
55        node = MarzNodeGRPCIO(
56            db_node.id,
57            db_node.address,
58            db_node.port,
59            usage_coefficient=db_node.usage_coefficient,
60        )
61    else:
62        node = MarzNodeGRPCLIB(
63            db_node.id,
64            db_node.address,
65            db_node.port,
66            certificate.key,
67            certificate.certificate,
68            usage_coefficient=db_node.usage_coefficient,
69        )
70    marznode.nodes[db_node.id] = node
71
72
73__all__ = ["update_user", "add_node", "remove_node"]