Merge pull request #129 from wuba/develop

Develop

modules/config/views.py

@@ -22,12 +22,14 @@ class ConfigViewSet(mixins.ListModelMixin, GenericViewSet):
 
     def get_paginated_response(self, data):
         _data = {}
-        mapping = {"1": True, "true": True, "0": False, "false": False}
+        mapping = {"1": True, "true": True, "0": False, "false": False, "True": True, "False": False}
         for i in data:
+            print(i)
             if i["name"] == "REGISTER_TYPE":
                 _data[i["name"]] = int(i['value'])
             else:
                 _data[i["name"]] = mapping.get(i['value'], i['value'])
+        print(_data)
         return Response(_data, status=status.HTTP_200_OK)
 
     @action(methods=["POST"], detail=False, permission_classes=[IsAdminUser, ])
@@ -76,6 +78,7 @@ def check_version(self, requests, *args, **kwargs):
         except Exception as e:
             return Response({"code": 0, "message": f"检查更新错误,原因:{e}"}, status=status.HTTP_200_OK)
 
+
 class DnsConfigViewSet(mixins.ListModelMixin, GenericViewSet):
     queryset = DnsConfig.objects.all().order_by("id")
     serializer_class = DnsConfigSerializer

modules/message/views.py

@@ -102,7 +102,7 @@ def show_dashboard_url(user_id):
         url_list = [
             {
                 "task_name": tci.task.name,
-                tci.template.name: get_payload(tci.task_config.key, tci.template.payload)
+                tci.template.name: get_payload(tci.task_config, tci.template.payload)
             } for tci in task_config_item_record[:30]
         ]
         return url_list

modules/task/migrations/0003_taskconfigitem_url_template.py

@@ -0,0 +1,20 @@
+# Generated by Django 3.2.9 on 2023-04-26 03:48
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('template', '0004_urltemplate'),
+        ('task', '0002_auto_20220826_1025'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='taskconfigitem',
+            name='url_template',
+            field=models.ForeignKey(db_column=False, help_text='链接模板', null=True, on_delete=django.db.models.deletion.CASCADE, related_name='url_template_task_config_item', to='template.urltemplate'),
+        ),
+    ]

modules/task/migrations/0004_auto_20230523_1410.py

@@ -0,0 +1,24 @@
+# Generated by Django 3.2.9 on 2023-05-23 06:10
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('template', '0004_urltemplate'),
+        ('task', '0003_taskconfigitem_url_template'),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name='taskconfigitem',
+            name='url_template',
+        ),
+        migrations.AddField(
+            model_name='taskconfig',
+            name='url_template',
+            field=models.ForeignKey(db_column=False, help_text='链接模板', null=True, on_delete=django.db.models.deletion.CASCADE, related_name='url_template_task_config_item', to='template.urltemplate'),
+        ),
+    ]

modules/task/models.py

@@ -1,7 +1,7 @@
 from django.contrib.auth.models import User
 from django.db import models
 from modules.task.constants import TASK_TMP, TASK_STATUS, SHOW_DASHBOARD
-from modules.template.models import Template, TemplateConfigItem
+from modules.template.models import Template, TemplateConfigItem, UrlTemplate
 
 
 class Task(models.Model):
@@ -33,6 +33,10 @@ class TaskConfig(models.Model):
                              help_text="所属任务")
     key = models.CharField(max_length=32, help_text='key')
 
+    url_template = models.ForeignKey(UrlTemplate, related_name='url_template_task_config_item',
+                                     on_delete=models.CASCADE, db_column=False,
+                                     null=True, help_text="链接模板")
+
     class Meta:
         db_table = 'task_config'
 

modules/task/serializers.py

@@ -22,6 +22,7 @@ class CreateTaskConfigItemSerializer(serializers.Serializer):
     template = serializers.IntegerField(required=True, help_text="组件id")
     template_config_item_list = serializers.JSONField(required=True, help_text="组件配置列表")
     task = serializers.IntegerField(required=True, help_text="任务id")
+    url_template = serializers.IntegerField(required=True, help_text="url模板id")
 
 
 class TaskConfigSerializer(serializers.ModelSerializer):

