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