Refactor user configuration logic in proxy models

Ehco1996 · Jan 25, 2024 · c4ba0e2 · c4ba0e2
1 parent 459984e
commit c4ba0e2
Showing 1 changed file with 17 additions and 3 deletions.
diff --git 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)