diff --git a/deploy/commands/management/commands/unlock_user.py b/deploy/commands/management/commands/unlock_user.py index c00ea6c5..cd94193f 100644 --- a/deploy/commands/management/commands/unlock_user.py +++ b/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." + ) + ) diff --git a/dongtai_common/migrations/0034_auto_20231007_1200.py b/dongtai_common/migrations/0034_auto_20231007_1200.py index 76d16039..93cdbca6 100644 --- a/dongtai_common/migrations/0034_auto_20231007_1200.py +++ b/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; """ diff --git a/dongtai_conf/celery.py b/dongtai_conf/celery.py index 9924a232..e2b84d9a 100644 --- a/dongtai_conf/celery.py +++ b/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 diff --git a/dongtai_conf/settings.py b/dongtai_conf/settings.py index eb60d1ca..a8bdb71e 100644 --- a/dongtai_conf/settings.py +++ b/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) diff --git a/dongtai_protocol/report/log_service.py b/dongtai_protocol/report/log_service.py index a0bc2877..7d8e6d79 100644 --- a/dongtai_protocol/report/log_service.py +++ b/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) diff --git a/dongtai_protocol/views/agent_register.py b/dongtai_protocol/views/agent_register.py index d81e3ad3..fe4fafda 100644 --- a/dongtai_protocol/views/agent_register.py +++ b/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 diff --git a/dongtai_protocol/views/report_upload.py b/dongtai_protocol/views/report_upload.py index 79b2aec7..84f635fb 100644 --- a/dongtai_protocol/views/report_upload.py +++ b/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服务端交互协议"], diff --git a/dongtai_web/aggr_vul/tasks.py b/dongtai_web/aggr_vul/tasks.py index 1a561a88..3d696415 100644 --- a/dongtai_web/aggr_vul/tasks.py +++ b/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, ) diff --git a/dongtai_web/views/utils/commonstats.py b/dongtai_web/views/utils/commonstats.py index 6f721b91..394b91ee 100644 --- a/dongtai_web/views/utils/commonstats.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index ae0a9afe..97ec46ae 100644 --- a/pyproject.toml +++ b/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"]