feat: some bugfix and feature

deploy/commands/management/commands/load_hook_strategy.py

@@ -13,6 +13,7 @@
 from dongtai_common.utils.validate import save_hook_stratefile_sha1sum
 from dongtai_conf.settings import BASE_DIR
 from dongtai_protocol.views.hook_profiles import LANGUAGE_DICT
+from tqdm import tqdm
 
 
 class Command(BaseCommand):

deploy/commands/management/commands/unlock_user.py

@@ -14,3 +14,10 @@ def handle(self, *args, **options):
         users = User.objects.filter(pk__in=options["id"]).all() if options["id"] else User.objects.all()
         users.update(failed_login_count=0)
         self.stdout.write(self.style.SUCCESS("Successfully Unlock Users"))
+        self.stdout.write(self.style.NOTICE("除了使用这个命令外，您还可以通过配置 TOTP 快捷解锁管理员用户。"))
+        self.stdout.write(
+            self.style.NOTICE(
+                "In addition to using this command, "
+                "you can also configure TOTP to quickly unlock the administrator user."
+            )
+        )

dongtai_common/engine/compatibility.py

@@ -24,6 +24,10 @@ def method_pool_3_to_2(dic: dict) -> dict:
         dic["taintPosition"]["target"] = []
     for pv in dic["parameterValues"]:
         pdict[pv["index"]] = pv["value"]
+    if "objValue" not in dic:
+        dic["objValue"] = ""
+    if "retValue" not in dic:
+        dic["retValue"] = ""
     sourceValues = []
     targetValues = []
     for position in dic["taintPosition"]["source"]:

dongtai_common/migrations/0034_auto_20231007_1200.py

@@ -11,7 +11,7 @@ class Migration(migrations.Migration):
     operations = [
         migrations.RunSQL(
             """
-UPDATE auth_user  SET phone = "" WHERE phone IS NULL;
+UPDATE auth_user SET phone = 15000000000 WHERE phone IS NULL;
 ALTER TABLE `auth_user` MODIFY `phone` varchar(255) DEFAULT '' NOT  NULL;
 
             """

dongtai_conf/celery.py

@@ -152,6 +152,11 @@
         "exchange": "dongtai-periodic-task",
         "routing_key": "dongtai-periodic-task",
     },
+    # dongtai-export-report-task 的 /tmp/logstash 目录和 server 共享
+    "dongtai_web.aggr_vul.tasks.update_vul_tantivy_index": {
+        "exchange": "dongtai-export-report-task",
+        "routing_key": "dongtai-export-report-task",
+    },
 }
 configs["CELERY_ENABLE_UTC"] = False
 configs["timezone"] = settings.TIME_ZONE

dongtai_conf/settings.py

@@ -1102,3 +1102,10 @@ def set_asyncio_policy():
 AUTH_LDAP_READY = AUTH_LDAP_SERVER_URI != ""
 # useless
 AUTH_LDAP_USER_DN_TEMPLATE = "uid=%(user)s,ou=users,dc=example,dc=com"
+
+
+# report upload throttle
+REPORT_UPLOAD_THROTTLE = config.get("throttle", "report_upload", fallback="")
+
+# log service timeout
+LOG_SERVICE_TIMEOUT = config.getint("log_service", "port", fallback=10)

dongtai_protocol/report/log_service.py

@@ -1,6 +1,8 @@
 import logging
 import socket
 
+from dongtai_conf.settings import LOG_SERVICE_TIMEOUT
+
 logger = logging.getLogger("dongtai.openapi")
 
 
@@ -16,10 +18,9 @@ def create_socket(self):
             return None
 
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.settimeout(5)
+        sock.settimeout(LOG_SERVICE_TIMEOUT)
         try:
             sock.connect((self.host, self.port))
-            sock.setblocking(False)
             self.socket = sock
         except OSError:
             logger.exception(f"failed to connect log service {self.host}:{self.port}")
@@ -39,7 +40,7 @@ def send(self, message):
             if not self.socket:
                 self.create_socket()
             if self.socket:
-                self.socket.sendall(bytes(message + "\n", encoding="utf-8"), socket.MSG_DONTWAIT)
+                self.socket.sendall(bytes(message + "\n", encoding="utf-8"))
                 return True
         except Exception as e:
             logger.exception("failed to send message to log service", exc_info=e)

dongtai_protocol/views/agent_register.py

@@ -146,7 +146,7 @@ def register_server(
         try:
             port = int(server_port)
         except Exception:
-            logger.exception(_("The server port does not exist, has been set to the default: 0"))
+            logger.info(_("The server port does not exist, has been set to the default: 0"))
             port = 0
 
         server_id = agent.server_id

dongtai_protocol/views/report_upload.py

@@ -4,18 +4,31 @@
 import logging
 
 from drf_spectacular.utils import extend_schema
+from rest_framework.throttling import SimpleRateThrottle
 
 from dongtai_common.endpoint import OpenApiEndPoint, R
+from dongtai_conf.settings import REPORT_UPLOAD_THROTTLE
 from dongtai_protocol.decrypter import parse_data
 from dongtai_protocol.report.report_handler_factory import ReportHandler
 
 logger = logging.getLogger("dongtai.openapi")
 
 
+class CustomRateThrottle(SimpleRateThrottle):
+    scope = "report_upload"
+    rate = REPORT_UPLOAD_THROTTLE
+
+    def get_cache_key(self, request, view):
+        return self.cache_format % {"scope": self.scope, "ident": self.get_ident(request)}
+
+
 class ReportUploadEndPoint(OpenApiEndPoint):
     name = "api-v1-report-upload"
     description = "agent上传报告"
 
+    if REPORT_UPLOAD_THROTTLE:
+        throttle_classes = [CustomRateThrottle]
+
     @extend_schema(
         summary="Agent 上传报告",
         tags=["Agent服务端交互协议"],

dongtai_web/aggr_vul/tasks.py

@@ -22,7 +22,7 @@ def update_vul_tantivy_index_receiver(sender, instance, **kwargs):
 
 
 @shared_task(
-    queue="dongtai-periodic-task",
+    queue="dongtai-export-report-task",
     base=Singleton,
     lock_expiry=20,
 )

dongtai_web/views/utils/commonstats.py

@@ -210,7 +210,7 @@ def get_summary_by_project(project_id: int, project_version_id: int):
         else:
             day_num_dict[i["day_label"]] = [i]
     day_num_data = []
-    last_timestamp: int = 0
+    last_timestamp: int = current_timestamp - 60 * 60 * 24 * 7
     for day_label_i in range(len(daylist)):
         timestamp, day_label = daylist[day_label_i]
         if day_label in day_num_dict: