From 50384d599f61032f8a39e3fb2d14a6ada4b0271d Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Wed, 11 Oct 2023 14:06:42 +0800 Subject: [PATCH 1/8] feat: new metric handler. --- .../report/handler/metric_handler.py | 227 ++++++++++++++++++ 1 file changed, 227 insertions(+) create mode 100644 dongtai_protocol/report/handler/metric_handler.py diff --git a/dongtai_protocol/report/handler/metric_handler.py b/dongtai_protocol/report/handler/metric_handler.py new file mode 100644 index 00000000..c63f85a2 --- /dev/null +++ b/dongtai_protocol/report/handler/metric_handler.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python +# datetime:2020/10/23 11:56 +import logging +import time +from typing import Any + +from celery import shared_task +from celery_singleton import Singleton +from django.core.cache import cache +from django.db.models import Q, QuerySet +from django.utils.translation import gettext_lazy as _ + +from dongtai_common.models.agent import IastAgent +from dongtai_common.models.heartbeat import IastHeartbeat +from dongtai_common.models.project import VulValidation +from dongtai_common.models.replay_queue import IastReplayQueue +from dongtai_common.utils import const +from dongtai_common.utils.systemsettings import get_vul_validate +from dongtai_protocol.report.handler.report_handler_interface import IReportHandler +from dongtai_protocol.report.report_handler_factory import ReportHandler +from dongtai_web.vul_log.vul_log import log_recheck_vul + +logger = logging.getLogger("dongtai.openapi") + + +def update_agent_cache(agent_id, data): + cache.set(f"heartbeat-{agent_id}", data, timeout=521) + + +def check_agent_incache(agent_id): + return bool(cache.get(f"heartbeat-{agent_id}")) + + +@shared_task(base=Singleton, unique_on=["agent_id"], lock_expiry=20) +def update_heartbeat(agent_id: int, defaults: dict[str, Any]): + IastHeartbeat.objects.update_or_create(agent_id=agent_id, defaults=defaults) + IastAgent.objects.update_or_create(pk=agent_id, defaults={"is_running": 1, "online": 1}) + + +@ReportHandler.register(const.REPORT_HEART_BEAT) +class HeartBeatHandler(IReportHandler): + def __init__(self): + super().__init__() + self.req_count = None + self.cpu = None + self.memory = None + self.network = None + self.report_queue = None + self.method_queue = None + self.replay_queue = None + self.return_queue = None + + def parse(self): + self.cpu = self.detail.get("cpu") + self.memory = self.detail.get("memory") + self.disk = self.detail.get("disk") + self.req_count = self.detail.get("reqCount", None) + self.report_queue = self.detail.get("reportQueue", 0) + self.method_queue = self.detail.get("methodQueue", 0) + self.replay_queue = self.detail.get("replayQueue", 0) + self.return_queue = self.detail.get("returnQueue", None) + + def has_permission(self): + self.agent = IastAgent.objects.filter(id=self.agent_id).first() + return self.agent + + def save_heartbeat(self): + default_dict = {"dt": int(time.time())} + if self.return_queue == 1: + default_dict["req_count"] = self.req_count + default_dict["report_queue"] = self.report_queue + default_dict["method_queue"] = self.method_queue + default_dict["replay_queue"] = self.replay_queue + elif self.return_queue == 0: + if self.req_count is not None: + default_dict["req_count"] = self.req_count + default_dict["memory"] = self.memory + default_dict["cpu"] = self.cpu + default_dict["disk"] = self.disk + update_heartbeat.delay(agent_id=self.agent_id, defaults=default_dict) + else: + default_dict["memory"] = self.memory + default_dict["cpu"] = self.cpu + default_dict["req_count"] = self.req_count + default_dict["report_queue"] = self.report_queue + default_dict["method_queue"] = self.method_queue + default_dict["replay_queue"] = self.replay_queue + default_dict["disk"] = self.disk + update_heartbeat.delay(agent_id=self.agent_id, defaults=default_dict) + update_agent_cache(self.agent_id, default_dict) + + def get_result(self, msg=None): + logger.info(f"return_queue: {self.return_queue}") + if (self.return_queue is None or self.return_queue == 1) and vul_recheck_state(self.agent_id): + try: + project_agents = ( + IastAgent.objects.values_list("id", flat=True) + .filter( + bind_project_id=self.agent.bind_project_id, + language=self.agent.language, + ) + .union( + addtional_agenti_ids_query_filepath_simhash( + self.agent.filepathsimhash, language=self.agent.language + ), + addtional_agent_ids_query_deployway_and_path( + self.agent.servicetype, + self.agent.server.path, + self.agent.server.hostname, + language=self.agent.language, + ), + ) + ) + project_agents = list(project_agents) + if project_agents is None: + logger.info(_("There is no probe under the project")) + logger.info(f"project_agent_ids : {project_agents}") + replay_queryset = IastReplayQueue.objects.values( + "id", + "relation_id", + "uri", + "method", + "scheme", + "header", + "params", + "body", + "replay_type", + ).filter(agent_id__in=project_agents, state__in=[const.WAITING, const.SOLVING],)[:200] + if len(replay_queryset) == 0: + logger.info(_("Replay request does not exist")) + + ( + success_ids, + success_vul_ids, + failure_ids, + failure_vul_ids, + replay_requests, + ) = ([], [], [], [], []) + for replay_request in replay_queryset: + if replay_request["uri"]: + replay_requests.append(replay_request) + success_ids.append(replay_request["id"]) + if replay_request["replay_type"] == const.VUL_REPLAY: + success_vul_ids.append(replay_request["relation_id"]) + else: + failure_ids.append(replay_request["id"]) + if replay_request["replay_type"] == const.VUL_REPLAY: + failure_vul_ids.append(replay_request["relation_id"]) + + timestamp = int(time.time()) + IastReplayQueue.objects.filter(id__in=success_ids, state=const.SOLVING).update( + update_time=timestamp, state=const.SOLVED + ) + IastReplayQueue.objects.filter(id__in=success_ids, state=const.WAITING).update( + update_time=timestamp, state=const.SOLVING + ) + IastReplayQueue.objects.filter(id__in=failure_ids).update(update_time=timestamp, state=const.SOLVED) + + log_recheck_vul( + self.agent.user.id, + self.agent.user.username, + success_vul_ids, + "验证中", + ) + logger.info(_("Reproduction request issued successfully")) + logger.debug([i["id"] for i in replay_requests]) + except Exception as e: + logger.info( + _("Replay request query failed, reason: {}").format(e), + exc_info=True, + ) + else: + return replay_requests + + return [] + + def save(self): + self.save_heartbeat() + + def get_agent(self, agent_id): + return IastAgent.objects.filter(id=agent_id).first() + + +def get_k8s_deployment_id(hostname: str) -> str: + return hostname[hostname.rindex("-")] + + +def addtional_agent_ids_query_deployway_and_path(deployway: str, path: str, hostname: str, language: str) -> QuerySet: + if deployway == "k8s": + deployment_id = get_k8s_deployment_id(hostname) + logger.info(f"deployment_id : {deployment_id}") + server_q = ( + Q(server__hostname__startswith=deployment_id) + & Q(server__path=path) + & Q(server__path="") + & ~Q(server__hostname="") + ) + elif deployway == "docker": + server_q = Q(server__path=path) & ~Q(server__path="") + else: + server_q = ( + Q(server__path=str(path)) + & Q(server__hostname=str(hostname)) + & ~Q(server__path="") + & ~Q(server__hostname="") + ) + final_q = server_q & Q(language=language) + return IastAgent.objects.filter(final_q).values_list("id", flat=True) + + +def addtional_agenti_ids_query_filepath_simhash(filepathsimhash: str, language: str) -> QuerySet: + return IastAgent.objects.filter(filepathsimhash=filepathsimhash, language=language).values_list("id", flat=True) + + +def get_project_vul_validation_state(agent_id): + state = IastAgent.objects.filter(pk=agent_id).values_list("bind_project__vul_validation", flat=True).first() + if state is None: + state = VulValidation.FOLLOW_GLOBAL + return state + + +def vul_recheck_state(agent_id): + project_level_validation = get_project_vul_validation_state(agent_id) + global_state = get_vul_validate() + if project_level_validation == VulValidation.FOLLOW_GLOBAL: + return global_state + return project_level_validation == VulValidation.ENABLE From 9a35f3fba5a6056a74595c2f89b33128684630ce Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Fri, 13 Oct 2023 11:03:40 +0800 Subject: [PATCH 2/8] feat: new metric handler. --- dongtai_common/models/agent.py | 12 +- dongtai_common/utils/const.py | 1 + .../report/handler/metric_handler.py | 188 ++---------------- dongtai_protocol/views/agent_register.py | 2 +- 4 files changed, 26 insertions(+), 177 deletions(-) diff --git a/dongtai_common/models/agent.py b/dongtai_common/models/agent.py index b5ae22a6..3efab79c 100644 --- a/dongtai_common/models/agent.py +++ b/dongtai_common/models/agent.py @@ -50,6 +50,7 @@ class IastAgent(models.Model): events = models.JSONField(default=get_events) department = models.ForeignKey(Department, models.DO_NOTHING) allow_report = models.IntegerField(default=1) + jvm_user_dir = models.CharField(max_length=1024, default="") class Meta: managed = get_managed() @@ -91,6 +92,11 @@ class Meta: db_table = "iast_agent_event" -# class IastAgent(models.Model): -# -# class Meta: +class IastAgentDiskList(models.Model): + agent = models.ForeignKey(IastAgent, on_delete=models.CASCADE) + create_at = models.DateTimeField(blank=True, null=True, auto_now_add=True) + data = models.JSONField(default=dict) + + class Meta: + managed = get_managed() + db_table = "iast_agent_disk_list" diff --git a/dongtai_common/utils/const.py b/dongtai_common/utils/const.py index d553cd10..b819569d 100644 --- a/dongtai_common/utils/const.py +++ b/dongtai_common/utils/const.py @@ -3,6 +3,7 @@ # report REPORT_HEART_BEAT = 0x01 +REPORT_METRIC = 0x02 REPORT_SCA = 0x11 REPORT_VULN_NORNAL = 0x21 REPORT_VULN_DYNAMIC = 0x22 diff --git a/dongtai_protocol/report/handler/metric_handler.py b/dongtai_protocol/report/handler/metric_handler.py index c63f85a2..36e243c7 100644 --- a/dongtai_protocol/report/handler/metric_handler.py +++ b/dongtai_protocol/report/handler/metric_handler.py @@ -10,7 +10,10 @@ from django.db.models import Q, QuerySet from django.utils.translation import gettext_lazy as _ -from dongtai_common.models.agent import IastAgent +from dongtai_common.models.agent import ( + IastAgent, + IastAgentDiskList, +) from dongtai_common.models.heartbeat import IastHeartbeat from dongtai_common.models.project import VulValidation from dongtai_common.models.replay_queue import IastReplayQueue @@ -19,26 +22,14 @@ from dongtai_protocol.report.handler.report_handler_interface import IReportHandler from dongtai_protocol.report.report_handler_factory import ReportHandler from dongtai_web.vul_log.vul_log import log_recheck_vul +from dongtai_protocol.report.handler.heartbeat_handler import update_heartbeat +from dongtai_common.models.agent import IastAgentDiskList logger = logging.getLogger("dongtai.openapi") -def update_agent_cache(agent_id, data): - cache.set(f"heartbeat-{agent_id}", data, timeout=521) - - -def check_agent_incache(agent_id): - return bool(cache.get(f"heartbeat-{agent_id}")) - - -@shared_task(base=Singleton, unique_on=["agent_id"], lock_expiry=20) -def update_heartbeat(agent_id: int, defaults: dict[str, Any]): - IastHeartbeat.objects.update_or_create(agent_id=agent_id, defaults=defaults) - IastAgent.objects.update_or_create(pk=agent_id, defaults={"is_running": 1, "online": 1}) - - -@ReportHandler.register(const.REPORT_HEART_BEAT) -class HeartBeatHandler(IReportHandler): +@ReportHandler.register(const.REPORT_METRIC) +class ReportMetricHandler(IReportHandler): def __init__(self): super().__init__() self.req_count = None @@ -53,12 +44,8 @@ def __init__(self): def parse(self): self.cpu = self.detail.get("cpu") self.memory = self.detail.get("memory") - self.disk = self.detail.get("disk") - self.req_count = self.detail.get("reqCount", None) - self.report_queue = self.detail.get("reportQueue", 0) - self.method_queue = self.detail.get("methodQueue", 0) - self.replay_queue = self.detail.get("replayQueue", 0) - self.return_queue = self.detail.get("returnQueue", None) + self.disk = self.detail.get("disk") # extend now + self.disk_list = self.detail.get("diskList") def has_permission(self): self.agent = IastAgent.objects.filter(id=self.agent_id).first() @@ -66,112 +53,13 @@ def has_permission(self): def save_heartbeat(self): default_dict = {"dt": int(time.time())} - if self.return_queue == 1: - default_dict["req_count"] = self.req_count - default_dict["report_queue"] = self.report_queue - default_dict["method_queue"] = self.method_queue - default_dict["replay_queue"] = self.replay_queue - elif self.return_queue == 0: - if self.req_count is not None: - default_dict["req_count"] = self.req_count - default_dict["memory"] = self.memory - default_dict["cpu"] = self.cpu - default_dict["disk"] = self.disk - update_heartbeat.delay(agent_id=self.agent_id, defaults=default_dict) - else: - default_dict["memory"] = self.memory - default_dict["cpu"] = self.cpu - default_dict["req_count"] = self.req_count - default_dict["report_queue"] = self.report_queue - default_dict["method_queue"] = self.method_queue - default_dict["replay_queue"] = self.replay_queue - default_dict["disk"] = self.disk - update_heartbeat.delay(agent_id=self.agent_id, defaults=default_dict) - update_agent_cache(self.agent_id, default_dict) + default_dict["memory"] = self.memory + default_dict["cpu"] = self.cpu + default_dict["disk"] = self.disk + update_heartbeat.delay(agent_id=self.agent_id, defaults=default_dict) + IastAgentDiskList.objects.create(agent=self.agent_id, data=self.disk_list) def get_result(self, msg=None): - logger.info(f"return_queue: {self.return_queue}") - if (self.return_queue is None or self.return_queue == 1) and vul_recheck_state(self.agent_id): - try: - project_agents = ( - IastAgent.objects.values_list("id", flat=True) - .filter( - bind_project_id=self.agent.bind_project_id, - language=self.agent.language, - ) - .union( - addtional_agenti_ids_query_filepath_simhash( - self.agent.filepathsimhash, language=self.agent.language - ), - addtional_agent_ids_query_deployway_and_path( - self.agent.servicetype, - self.agent.server.path, - self.agent.server.hostname, - language=self.agent.language, - ), - ) - ) - project_agents = list(project_agents) - if project_agents is None: - logger.info(_("There is no probe under the project")) - logger.info(f"project_agent_ids : {project_agents}") - replay_queryset = IastReplayQueue.objects.values( - "id", - "relation_id", - "uri", - "method", - "scheme", - "header", - "params", - "body", - "replay_type", - ).filter(agent_id__in=project_agents, state__in=[const.WAITING, const.SOLVING],)[:200] - if len(replay_queryset) == 0: - logger.info(_("Replay request does not exist")) - - ( - success_ids, - success_vul_ids, - failure_ids, - failure_vul_ids, - replay_requests, - ) = ([], [], [], [], []) - for replay_request in replay_queryset: - if replay_request["uri"]: - replay_requests.append(replay_request) - success_ids.append(replay_request["id"]) - if replay_request["replay_type"] == const.VUL_REPLAY: - success_vul_ids.append(replay_request["relation_id"]) - else: - failure_ids.append(replay_request["id"]) - if replay_request["replay_type"] == const.VUL_REPLAY: - failure_vul_ids.append(replay_request["relation_id"]) - - timestamp = int(time.time()) - IastReplayQueue.objects.filter(id__in=success_ids, state=const.SOLVING).update( - update_time=timestamp, state=const.SOLVED - ) - IastReplayQueue.objects.filter(id__in=success_ids, state=const.WAITING).update( - update_time=timestamp, state=const.SOLVING - ) - IastReplayQueue.objects.filter(id__in=failure_ids).update(update_time=timestamp, state=const.SOLVED) - - log_recheck_vul( - self.agent.user.id, - self.agent.user.username, - success_vul_ids, - "验证中", - ) - logger.info(_("Reproduction request issued successfully")) - logger.debug([i["id"] for i in replay_requests]) - except Exception as e: - logger.info( - _("Replay request query failed, reason: {}").format(e), - exc_info=True, - ) - else: - return replay_requests - return [] def save(self): @@ -179,49 +67,3 @@ def save(self): def get_agent(self, agent_id): return IastAgent.objects.filter(id=agent_id).first() - - -def get_k8s_deployment_id(hostname: str) -> str: - return hostname[hostname.rindex("-")] - - -def addtional_agent_ids_query_deployway_and_path(deployway: str, path: str, hostname: str, language: str) -> QuerySet: - if deployway == "k8s": - deployment_id = get_k8s_deployment_id(hostname) - logger.info(f"deployment_id : {deployment_id}") - server_q = ( - Q(server__hostname__startswith=deployment_id) - & Q(server__path=path) - & Q(server__path="") - & ~Q(server__hostname="") - ) - elif deployway == "docker": - server_q = Q(server__path=path) & ~Q(server__path="") - else: - server_q = ( - Q(server__path=str(path)) - & Q(server__hostname=str(hostname)) - & ~Q(server__path="") - & ~Q(server__hostname="") - ) - final_q = server_q & Q(language=language) - return IastAgent.objects.filter(final_q).values_list("id", flat=True) - - -def addtional_agenti_ids_query_filepath_simhash(filepathsimhash: str, language: str) -> QuerySet: - return IastAgent.objects.filter(filepathsimhash=filepathsimhash, language=language).values_list("id", flat=True) - - -def get_project_vul_validation_state(agent_id): - state = IastAgent.objects.filter(pk=agent_id).values_list("bind_project__vul_validation", flat=True).first() - if state is None: - state = VulValidation.FOLLOW_GLOBAL - return state - - -def vul_recheck_state(agent_id): - project_level_validation = get_project_vul_validation_state(agent_id) - global_state = get_vul_validate() - if project_level_validation == VulValidation.FOLLOW_GLOBAL: - return global_state - return project_level_validation == VulValidation.ENABLE diff --git a/dongtai_protocol/views/agent_register.py b/dongtai_protocol/views/agent_register.py index 0e094e8b..fd997912 100644 --- a/dongtai_protocol/views/agent_register.py +++ b/dongtai_protocol/views/agent_register.py @@ -252,7 +252,7 @@ def post(self, request: Request): version_name = param.get("projectVersion", "V1.0") version_name = version_name if version_name else "V1.0" template_id = param.get("projectTemplateId", None) - + jvm_user_dir = param.get("jvmUserDirectory", "") if template_id is not None: template = IastProjectTemplate.objects.filter(pk=template_id).first() if not template: From 84778d0d3ad07359ac7733ce4129abdb2ef70dae Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Fri, 13 Oct 2023 11:03:47 +0800 Subject: [PATCH 3/8] feat: new metric handler. --- .../migrations/0036_auto_20231013_1051.py | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 dongtai_common/migrations/0036_auto_20231013_1051.py diff --git a/dongtai_common/migrations/0036_auto_20231013_1051.py b/dongtai_common/migrations/0036_auto_20231013_1051.py new file mode 100644 index 00000000..c7f5832c --- /dev/null +++ b/dongtai_common/migrations/0036_auto_20231013_1051.py @@ -0,0 +1,32 @@ +# Generated by Django 3.2.20 on 2023-10-13 10:51 + +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + ('dongtai_common', '0035_alter_user_phone'), + ] + + operations = [ + migrations.AddField( + model_name='iastagent', + name='jvm_user_dir', + field=models.CharField(default='', max_length=1024), + ), + migrations.CreateModel( + name='IastAgentDiskList', + fields=[ + ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('create_at', models.DateTimeField(auto_now_add=True, null=True)), + ('data', models.JSONField(default=dict)), + ('agent', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='dongtai_common.iastagent')), + ], + options={ + 'db_table': 'iast_agent_disk_list', + 'managed': True, + }, + ), + ] From b558ada22df0c49716f77511feabe3438e8f29c1 Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Fri, 13 Oct 2023 12:25:09 +0800 Subject: [PATCH 4/8] feat: new metric handler. --- .../migrations/0036_auto_20231013_1051.py | 26 ++++++++++--------- .../report/handler/metric_handler.py | 15 +---------- dongtai_protocol/views/agent_register.py | 2 +- 3 files changed, 16 insertions(+), 27 deletions(-) diff --git a/dongtai_common/migrations/0036_auto_20231013_1051.py b/dongtai_common/migrations/0036_auto_20231013_1051.py index c7f5832c..215f791b 100644 --- a/dongtai_common/migrations/0036_auto_20231013_1051.py +++ b/dongtai_common/migrations/0036_auto_20231013_1051.py @@ -5,28 +5,30 @@ class Migration(migrations.Migration): - dependencies = [ - ('dongtai_common', '0035_alter_user_phone'), + ("dongtai_common", "0035_alter_user_phone"), ] operations = [ migrations.AddField( - model_name='iastagent', - name='jvm_user_dir', - field=models.CharField(default='', max_length=1024), + model_name="iastagent", + name="jvm_user_dir", + field=models.CharField(default="", max_length=1024), ), migrations.CreateModel( - name='IastAgentDiskList', + name="IastAgentDiskList", fields=[ - ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), - ('create_at', models.DateTimeField(auto_now_add=True, null=True)), - ('data', models.JSONField(default=dict)), - ('agent', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='dongtai_common.iastagent')), + ("id", models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")), + ("create_at", models.DateTimeField(auto_now_add=True, null=True)), + ("data", models.JSONField(default=dict)), + ( + "agent", + models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to="dongtai_common.iastagent"), + ), ], options={ - 'db_table': 'iast_agent_disk_list', - 'managed': True, + "db_table": "iast_agent_disk_list", + "managed": True, }, ), ] diff --git a/dongtai_protocol/report/handler/metric_handler.py b/dongtai_protocol/report/handler/metric_handler.py index 36e243c7..db1e3f0d 100644 --- a/dongtai_protocol/report/handler/metric_handler.py +++ b/dongtai_protocol/report/handler/metric_handler.py @@ -2,28 +2,15 @@ # datetime:2020/10/23 11:56 import logging import time -from typing import Any - -from celery import shared_task -from celery_singleton import Singleton -from django.core.cache import cache -from django.db.models import Q, QuerySet -from django.utils.translation import gettext_lazy as _ from dongtai_common.models.agent import ( IastAgent, IastAgentDiskList, ) -from dongtai_common.models.heartbeat import IastHeartbeat -from dongtai_common.models.project import VulValidation -from dongtai_common.models.replay_queue import IastReplayQueue from dongtai_common.utils import const -from dongtai_common.utils.systemsettings import get_vul_validate +from dongtai_protocol.report.handler.heartbeat_handler import update_heartbeat from dongtai_protocol.report.handler.report_handler_interface import IReportHandler from dongtai_protocol.report.report_handler_factory import ReportHandler -from dongtai_web.vul_log.vul_log import log_recheck_vul -from dongtai_protocol.report.handler.heartbeat_handler import update_heartbeat -from dongtai_common.models.agent import IastAgentDiskList logger = logging.getLogger("dongtai.openapi") diff --git a/dongtai_protocol/views/agent_register.py b/dongtai_protocol/views/agent_register.py index fd997912..bb0d908e 100644 --- a/dongtai_protocol/views/agent_register.py +++ b/dongtai_protocol/views/agent_register.py @@ -252,7 +252,7 @@ def post(self, request: Request): version_name = param.get("projectVersion", "V1.0") version_name = version_name if version_name else "V1.0" template_id = param.get("projectTemplateId", None) - jvm_user_dir = param.get("jvmUserDirectory", "") + param.get("jvmUserDirectory", "") if template_id is not None: template = IastProjectTemplate.objects.filter(pk=template_id).first() if not template: From 380f37eb5709b60596a4d616fe54035285a8fe98 Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Fri, 13 Oct 2023 15:44:04 +0800 Subject: [PATCH 5/8] feat: new metric handler. --- dongtai_common/migrations/0036_auto_20231013_1051.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dongtai_common/migrations/0036_auto_20231013_1051.py b/dongtai_common/migrations/0036_auto_20231013_1051.py index 215f791b..9e32658e 100644 --- a/dongtai_common/migrations/0036_auto_20231013_1051.py +++ b/dongtai_common/migrations/0036_auto_20231013_1051.py @@ -1,7 +1,7 @@ # Generated by Django 3.2.20 on 2023-10-13 10:51 -from django.db import migrations, models import django.db.models.deletion +from django.db import migrations, models class Migration(migrations.Migration): From 93303249c9b90f43be325294c7fe18fc8ec87d81 Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Mon, 16 Oct 2023 11:23:03 +0800 Subject: [PATCH 6/8] feat: new metric handler. --- dongtai_protocol/views/agent_register.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/dongtai_protocol/views/agent_register.py b/dongtai_protocol/views/agent_register.py index bb0d908e..aabe841e 100644 --- a/dongtai_protocol/views/agent_register.py +++ b/dongtai_protocol/views/agent_register.py @@ -43,7 +43,7 @@ class AgentRegisterEndPoint(OpenApiEndPoint): description = "引擎注册" @staticmethod - def register_agent(token, version, language, project_name, user, project_version): + def register_agent(token, version, language, project_name, user, project_version, jvm_user_dir=""): project = IastProject.objects.values("id").filter(name=project_name).first() is_audit = AgentRegisterEndPoint.get_is_audit() project_id = -1 @@ -76,6 +76,7 @@ def register_agent(token, version, language, project_name, user, project_version language=language, is_audit=is_audit, allow_report=allow_report, + jvm_user_dir=jvm_user_dir, ) else: IastAgent.objects.filter(pk=agent_id).update( @@ -252,7 +253,7 @@ def post(self, request: Request): version_name = param.get("projectVersion", "V1.0") version_name = version_name if version_name else "V1.0" template_id = param.get("projectTemplateId", None) - param.get("jvmUserDirectory", "") + jvm_user_dir = param.get("jvmUserDirectory", "") if template_id is not None: template = IastProjectTemplate.objects.filter(pk=template_id).first() if not template: @@ -294,6 +295,7 @@ def post(self, request: Request): version=version, project_version=project_version, user=user, + jvm_user_dir=jvm_user_dir, ) else: agent_id = self.register_agent( @@ -303,6 +305,7 @@ def post(self, request: Request): version=version, user=user, project_version=None, + jvm_user_dir=jvm_user_dir, ) self.register_server( @@ -367,6 +370,7 @@ def __register_agent( language, is_audit, allow_report, + jvm_user_dir="", ): if exist_project: IastAgent.objects.filter( @@ -390,6 +394,7 @@ def __register_agent( is_audit=is_audit, allow_report=allow_report, department_id=1, + jvm_user_dir=jvm_user_dir, ) return agent.id From 2a73378f7fcbd98368bdd23ecc0f684d43c7da68 Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Mon, 16 Oct 2023 12:31:03 +0800 Subject: [PATCH 7/8] fix: change log level. --- .../report/handler/metric_handler.py | 39 ++++++++++++++++++- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/dongtai_protocol/report/handler/metric_handler.py b/dongtai_protocol/report/handler/metric_handler.py index db1e3f0d..ffa9c820 100644 --- a/dongtai_protocol/report/handler/metric_handler.py +++ b/dongtai_protocol/report/handler/metric_handler.py @@ -11,6 +11,7 @@ from dongtai_protocol.report.handler.heartbeat_handler import update_heartbeat from dongtai_protocol.report.handler.report_handler_interface import IReportHandler from dongtai_protocol.report.report_handler_factory import ReportHandler +import os logger = logging.getLogger("dongtai.openapi") @@ -31,7 +32,6 @@ def __init__(self): def parse(self): self.cpu = self.detail.get("cpu") self.memory = self.detail.get("memory") - self.disk = self.detail.get("disk") # extend now self.disk_list = self.detail.get("diskList") def has_permission(self): @@ -42,7 +42,7 @@ def save_heartbeat(self): default_dict = {"dt": int(time.time())} default_dict["memory"] = self.memory default_dict["cpu"] = self.cpu - default_dict["disk"] = self.disk + default_dict["disk"] = get_disk_from_disk_list(self.disk_list, self.agent.jvm_user_dir) update_heartbeat.delay(agent_id=self.agent_id, defaults=default_dict) IastAgentDiskList.objects.create(agent=self.agent_id, data=self.disk_list) @@ -54,3 +54,38 @@ def save(self): def get_agent(self, agent_id): return IastAgent.objects.filter(id=agent_id).first() + + +def get_disk_from_disk_list(disk_list: list, agent_jvm_dir: str): + total_space_bytes = 0 + usable_space_bytes = 0 + for disk in disk_list: + for partition in disk["partitionList"]: + if is_dir_under_another_dir(agent_jvm_dir, partition["mountPoint"]): + return get_disk_rate_dict(partition["totalSpaceBytes"], partition["usableSpaceBytes"]) + total_space_bytes += partition["totalSpaceBytes"] + usable_space_bytes += partition["usableSpaceBytes"] + return get_disk_rate_dict(total_space_bytes, usable_space_bytes) + + +def get_disk_rate_dict(total_space_bytes: int, usable_space_bytes: int): + return (total_space_bytes - usable_space_bytes) / total_space_bytes + + +def is_dir_under_another_dir(subdir: str, parent_dir: str) -> bool: + """Returns True if the subdirectory is under the parent directory, False otherwise. + + Args: + subdir: The subdirectory path. + parent_dir: The parent directory path. + + Returns: + True if the subdirectory is under the parent directory, False otherwise. + """ + + # Normalize the paths. + subdir = os.path.normpath(subdir) + parent_dir = os.path.normpath(parent_dir) + + # Check if the subdirectory path starts with the parent directory path. + return subdir.startswith(parent_dir) From 42fb9a6c15afcc9ffb6847cb77a15feeaff38cf4 Mon Sep 17 00:00:00 2001 From: bidaya0 Date: Mon, 16 Oct 2023 12:33:10 +0800 Subject: [PATCH 8/8] fix: change log level. --- dongtai_protocol/report/handler/metric_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dongtai_protocol/report/handler/metric_handler.py b/dongtai_protocol/report/handler/metric_handler.py index ffa9c820..fae42671 100644 --- a/dongtai_protocol/report/handler/metric_handler.py +++ b/dongtai_protocol/report/handler/metric_handler.py @@ -1,6 +1,7 @@ #!/usr/bin/env python # datetime:2020/10/23 11:56 import logging +import os import time from dongtai_common.models.agent import ( @@ -11,7 +12,6 @@ from dongtai_protocol.report.handler.heartbeat_handler import update_heartbeat from dongtai_protocol.report.handler.report_handler_interface import IReportHandler from dongtai_protocol.report.report_handler_factory import ReportHandler -import os logger = logging.getLogger("dongtai.openapi")