Merge pull request #121 from wuba/develop

Develop

modules/account/views.py

@@ -97,7 +97,6 @@ class UserViewSet(mixins.ListModelMixin, mixins.UpdateModelMixin, GenericViewSet
     permission_classes = (IsAdminUser,)
 
     @action(methods=["POST"], detail=False, permission_classes=[AllowAny])
-    @transaction.atomic()
     def register(self, request, *args, **kwargs):
         """
         注册用户
@@ -107,13 +106,14 @@ def register(self, request, *args, **kwargs):
         serializer.is_valid(raise_exception=True)
         username = request.data["username"]
         password = request.data["password"]
-        user = User.objects.create_user(username=username, password=password)
-        invite_code = request.data.get("invite_code", "")
-        if setting.REGISTER_TYPE == REGISTER_TYPE.INVITE:  # 判断是开放邀请注册
-            InviteCode.objects.filter(code=invite_code).delete()
         apikey = generate_code(32)
-        ApiKey.objects.create(user=user, key=apikey)
-        TaskCreator().create_initial_task(user.id)
+        invite_code = request.data.get("invite_code", "")
+        with transaction.atomic():
+            user = User.objects.create_user(username=username, password=password)
+            if setting.REGISTER_TYPE == REGISTER_TYPE.INVITE:  # 判断是开放邀请注册
+                InviteCode.objects.filter(code=invite_code).delete()
+            ApiKey.objects.create(user=user, key=apikey)
+            TaskCreator().create_initial_task(user.id)
         response_data = {
             "username": username,
             "apikey": apikey

modules/message/views.py

@@ -6,24 +6,21 @@
 from django.db.models.functions import Cast
 from django.db.models import DateField
 from modules.api.models import ApiKey
-from modules.message.models import Message
 from modules.message.serializers import MessageFilter, MessageSerializer
-from modules.task.models import Task, TaskConfigItem
+from modules.task.models import Task
 from modules.template.choose_template import match_template
 from modules.template.constants import PRIVATE_TYPES
-from modules.message.constants import MESSAGE_TYPES
 from modules.template.models import Template
 from rest_framework import filters, mixins, status
 from rest_framework.decorators import action
 from rest_framework.permissions import AllowAny, IsAuthenticated
 from rest_framework.response import Response
 from rest_framework.viewsets import GenericViewSet
 from django.db.models import Avg, Max, Min, Count, Sum, Q
-from utils.helper import get_payload, send_message, get_message_type_name, reconstruct_request, get_param_message, \
-    send_email_message
+from utils.helper import get_payload, get_message_type_name, reconstruct_request, get_param_message
 from modules.task.constants import TASK_TMP
-from modules.config import setting
 from utils.helper import is_base64
+from modules.template.depend.base import *
 
 
 class MessageView(GenericViewSet, mixins.ListModelMixin, mixins.DestroyModelMixin):
@@ -182,7 +179,6 @@ def index(request):
     message = is_base64(message_to_base64)
     host = request.get_host()  # 项目域名
     domain_key = host.split('.')[0]
-    url = host + '/' + path
     remote_addr = request.META.get('REMOTE_ADDR', '')  # 请求ip
     regex = re.compile('^HTTP_')
     headers = dict((regex.sub('', header), value) for (header, value) in request.META.items() if
@@ -197,26 +193,13 @@ def index(request):
                 template_response = match_template(task_config_item, param_list)
                 return template_response
             else:
-                Message.objects.create(domain=url, remote_addr=remote_addr, uri=path, header=headers,
-                                       message_type=MESSAGE_TYPES.HTTP, content=message,
-                                       task_id=task_config_item.task_id, html=raw_response,
-                                       template_id=task_config_item.template_id)
-                send_message(url=url, remote_addr=remote_addr, uri=path, header=headers,
-                             message_type=MESSAGE_TYPES.HTTP, content=message, task_id=task_config_item.task_id,
-                             raw=raw_response)
-
+                message_callback(domain=host, remote_addr=remote_addr, task_config_item=task_config_item, uri=path,
+                                 header=headers, message_type=MESSAGE_TYPES.HTTP, content=message, raw=raw_response)
     # http 请求日志
     elif len(domain_key) == 4 and domain_key != setting.PLATFORM_DOMAIN.split('.')[0]:
         task_config_item = TaskConfigItem.objects.filter(task_config__key__iexact=domain_key,
                                                          task__status=1).first()
         if task_config_item and task_config_item.template.name == "HTTP":
-            username = task_config_item.task.user.username
-            send_email_message(username, remote_addr)
-            Message.objects.create(domain=host, remote_addr=remote_addr, uri=path, header=headers,
-                                   message_type=MESSAGE_TYPES.HTTP, content=message,
-                                   task_id=task_config_item.task_id,
-                                   template_id=task_config_item.template_id, html=raw_response, )
-            send_message(url=host, remote_addr=remote_addr, uri=path, header=headers,
-                         message_type=MESSAGE_TYPES.HTTP, content=message, task_id=task_config_item.task_id,
-                         raw=raw_response)
+            message_callback(domain=host, remote_addr=remote_addr, task_config_item=task_config_item, uri=path,
+                             header=headers, message_type=MESSAGE_TYPES.HTTP, content=message, raw=raw_response)
     return HttpResponse('', content_type='text/html;charset=utf-8')

modules/template/depend/base.py

@@ -3,9 +3,20 @@
 from rest_framework.response import Response
 from modules.config.models import Config
 from modules.config import setting
+from modules.message.models import Message
 from modules.task.models import TaskConfigItem
 import requests
 import json
+from utils.helper import send_email_message, send_message
+from modules.task.constants import TASK_STATUS
+from modules.message.constants import MESSAGE_TYPES
+from django.db import connections, transaction
+
+
+def close_old_connections():
+    for conn in connections.all():
+        conn.close_if_unusable_or_obsolete()
+
 
 class BaseTemplate:
     info = [{
@@ -58,3 +69,41 @@ def replace_code(self, code=""):
     @abstractmethod
     def generate(self, key, config):
         pass
+
+@transaction.non_atomic_requests
+def hit(key, template_name, iexact=False):
+    """
+    命中
+    """
+    close_old_connections()
+    found = False
+    if iexact:
+        found_item = TaskConfigItem.objects.filter(
+            task_config__key__iexact=key, task__status=TASK_STATUS.OPEN
+        ).first()
+    else:
+        found_item = TaskConfigItem.objects.filter(
+            task_config__key=key, task__status=TASK_STATUS.OPEN
+        ).first()
+
+    if found_item and found_item.template.name in template_name:
+        found = True
+
+    return found, found_item
+
+
+def message_callback(domain, remote_addr, task_config_item, uri, header, message_type, content, raw=""):
+    """
+    命中回调
+    """
+    try:
+        Message.objects.create(domain=domain, message_type=message_type,
+                               remote_addr=remote_addr,
+                               task_id=task_config_item.task_id,
+                               template_id=task_config_item.template_id, uri=uri, header=header, content=content,
+                               html=raw)
+        send_email_message(task_config_item.task.user.username, remote_addr)
+        send_message(url=domain, remote_addr=remote_addr, uri=uri, header=header,
+                     message_type=message_type, content=content, task_id=task_config_item.task_id, raw=raw)
+    except Exception as e:
+        print("message_callback error: %s" % e)

modules/template/depend/listen/dnslog.py

@@ -12,22 +12,18 @@
 import django
 from itertools import cycle
 
+from twisted.internet import reactor
+from twisted.names import dns, server
 from twisted.names.dns import DNSDatagramProtocol
 
 PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__) + "../../../../../")
 sys.path.append(PROJECT_ROOT)
 os.environ['DJANGO_SETTINGS_MODULE'] = 'antenna.settings'
 django.setup()
 
-from modules.config import setting
-from twisted.internet import reactor
-from twisted.names import dns, server
-from utils.helper import send_message, send_email_message
 from modules.message.constants import MESSAGE_TYPES
 from modules.config.models import DnsConfig
-from modules.message.models import Message
-from modules.task.models import TaskConfig, TaskConfigItem
-from modules.template.depend.base import BaseTemplate
+from modules.template.depend.base import *
 
 
 class DNSServerFactory(server.DNSServerFactory):
@@ -58,9 +54,10 @@ def __init__(self):
         # 初始化变量
         self._peer_address = None
         self.dns_config = {}
-        dns_recoed = DnsConfig.objects.all()
-        self.dns_config_domain = [_dns.domain for _dns in dns_recoed]
-        for _dns in dns_recoed:
+        close_old_connections()
+        self.dns_recoed = DnsConfig.objects.all()
+        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)
 
     @property
@@ -93,26 +90,24 @@ def _doDynamicResponse(self, query):
         self.dns_config_domain.sort(key=lambda x: x.startswith("*"))  # 进行排序
         print("请求解析域名：", name.decode("utf-8"), flush=True)
         for domain in self.dns_config_domain:
-            print("匹配域名", domain, "匹配结果:", fnmatch.fnmatch(name.decode("utf-8"), domain))
+            print("匹配域名", domain, "匹配结果:", fnmatch.fnmatch(name.decode("utf-8").lower(), domain.lower()))
             if fnmatch.fnmatch(name.decode("utf-8").lower(), domain.lower()):
+                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]), encoding="utf-8")), ttl=0))
+                    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:
-                    task_config_item = TaskConfigItem.objects.filter(task_config__key__iexact=udomain[0],
-                                                                     task__status=1).first()
-                    if task_config_item and task_config_item.template.name == "DNS":
-                        username = task_config_item.task.user.username
-                        Message.objects.create(domain=name.decode("utf-8"), message_type=MESSAGE_TYPES.DNS,
-                                               remote_addr=addr,
-                                               task_id=task_config_item.task_id,
-                                               template_id=task_config_item.template_id)
-                        print("完成保存")
-                        send_email_message(username, addr)
-                        send_message(url=name.decode("utf-8"), remote_addr=addr, uri='', header='',
-                                     message_type=MESSAGE_TYPES.DNS, content='', task_id=task_config_item.task_id)
+                    flag, task_config_item = hit(udomain[0], template_name=["DNS"], iexact=True)  # 不区分大小写
+                    if flag:
+                        message_callback(domain=name.decode("utf-8"), remote_addr=addr,
+                                         task_config_item=task_config_item, uri='',
+                                         header='', message_type=MESSAGE_TYPES.DNS, content='')  # 命中回调
                 break
         authority = []
         additional = []
@@ -165,4 +160,4 @@ def __init__(self):
 
 
 if __name__ == '__main__':
-    raise SystemExit(main())
+    main()

modules/template/depend/listen/ftplog.py

@@ -12,10 +12,7 @@
 django.setup()
 
 from modules.message.constants import MESSAGE_TYPES
-from modules.message.models import Message
-from modules.task.models import TaskConfigItem
-from modules.template.depend.base import BaseTemplate
-from utils.helper import send_message, send_email_message
+from modules.template.depend.base import *
 from modules.config import setting
 
 WELCOME_MSG = b'220 (vsFTPd 2.0.5) '
@@ -43,19 +40,11 @@ def connectionMade(self):
 
     def connectionLost(self, line):
         self.sendLine(GOODBYE_MSG)
-        task_config_item = TaskConfigItem.objects.filter(task_config__key=self.key,
-                                                         task__status=1).first()
-        if task_config_item and (
-            task_config_item.template.name == "FTP" or task_config_item.template.name == "XXE"):
-            username = task_config_item.task.user.username
-            send_email_message(username, self.remote_addr)
-            Message.objects.create(domain=setting.PLATFORM_DOMAIN, message_type=MESSAGE_TYPES.FTP,
-                                   remote_addr=self.remote_addr,
-                                   task_id=task_config_item.task_id, template_id=task_config_item.template_id,
-                                   content=self.content)
-            send_message(url=setting.PLATFORM_DOMAIN, remote_addr=self.remote_addr, uri='', header='',
-                         message_type=MESSAGE_TYPES.HTTP, content=self.content,
-                         task_id=task_config_item.task_id)
+        flag, task_config_item = hit(self.key, template_name=["FTP", "XXE"], iexact=False)
+        if flag:
+            message_callback(domain=setting.PLATFORM_DOMAIN, remote_addr=self.remote_addr,
+                             task_config_item=task_config_item, uri='', header='', message_type=MESSAGE_TYPES.FTP,
+                             content=self.content)
 
     def lineReceived(self, line):
         if self.state == 'get_name':

modules/template/depend/listen/httplog.py

@@ -11,7 +11,7 @@
 os.environ['DJANGO_SETTINGS_MODULE'] = 'antenna.settings'
 django.setup()
 
-from modules.template.depend.base import BaseTemplate
+from modules.template.depend.base import *
 from modules.message.constants import MESSAGE_TYPES
 from modules.message.models import Message
 from modules.task.models import TaskConfigItem
@@ -68,20 +68,10 @@ def is_payload(self, path):
                              message_type=MESSAGE_TYPES.HTTP, content=self.html, task_id=task_config_item.task_id)
 
     def connectionLost(self, reason):
-        task_config_item = TaskConfigItem.objects.filter(task_config__key=self.key,
-                                                         task__status=1).first()
-        if task_config_item and task_config_item.template.name == "HTTP":
-            username = task_config_item.task.user.username
-            send_email_message(username, self.remote_addr)
-            Message.objects.create(domain=self.domain + '/' + self.uri, message_type=MESSAGE_TYPES.HTTPS,
-                                   remote_addr=self.remote_addr,
-                                   task_id=task_config_item.task_id,
-                                   uri=self.uri,
-                                   template_id=task_config_item.template_id,
-                                   content=self.content)
-            send_message(url=self.domain + '/' + self.uri, remote_addr=self.remote_addr, uri=self.uri, header='',
-                         message_type=MESSAGE_TYPES.HTTPS, content=self.content,
-                         task_id=task_config_item.task_id)
+        flag, task_config_item = hit(self.key, template_name=["HTTP"], iexact=False)
+        if flag:
+            message_callback(domain=self.domain, remote_addr=self.remote_addr, task_config_item=task_config_item,
+                             uri=self.uri, header='', message_type=MESSAGE_TYPES.HTTP, content=self.content)
 
 
 class HttpTemplate(BaseTemplate):