From c4ba0e2120caa05e99105319328f59966cfb61a4 Mon Sep 17 00:00:00 2001 From: Ehco1996 Date: Thu, 25 Jan 2024 19:52:45 +0800 Subject: [PATCH] Refactor user configuration logic in proxy models --- apps/proxy/models.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/apps/proxy/models.py b/apps/proxy/models.py index 1bb61898c3..bf7d8df9ad 100644 --- a/apps/proxy/models.py +++ b/apps/proxy/models.py @@ -542,9 +542,14 @@ def to_node_config(self, node: ProxyNode): return configs def to_user_config(self, node: ProxyNode, user: User): - enable = node.enable and user.total_traffic > ( + have_shared_traffic = user.total_traffic > ( user.download_traffic + user.upload_traffic ) + have_oc_traffic = False + oc = UserProxyNodeOccupancy.get_by_proxy_node_and_user(node, user) + if oc: + have_oc_traffic = not oc.out_of_usage() + enable = node.enable and (have_shared_traffic or have_oc_traffic) return { "user_id": user.id, "password": user.proxy_password, @@ -596,9 +601,15 @@ def to_node_config(self, node: ProxyNode): return configs def to_user_config(self, node: ProxyNode, user: User): - enable = node.enable and user.total_traffic > ( + have_shared_traffic = user.total_traffic > ( user.download_traffic + user.upload_traffic ) + have_oc_traffic = False + oc = UserProxyNodeOccupancy.get_by_proxy_node_and_user(node, user) + if oc: + have_oc_traffic = not oc.out_of_usage() + enable = node.enable and (have_shared_traffic or have_oc_traffic) + return { "user_id": user.id, "password": user.proxy_password, @@ -1045,11 +1056,14 @@ def get_occupied_node_ids(cls): def get_node_occupancy_user_ids(cls, node: ProxyNode): return cls._valid_occupancy_query().filter(proxy_node=node).values("user_id") - @classmethod @classmethod def get_user_occupied_node_ids(cls, user: User): return cls._valid_occupancy_query().filter(user=user).values("proxy_node_id") + @classmethod + def get_by_proxy_node_and_user(cls, proxy_node: ProxyNode, user: User): + return cls.objects.filter(proxy_node=proxy_node, user=user).first() + @classmethod def get_node_occupancies(cls, node: ProxyNode): return cls._valid_occupancy_query().filter(proxy_node=node)