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..9e32658e --- /dev/null +++ b/dongtai_common/migrations/0036_auto_20231013_1051.py @@ -0,0 +1,34 @@ +# Generated by Django 3.2.20 on 2023-10-13 10:51 + +import django.db.models.deletion +from django.db import migrations, models + + +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, + }, + ), + ] 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 new file mode 100644 index 00000000..fae42671 --- /dev/null +++ b/dongtai_protocol/report/handler/metric_handler.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python +# datetime:2020/10/23 11:56 +import logging +import os +import time + +from dongtai_common.models.agent import ( + IastAgent, + IastAgentDiskList, +) +from dongtai_common.utils import const +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 + +logger = logging.getLogger("dongtai.openapi") + + +@ReportHandler.register(const.REPORT_METRIC) +class ReportMetricHandler(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_list = self.detail.get("diskList") + + 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())} + default_dict["memory"] = self.memory + default_dict["cpu"] = self.cpu + 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) + + def get_result(self, msg=None): + return [] + + def save(self): + self.save_heartbeat() + + 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) diff --git a/dongtai_protocol/views/agent_register.py b/dongtai_protocol/views/agent_register.py index 0e094e8b..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) - + 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