modules/task/views.py

@@ -3,7 +3,7 @@
 
 from django.db import transaction
 from modules.api.models import ApiKey
-from modules.template.models import Template
+from modules.template.models import Template, UrlTemplate
 
 from django_filters.rest_framework import DjangoFilterBackend
 from modules.task.constants import TASK_STATUS, TASK_TMP
@@ -137,67 +137,74 @@ def get_result_data(data):
         """
         修改任务详情接口返回格式
         """
-        if not data:
-            return Response(data={}, status=status.HTTP_200_OK)
-        task_id = data[0]["task"]
-        task_record = Task.objects.get(id=task_id)
-
-        task_configs = TaskConfig.objects.in_bulk(item["task_config"] for item in data)
-        templates = Template.objects.in_bulk(item["template"] for item in data)
-
-        listen_data_list = [{
-            "task": item["task"],
-            "template": item["template"],
-            "template_name": templates[item["template"]].name,
-            "template_type": templates[item["template"]].type,
-            "template_choice_type": templates[item["template"]].choice_type,
-            "task_config_id": item["task_config"],
-            "key": get_payload(task_configs[item["task_config"]].key, templates[item["template"]].payload),
-            "task_config_item_list": [{
-                "template_config_item": item["template_config_item"],
-                "id": item["id"],
-                "value": item["value"]
-            }]
-        } for item in data if templates[item["template"]].type == 1]
-
-        payload_data_list = []
-        for item in data:
-            task_config = task_configs[item["task_config"]]
-            template = templates[item["template"]]
-            if template.type != 1:
-                url = get_payload(task_config.key, template.payload)
-                task_config_item_list = {
+        try:
+            if not data:
+                return Response(data={}, status=status.HTTP_200_OK)
+            task_id = data[0]["task"]
+            task_record = Task.objects.get(id=task_id)
+
+            task_configs = TaskConfig.objects.in_bulk(item["task_config"] for item in data)
+            templates = Template.objects.in_bulk(item["template"] for item in data)
+            listen_data_list = [{
+                "task": item["task"],
+                "template": item["template"],
+                "template_name": templates[item["template"]].name,
+                "template_type": templates[item["template"]].type,
+                "template_choice_type": templates[item["template"]].choice_type,
+                "task_config_id": item["task_config"],
+                "url_template": task_configs[item["task_config"]].url_template_id if task_configs[
+                    item["task_config"]].url_template_id else UrlTemplate.objects.filter(template_id=item["template"]).first().id,
+                "key": get_payload(task_configs[item["task_config"]], templates[item["template"]].payload),
+                "task_config_item_list": [{
                     "template_config_item": item["template_config_item"],
                     "id": item["id"],
                     "value": item["value"]
-                }
-                for payload_data in payload_data_list:
-                    if payload_data["task_config_id"] == item["task_config"]:
-                        payload_data["task_config_item_list"].append(task_config_item_list)
-                        break
-                else:
-                    payload_data_list.append({
-                        "task": item["task"],
-                        "template": item["template"],
-                        "template_name": template.name,
-                        "template_type": template.type,
-                        "template_choice_type": template.choice_type,
-                        "task_config_id": item["task_config"],
-                        "key": url,
-                        "task_config_item_list": [task_config_item_list]
-                    })
-
-        result = {
-            "task_info": {
-                "task_id": task_id,
-                "task_name": task_record.name,
-                "callback_url": task_record.callback_url,
-                "callback_url_headers": task_record.callback_url_headers,
-                "show_dashboard": bool(task_record.show_dashboard)},
-            "listen_template_info": listen_data_list,
-            "payload_template_info": payload_data_list,
-        }
-        return Response(result, status=status.HTTP_200_OK)
+                }]
+            } for item in data if templates[item["template"]].type == 1]
+
+            payload_data_list = []
+            for item in data:
+                task_config = task_configs[item["task_config"]]
+                template = templates[item["template"]]
+                if template.type != 1:
+                    url = get_payload(task_config, template.payload)
+                    task_config_item_list = {
+                        "template_config_item": item["template_config_item"],
+                        "id": item["id"],
+                        "value": item["value"]
+                    }
+                    for payload_data in payload_data_list:
+                        if payload_data["task_config_id"] == item["task_config"]:
+                            payload_data["task_config_item_list"].append(task_config_item_list)
+                            break
+                    else:
+                        payload_data_list.append({
+                            "task": item["task"],
+                            "template": item["template"],
+                            "template_name": template.name,
+                            "template_type": template.type,
+                            "template_choice_type": template.choice_type,
+                            "task_config_id": item["task_config"],
+                            "url_template": task_configs[item["task_config"]].url_template_id if task_configs[
+                                item["task_config"]].url_template_id else UrlTemplate.objects.filter(template_id=item["template"]).first().id,
+                            "key": url,
+                            "task_config_item_list": [task_config_item_list]
+                        })
+
+            result = {
+                "task_info": {
+                    "task_id": task_id,
+                    "task_name": task_record.name,
+                    "callback_url": task_record.callback_url,
+                    "callback_url_headers": task_record.callback_url_headers,
+                    "show_dashboard": bool(task_record.show_dashboard)},
+                "listen_template_info": listen_data_list,
+                "payload_template_info": payload_data_list,
+            }
+            return Response(result, status=status.HTTP_200_OK)
+        except Exception as e:
+            print(repr(e))
+            return Response({"code": 0, "message": f"错误原因:{e}"}, status=status.HTTP_200_OK)
 
     def list(self, request, *args, **kwargs):
         queryset = self.filter_queryset(self.get_queryset())
