Refactor code and add occupancy configuration

apps/mixin.py

@@ -29,9 +29,6 @@ def truncate(cls):
         with connection.cursor() as cursor:
             cursor.execute("TRUNCATE TABLE {}".format(cls._meta.db_table))
 
-    def to_dict(self):
-        return {field.name: getattr(self, field.name) for field in self._meta.fields}
-
 
 class BaseLogModel(BaseModel):
     created_at = models.DateTimeField(

apps/proxy/admin.py

@@ -23,6 +23,17 @@ class TrojanConfigInline(admin.StackedInline):
     fields = ["proxy_node", "multi_user_port", "fallback_addr"]
 
 
+class OccupancyConfigInline(admin.StackedInline):
+    model = models.OccupancyConfig
+    verbose_name = "占用配置"
+    fields = [
+        "proxy_node",
+        "occupancy_price",
+        "occupancy_traffic",
+        "occupancy_user_limit",
+    ]
+
+
 class RelayRuleInline(admin.TabularInline):
     model = models.RelayRule
     verbose_name = "中转规则配置"
@@ -80,8 +91,13 @@ class ProxyNodeAdmin(admin.ModelAdmin):
         "sequence",
         "api_endpoint",
     ]
-    inlines = [RelayRuleInline]
-    all_inlines = [TrojanConfigInline, SSConfigInline, RelayRuleInline]
+    inlines = [RelayRuleInline, OccupancyConfigInline]
+    all_inlines = [
+        TrojanConfigInline,
+        SSConfigInline,
+        RelayRuleInline,
+        OccupancyConfigInline,
+    ]
     list_editable = ["sequence"]
     list_filter = ["node_type", "country", "provider_remark"]
     actions = ["reset_port", "clear_traffic_logs", "toggle_enable"]
@@ -208,7 +224,23 @@ def total_traffic(self, instance):
     total_traffic.short_description = "流量"
 
 
+class UserProxyNodeOccupancyAdmin(admin.ModelAdmin):
+    list_display = [
+        "proxy_node",
+        "user",
+        "start_time",
+        "end_time",
+        "traffic_used",
+        "out_of_traffic",
+    ]
+    search_fields = ["user__username"]
+    list_filter = ["proxy_node", "user"]
+    list_per_page = 10
+    show_full_result_count = False
+
+
 # Register your models here.
 admin.site.register(models.ProxyNode, ProxyNodeAdmin)
 admin.site.register(models.RelayNode, RelayNodeAdmin)
 admin.site.register(models.UserTrafficLog, UserTrafficLogAdmin)
+admin.site.register(models.UserProxyNodeOccupancy, UserProxyNodeOccupancyAdmin)

apps/proxy/migrations/0022_proxynode_cost_price_relaynode_cost_price_and_more.py

@@ -1,4 +1,4 @@
-# Generated by Django 4.2.6 on 2023-12-20 08:17
+# Generated by Django 4.2.6 on 2023-12-20 08:35
 
 import django.db.models.deletion
 from django.conf import settings
@@ -39,7 +39,7 @@ class Migration(migrations.Migration):
                     ),
                 ),
                 (
-                    "occupancy_price_30day",
+                    "occupancy_price",
                     models.DecimalField(
                         decimal_places=2, max_digits=10, verbose_name="占用 30 天价格"
                     ),
@@ -54,7 +54,7 @@ class Migration(migrations.Migration):
                 ),
                 (
                     "proxy_node",
-                    models.ForeignKey(
+                    models.OneToOneField(
                         on_delete=django.db.models.deletion.CASCADE,
                         to="proxy.proxynode",
                         verbose_name="代理节点",
@@ -67,7 +67,7 @@ class Migration(migrations.Migration):
             },
         ),
         migrations.CreateModel(
-            name="UserProxyNodeOccupancyRecord",
+            name="UserProxyNodeOccupancy",
             fields=[
                 (
                     "id",
@@ -89,7 +89,7 @@ class Migration(migrations.Migration):
                 ),
                 (
                     "out_of_traffic",
-                    models.BooleanField(default=False, verbose_name="是否超出流量"),
+                    models.BooleanField(default=False, verbose_name="流量溢出"),
                 ),
                 (
                     "occupancy_config_snapshot",
@@ -113,12 +113,12 @@ class Migration(migrations.Migration):
                 ),
             ],
             options={
-                "verbose_name": "用户代理节点占用记录",
-                "verbose_name_plural": "用户代理节点占用记录",
+                "verbose_name": "占用记录",
+                "verbose_name_plural": "占用记录",
                 "index_together": {
                     ("out_of_traffic", "user", "end_time"),
-                    ("out_of_traffic", "end_time"),
                     ("out_of_traffic", "proxy_node", "end_time"),
+                    ("out_of_traffic", "end_time"),
                 },
             },
         ),

apps/proxy/models.py

@@ -259,12 +259,10 @@ def get_user_active_nodes(cls, user):
         base_query = cls.get_active_nodes()
         query = base_query.filter(level__gte=user.level)
         # 2. filter out nodes that has been occupied by other users
-        occupied_node_ids = UserProxyNodeOccupancyRecord.get_occupied_node_ids()
+        occupied_node_ids = UserProxyNodeOccupancy.get_occupied_node_ids()
         query = query.exclude(id__in=occupied_node_ids)
         # 3. add nodes that has been occupied by this user
-        user_occupied_node_ids = (
-            UserProxyNodeOccupancyRecord.get_user_occupied_node_ids(user)
-        )
+        user_occupied_node_ids = UserProxyNodeOccupancy.get_user_occupied_node_ids(user)
         query = query | base_query.filter(id__in=user_occupied_node_ids)
         return query
 
@@ -283,7 +281,7 @@ def get_node_users(self):
         if not self.enable:
             return []
         # 2. node occupied by users, return users
-        occupancies_query = UserProxyNodeOccupancyRecord.get_node_occupancies(self)
+        occupancies_query = UserProxyNodeOccupancy.get_node_occupancies(self)
         if occupancies_query.count() > 0:
             user_ids = occupancies_query.values("user_id")
             return User.objects.filter(id__in=user_ids)
@@ -493,28 +491,28 @@ def to_node_config(self, node: ProxyNode):
         )
         ss_config = self
         ss_inbound = deepcopy(XRayTemplates.SS_INBOUND)
-        ss_inbound["listen"] = self.get_inbound_listen_host()
+        ss_inbound["listen"] = node.get_inbound_listen_host()
         ss_inbound["port"] = ss_config.multi_user_port
-        if self.enable_udp:
+        if node.enable_udp:
             ss_inbound["settings"]["network"] += ",udp"
         xray_config["inbounds"].append(ss_inbound)
         configs = {
             "xray_config": xray_config,
             "sync_traffic_endpoint": node.api_endpoint,
         }
-        configs.update(self.get_ehco_server_config())
+        configs.update(node.get_ehco_server_config())
         return configs
 
     def to_user_config(self, node: ProxyNode, user: User):
-        enable = self.enable and user.total_traffic > (
+        enable = node.enable and user.total_traffic > (
             user.download_traffic + user.upload_traffic
         )
         return {
             "user_id": user.id,
             "password": user.proxy_password,
             "enable": enable,
             "method": self.method,
-            "protocol": self.NODE_TYPE_SS,
+            "protocol": ProxyNode.NODE_TYPE_SS,
         }
 
 
@@ -545,7 +543,7 @@ def to_node_config(self, node: ProxyNode):
             node.ehco_log_level,
         )
         inbound = deepcopy(XRayTemplates.TROJAN_INBOUND)
-        inbound["listen"] = self.get_inbound_listen_host()
+        inbound["listen"] = node.get_inbound_listen_host()
         inbound["port"] = self.multi_user_port
         inbound["settings"]["fallbacks"][0]["dest"] = self.fallback_addr
         if node.enable_udp:
@@ -560,14 +558,14 @@ def to_node_config(self, node: ProxyNode):
         return configs
 
     def to_user_config(self, node: ProxyNode, user: User):
