Merge pull request #1915 from st1020/feat/some-bugfix-and-feature

feat: some bugfix and feature

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/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:

pyproject.toml

@@ -47,6 +47,7 @@ ignore = [
     "TRY003",  # Avoid specifying long messages outside the exception class
     "TRY301",  # Abstract raise to an inner function
     "PERF203", # try-except within a loop incurs performance overhead
+    "RUF001",  # String contains ambiguous
 ]
 # target version is python 3.9 because Cython not support some python 3.10 feature
 target-version = "py39"
@@ -56,3 +57,6 @@ target-version = "py39"
 "**/tests.py" = ["T201"]
 "deploy/*.py" = ["T201"]
 "dongtai_conf/*.py" = ["T201"]
+
+[tool.pyright]
+ignore = [".venv", "**/urls.py"]