@@ -210,10 +217,10 @@ def create(self, request, *args, **kwargs):
         {
     "task": 1,
     "template": 32,
+    "url_template":1,
     "template_config_item_list": [
         {
             "template_config_item": 52,
-            "id": 7,
             "value": {
                 "ip": "219.137.78.61",
                 "port": 75
@@ -226,10 +233,11 @@ def create(self, request, *args, **kwargs):
         serializer.is_valid(raise_exception=True)
         task_id = serializer.data["task"]
         template_id = serializer.data["template"]
+        url_template_id = serializer.data["url_template"]
         code = generate_code(4)
 
         with transaction.atomic():
-            task_config = TaskConfig.objects.create(task_id=task_id, key=code)
+            task_config = TaskConfig.objects.create(task_id=task_id, key=code, url_template_id=url_template_id)
             template_config_item_list = request.data["template_config_item_list"]
             task_config_items = [
                 TaskConfigItem(
@@ -255,6 +263,7 @@ def update_config(self, request, *args, **kwargs):
     "task": 242,
     "template": 4,
     "task_config":5,
+    "url_template":1,
     "template_config_item_list": [
         {
             "id":1,
@@ -272,6 +281,7 @@ def update_config(self, request, *args, **kwargs):
         template_id = serializer.data["template"]
         task_config_id = serializer.data["task_config"]
         template_config_item_list = request.data["template_config_item_list"]
+        url_template_id = request.data["url_template"]
         with transaction.atomic():
             task_config = TaskConfig.objects.filter(id=task_config_id, task__user_id=self.request.user.id).first()
             if not task_config:
@@ -283,7 +293,7 @@ def update_config(self, request, *args, **kwargs):
             TaskConfig.objects.filter(id=task_config_id).delete()
 
             # 创建新的任务配置
-            new_task_config = TaskConfig.objects.create(task_id=task_id, key=old_key)
+            new_task_config = TaskConfig.objects.create(task_id=task_id, key=old_key, url_template_id=url_template_id)
             new_task_config_id = new_task_config.id
 
             # 创建新的任务配置项
@@ -337,13 +347,11 @@ def api(self, request, *args, **kwargs):
 
         task_config_items = TaskConfigItem.objects.filter(
             task__user=key.user_id, task__status=TASK_STATUS.OPEN, task__show_dashboard=True
-        ).values('template__payload', 'task_config__key')
+        ).values('template__payload', 'task_config')
 
         payload_list = set()
         for task_config_item in task_config_items:
-            payload = get_payload(task_config_item['task_config__key'], task_config_item['template__payload'])
+            payload = get_payload(task_config_item['task_config'], task_config_item['template__payload'])
             if payload:
                 payload_list.add(payload)
         return Response(data={"payload": payload_list}, status=status.HTTP_200_OK)
-
-    from django.db import transaction

modules/template/depend/listen/dnslog.py

@@ -26,6 +26,11 @@
 from modules.template.depend.base import *
 
 
+def get_dns_record():
+    close_old_connections()
+    return DnsConfig.objects.all()
+
+
 class DNSServerFactory(server.DNSServerFactory):
     def handleQuery(self, message, protocol, address):
         # 感谢 suppress\excessive 同学发现python3.6版本运行bug，已fix
@@ -55,7 +60,7 @@ def __init__(self):
         self._peer_address = None
         self.dns_config = {}
         close_old_connections()
-        self.dns_recoed = DnsConfig.objects.all()
+        self.dns_recoed = get_dns_record()
         self.dns_config_domain = [_dns.domain for _dns in self.dns_recoed]
         for _dns in self.dns_recoed:
             self.dns_config[_dns.domain] = cycle(_dns.value)
@@ -92,14 +97,16 @@ def _doDynamicResponse(self, query):
         for domain in self.dns_config_domain:
             print("匹配域名", domain, "匹配结果:", fnmatch.fnmatch(name.decode("utf-8").lower(), domain.lower()))
             if fnmatch.fnmatch(name.decode("utf-8").lower(), domain.lower()):
+                close_old_connections()
                 if len(list(self.dns_recoed.get(domain=domain.lower()).value)) == 1:
                     ttl = 60
                 else:
                     ttl = 0
                 print("ttl:", ttl, flush=True)
                 answers.append(dns.RRHeader(
                     name=name,
-                    payload=dns.Record_A(address=bytes(next(self.dns_config[domain.lower()]), encoding="utf-8")), ttl=ttl))
+                    payload=dns.Record_A(address=bytes(next(self.dns_config[domain.lower()]), encoding="utf-8")),
+                    ttl=ttl))
                 # 存储数据
                 udomain = re.findall(r'\.?([^\.]+)\.%s' % setting.DNS_DOMAIN.strip("*."), name.decode("utf-8").lower())
                 if udomain:

modules/template/depend/listen/jndi.py

@@ -87,7 +87,7 @@ def domessage(self, coon, addr):
                     task_id = task_config_record.task_id
                     username = task_config_record.task.user.username
                     send_email_message(username, remote_addr)
-                    Message.objects.create(domain=self.domain, remote_addr=remote_addr, uri=path,
+                    Message.objects.create(domain=self.domain + "/" + self.key, remote_addr=remote_addr, uri=path,
                                            message_type=MESSAGE_TYPES.LDAP, task_id=task_id, template_id=10)
                     send_message(url=self.domain, remote_addr=remote_addr, uri=path, header='',
                                  message_type=MESSAGE_TYPES.LDAP, content='', task_id=task_id)
@@ -105,7 +105,7 @@ def domessage(self, coon, addr):
                     username = task_config_record.task.user.username
                     send_email_message(username, remote_addr)
                     task_id = task_config_record.task_id
-                    Message.objects.create(domain=self.domain, remote_addr=remote_addr, uri=path,
+                    Message.objects.create(domain=self.domain + "/" + self.key, remote_addr=remote_addr, uri=path,
                                            message_type=MESSAGE_TYPES.RMI, task_id=task_id, template_id=10)
                     send_message(url=self.domain, remote_addr=remote_addr, uri=path, header='',
                                  message_type=MESSAGE_TYPES.RMI, content='', task_id=task_id)

modules/template/migrations/0004_urltemplate.py

@@ -0,0 +1,25 @@
+# Generated by Django 3.2.9 on 2023-04-26 03:48
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('template', '0003_template_code'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='UrlTemplate',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('payload', models.TextField(default='http://{domain}/{key}', help_text='利用代码实例')),
+                ('template', models.ForeignKey(db_constraint=False, help_text='所属模板', on_delete=django.db.models.deletion.CASCADE, related_name='template_url_template', to='template.template')),
+            ],
+            options={
+                'db_table': 'template_url',
+            },
+        ),
+    ]