-        enable = self.enable and user.total_traffic > (
+        enable = node.enable and user.total_traffic > (
             user.download_traffic + user.upload_traffic
         )
         return {
             "user_id": user.id,
             "password": user.proxy_password,
             "enable": enable,
-            "protocol": self.NODE_TYPE_TROJAN,
+            "protocol": ProxyNode.NODE_TYPE_TROJAN,
         }
 
 
@@ -840,44 +838,62 @@ def total_traffic(self):
 
 
 class OccupancyConfig(BaseModel):
-    proxy_node = models.ForeignKey(
-        ProxyNode, on_delete=models.CASCADE, verbose_name="代理节点", db_index=True
+    proxy_node = models.OneToOneField(
+        ProxyNode,
+        on_delete=models.CASCADE,
+        verbose_name="代理节点",
+        db_index=True,
+        related_name="occupancy_config",
     )
-    occupancy_price_30day = models.DecimalField(
-        max_digits=10, decimal_places=2, verbose_name="占用 30 天价格"
+    occupancy_price = models.DecimalField(
+        max_digits=10, decimal_places=2, verbose_name="价格"
     )
-    occupancy_traffic = models.BigIntegerField(default=0, verbose_name="已用流量")
-    occupancy_user_limit = models.PositiveIntegerField(verbose_name="占用用户限制", default=0)
+    occupancy_traffic = models.BigIntegerField(default=0, verbose_name="流量")
+    occupancy_user_limit = models.PositiveIntegerField(verbose_name="用户数", default=0)
 
     class Meta:
         verbose_name = "占用配置"
         verbose_name_plural = "占用配置"
 
+    def __str__(self) -> str:
+        return f"占用配置:{self.id}"
+
     @classmethod
     def get_by_proxy_node(cls, node: ProxyNode):
         return cls.objects.filter(proxy_node=node).first()
 
+    def to_snapshot(self):
+        return {
+            "proxy_node_id": self.proxy_node.id,
+            "occupancy_price": self.occupancy_price,
+            "occupancy_traffic": self.occupancy_traffic,
+            "occupancy_user_limit": self.occupancy_user_limit,
+        }
 
-class UserProxyNodeOccupancyRecord(BaseModel):
+
+class UserProxyNodeOccupancy(BaseModel):
     user = models.ForeignKey(User, on_delete=models.CASCADE, verbose_name="用户")
     proxy_node = models.ForeignKey(
         ProxyNode, on_delete=models.CASCADE, verbose_name="代理节点"
     )
     start_time = models.DateTimeField(auto_now_add=True, verbose_name="开始占用时间")
     end_time = models.DateTimeField(null=False, blank=False, verbose_name="结束占用时间")
     traffic_used = models.BigIntegerField(default=0, verbose_name="已用流量")
-    out_of_traffic = models.BooleanField(default=False, verbose_name="是否超出流量")
+    out_of_traffic = models.BooleanField(default=False, verbose_name="流量溢出")
     occupancy_config_snapshot = models.JSONField(verbose_name="快照", default=dict)
 
     class Meta:
-        verbose_name = "用户代理节点占用记录"
-        verbose_name_plural = "用户代理节点占用记录"
+        verbose_name = "占用记录"
+        verbose_name_plural = "占用记录"
         index_together = (
             ["out_of_traffic", "end_time"],
             ["out_of_traffic", "user", "end_time"],
             ["out_of_traffic", "proxy_node", "end_time"],
         )
 
+    def __str__(self) -> str:
+        return f"用户占用配置:{self.id}"
+
     @classmethod
     @transaction.atomic
     def create_by_occupancy_config(
@@ -901,7 +917,7 @@ def create_by_occupancy_config(
             start_time=utils.get_current_datetime(),
             end_time=utils.get_current_datetime().add(days=30),
             traffic_used=occupancy_config.occupancy_traffic,
-            occupancy_config_snapshot=occupancy_config.to_dict(),
+            occupancy_config_snapshot=occupancy_config.to_snapshot(),
         )
 
     @classmethod
@@ -930,8 +946,8 @@ def get_user_occupied_node_ids(cls, user: User):
         return user_occupied_node_ids
 
     @classmethod
-    def get_node_occupies(cls, node: ProxyNode):
-        return UserProxyNodeOccupancyRecord.objects.filter(
+    def get_node_occupancies(cls, node: ProxyNode):
+        return UserProxyNodeOccupancy.objects.filter(
             out_of_traffic=False,
             proxy_node=node,
             end_time__gte=utils.get_current_datetime(),

apps/sspanel/tasks.py

@@ -4,7 +4,7 @@
 from django.core.mail import send_mail
 
 from apps import celery_app
-from apps.proxy.models import ProxyNode, UserProxyNodeOccupancyRecord, UserTrafficLog
+from apps.proxy.models import ProxyNode, UserProxyNodeOccupancy, UserTrafficLog
 from apps.sspanel import models as m
 from apps.utils import get_current_datetime
 
@@ -28,9 +28,7 @@ def sync_user_traffic_task(node_id, data):
     node: ProxyNode = ProxyNode.get_or_none(node_id)
     if not node:
         return
-    node_occurred_user_ids = UserProxyNodeOccupancyRecord.get_node_occupancy_user_ids(
-        node
-    )
+    node_occurred_user_ids = UserProxyNodeOccupancy.get_node_occupancy_user_ids(node)
     node_total_traffic = 0
     log_time = get_current_datetime()
     user_model_list = []
@@ -71,7 +69,7 @@ def sync_user_traffic_task(node_id, data):
 
         # 记录用户占用节点流量
         if user_id in node_occurred_user_ids:
-            UserProxyNodeOccupancyRecord.check_and_incr_traffic(
+            UserProxyNodeOccupancy.check_and_incr_traffic(
                 user_id=user_id, node_id=node_id, traffic=d + u
             )