feat: add event and service.

.github/workflows/teststate.yml

@@ -144,7 +144,7 @@ jobs:
 
       - name: Django Unit Testing
         run: |
-          pep8  --ignore=E501,E402,E231,E302,E301,E303,E251,E265,W291,E126,W391,W503,E128,E122,E241,E225,E401,E701,W191 --statistics .
+          pep8  --ignore=E501,E402,E231,E302,E301,E303,E251,E265,W291,E126,W391,W503,E128,E122,E241,E225,E701,W191 --statistics .
 
   Run-Unit-TEST:
     runs-on: ubuntu-latest

dongtai_common/models/agent.py

@@ -13,6 +13,8 @@
 from dongtai_common.utils.settings import get_managed
 from dongtai_common.models.project import IastProject
 from dongtai_common.models.project_version import IastProjectVersion
+import json
+
 
 class IastAgent(models.Model):
     token = models.CharField(max_length=255, blank=True, null=True)
@@ -38,10 +40,10 @@ class IastAgent(models.Model):
                                      null=True,
                                      default=-1)
     project_version = models.ForeignKey(IastProjectVersion,
-                                     on_delete=models.CASCADE,
-                                     blank=True,
-                                     null=True,
-                                     default=-1)
+                                        on_delete=models.CASCADE,
+                                        blank=True,
+                                        null=True,
+                                        default=-1)
     project_name = models.CharField(max_length=255, blank=True, null=True)
     online = models.PositiveSmallIntegerField(blank=True, default=0)
     language = models.CharField(max_length=10, blank=True, null=True)
@@ -59,12 +61,21 @@ class IastAgent(models.Model):
     actual_running_status = models.IntegerField(default=1, null=False)
     except_running_status = models.IntegerField(default=1, null=False)
     state_status = models.IntegerField(default=1, null=False)
-
+    events = models.JSONField(null=False, default=lambda: ['注册成功'])
 
     class Meta:
         managed = get_managed()
         db_table = 'iast_agent'
 
+    def append_events(self, event: str):
+        events_list = json.loads(self.events)
+        events_list.append(event)
+        self.events = json.dumps(events_list)
+        self.save()
+
+    def only_register(self):
+        events_list = json.loads(self.events)
+        return events_list == ['注册成功']
 
 #class IastAgent(models.Model):
 #

dongtai_common/models/server.py

@@ -30,6 +30,7 @@ class IastServer(models.Model):
                                 null=True,
                                 default='')
     pid = models.IntegerField(blank=True, null=True)
+    ipaddresslist = models.JSONField(null=False, default=list)
 
     class Meta:
         managed = get_managed()

dongtai_engine/signals/handlers/vul_handler.py

@@ -3,7 +3,8 @@
 # author: [email protected]
 # datetime: 2021/4/30 下午3:00
 # project: dongtai-engine
-import json,random
+import json
+import random
 import time
 import requests
 from celery.apps.worker import logger

dongtai_protocol/report/report_handler_factory.py

@@ -4,7 +4,10 @@
 # datetime:2020/10/23 12:00
 # software: PyCharm
 # project: webapi
-import logging, requests, json, time
+import logging
+import requests
+import json
+import time
 from django.utils.translation import gettext_lazy as _
 from dongtai_conf import settings
 from dongtai_protocol.report.log_service import LogService

dongtai_protocol/views/agent_download.py

@@ -5,8 +5,10 @@
 # software: PyCharm
 # project: lingzhi-agent-server
 import json
-import os, re
-import uuid, logging
+import os
+import re
+import uuid
+import logging
 
 from django.http import FileResponse
 from dongtai_common.endpoint import UserEndPoint, R
@@ -19,7 +21,8 @@
 from dongtai_conf.settings import BUCKET_NAME_BASE_URL, VERSION
 
 import shutil
-import tarfile, os
+import tarfile
+import os
 import threading
 import time
 

dongtai_protocol/views/agent_register.py

@@ -99,8 +99,10 @@ def get_runtime(envs):
         return ''
 
     @staticmethod
-    def register_server(agent_id, hostname, network, container_name, server_addr, server_port, cluster_name,cluster_version,
-                        server_path, server_env, pid):
+    def register_server(agent_id, hostname, network, container_name,
+                        server_addr, server_port, cluster_name,
+                        cluster_version, server_path, server_env, pid,
+                        server_ipaddresslist):
         """
         注册server，并关联server至agent
         :param agent_id:
@@ -131,12 +133,16 @@ def register_server(agent_id, hostname, network, container_name, server_addr, se
         try:
             port = int(server_port)
         except Exception as e:
-            logger.error(_('The server port does not exist, has been set to the default: 0'))
+            logger.error(
+                _('The server port does not exist, has been set to the default: 0'
+                  ))
             port = 0
 
         server_id = agent.server_id
 
-        server = IastServer.objects.filter(id=server_id).first() if server_id else None
+        server = IastServer.objects.filter(
+            id=server_id).first() if server_id else None
+        ipaddresslist = json.dumps(server_ipaddresslist)
         if server:
             server.hostname = hostname
             server.network = network
@@ -149,7 +155,19 @@ def register_server(agent_id, hostname, network, container_name, server_addr, se
             server.cluster_version = cluster_version
             server.status = 'online'
             server.update_time = int(time.time())
-            server.save(update_fields=['hostname', 'command', 'ip', 'port', 'env', 'status', 'update_time', 'cluster_name', 'cluster_version'])
+            server.ipaddresslist = ipaddresslist
+            server.save(update_fields=[
+                'hostname',
+                'command',
+                'ip',
+                'port',
+                'env',
+                'status',
+                'update_time',
+                'cluster_name',
+                'cluster_version',
+                "ipaddresslist",
+            ])
         else:
             server = IastServer.objects.create(
                 hostname=hostname,
@@ -167,34 +185,33 @@ def register_server(agent_id, hostname, network, container_name, server_addr, se
                 command=command,
                 runtime=AgentRegisterEndPoint.get_runtime(envs),
                 create_time=int(time.time()),
-                update_time=int(time.time())
+                update_time=int(time.time()),
+                ipaddresslist=ipaddresslist,
             )
             agent.server_id = server.id
             agent.save(update_fields=['server_id'])
             logger.info(_('Server record creation success'))
 
-    @extend_schema(
-        description='Agent Register, Data is Gzip',
-        parameters=[
-            DongTaiParameter.AGENT_NAME,
-            DongTaiParameter.LANGUAGE,
-            DongTaiParameter.VERSION,
-            DongTaiParameter.PROJECT_NAME,
-            DongTaiParameter.HOSTNAME,
-            DongTaiParameter.NETWORK,
-            DongTaiParameter.CONTAINER_NAME,
-            DongTaiParameter.SERVER_ADDR,
-            DongTaiParameter.SERVER_PORT,
-            DongTaiParameter.SERVER_PATH,
-            DongTaiParameter.SERVER_ENV,
-            DongTaiParameter.PID,
-            DongTaiParameter.AUTO_CREATE_PROJECT,
-        ],
-        responses=[
-            {204: None}
-        ],
-        methods=['POST']
-    )
+    @extend_schema(description='Agent Register, Data is Gzip',
+                   parameters=[
+                       DongTaiParameter.AGENT_NAME,
+                       DongTaiParameter.LANGUAGE,
+                       DongTaiParameter.VERSION,
+                       DongTaiParameter.PROJECT_NAME,
+                       DongTaiParameter.HOSTNAME,
+                       DongTaiParameter.NETWORK,
+                       DongTaiParameter.CONTAINER_NAME,
+                       DongTaiParameter.SERVER_ADDR,
+                       DongTaiParameter.SERVER_PORT,
+                       DongTaiParameter.SERVER_PATH,
+                       DongTaiParameter.SERVER_ENV,
+                       DongTaiParameter.PID,
+                       DongTaiParameter.AUTO_CREATE_PROJECT,
+                   ],
+                   responses=[{
+                       204: None
+                   }],
+                   methods=['POST'])
     def post(self, request: Request):
         try:
             param = parse_data(request.read())
@@ -282,6 +299,7 @@ def post(self, request: Request):
                 cluster_version=cluster_version,
                 server_env=server_env,
                 pid=pid,
+                server_ipaddresslist = get_ipaddresslist(network),
             )
 
             core_auto_start = 0
@@ -353,3 +371,20 @@ def get_ipaddress(network: str):
     except Exception as e:
         logger.error(e, exc_info=True)
         return ''
+
+
+def get_ipaddresslist(network: str) -> list:
+    try:
+        network_data = json.loads(network)
+        if isinstance(network_data, list):
+            res = []
+            for i in network_data:
+                res.append(i['ip'])
+            return res
+        if isinstance(network_data, dict):
+            return [network_data['ip']]
+    except KeyError as e:
+        logger.error(network_data,exc_info=e)
+    except Exception as e:
+        logger.error(e, exc_info=e)
+    return []

dongtai_protocol/views/except_action.py

@@ -10,10 +10,13 @@
 from urllib.parse import urlparse, urlunparse
 from dongtai_web.views.project_add import is_ip
 from rest_framework.viewsets import ViewSet
+
 logger = logging.getLogger('dongtai.openapi')
 
+EVENT_MAPPING = {1: "加载成功", 2: "降级成功", 3: "卸载成功", 4: "代码异常"}
+
 
-class AgentActionV2EndPoint(OpenApiEndPoint,ViewSet):
+class AgentActionV2EndPoint(OpenApiEndPoint, ViewSet):
 
     @extend_schema(description='Agent Update, Data is Gzip',
                    responses=[{
@@ -26,12 +29,18 @@ def actual_running_status(self, request):
             agent_id = int(param.get('agentId', None))
             actual_running_status = int(param.get('actualRunningStatus', None))
             state_status = int(param.get('stateStatus', None))
-        except Exception as e:
-            logger.error(e, exc_info=True)
+        except Exception as exception:
+            logger.error(exception, exc_info=True)
             return R.failure(msg="参数错误")
         agent = IastAgent.objects.filter(pk=agent_id).first()
         if not agent:
             return R.failure(msg=_("Agent not found"))
+        if agent.only_register(
+        ) or agent.actual_running_status != actual_running_status:
+            if not agent.events:
+                agent.append_events("注册成功")
+            if agent.actual_running_status in EVENT_MAPPING:
+                agent.append_events(EVENT_MAPPING[agent.actual_running_status])
         agent.actual_running_status = actual_running_status
         agent.state_status = state_status
         agent.save()

dongtai_protocol/views/report_upload.py

@@ -7,7 +7,8 @@
 
 from dongtai_common.endpoint import OpenApiEndPoint, R
 from drf_spectacular.utils import extend_schema
-import time,logging
+import time
+import logging
 from dongtai_protocol.api_schema import DongTaiParameter
 from dongtai_protocol.decrypter import parse_data
 from dongtai_protocol.report.report_handler_factory import ReportHandler

dongtai_web/base/update_project_version.py

@@ -3,7 +3,8 @@
 # author:sjh
 # software: PyCharm
 # project: lingzhi-webapi
-import logging, time
+import logging
+import time
 from dongtai_common.endpoint import UserEndPoint, R
 from dongtai_common.models.project_version import IastProjectVersion
 from dongtai_common.models.project import IastProject

dongtai_web/views/agents_v2.py

@@ -75,8 +75,10 @@ def pagenation_list(self, request):
                     agent['except_running_status'],
                     agent['online'],
                 )
-                agent['ipaddresses'] = ['172.22.22.11', '10.0.2.15']
-                agent['events'] = ['注册成功', '加载成功']
+                agent['ipaddresses'] = get_service_addrs(
+                    agent['server__ipaddresslist'], agent['server__port'])
+                if not agent['events']:
+                    agent['events'] = ['注册成功']
             data = {'agents': queryset, "summary": summary}
         except Exception as e:
             logger.error("agents pagenation_list error:{}".format(e),
@@ -102,6 +104,10 @@ def agent_stat(self, request):
         return R.success(data=res)
 
 
+def get_service_addrs(ip_list: list, port: int) -> list:
+    return list(map(lambda x: x + str(port), ip_list))
+
+
 def get_agent_stat(agent_id: int, user_id: int) -> dict:
     res = {}
     res['api_count'] = IastApiRoute.objects.filter(
@@ -189,7 +195,8 @@ def query_agent(filter_condiction=Q()) -> QuerySet:
     return IastAgent.objects.filter(filter_condiction).values(
         'alias', 'token', 'bind_project__name', 'bind_project__user__username',
         'language', 'server__ip', 'server__port', 'server__path',
-        'server__hostname', 'heartbeat__memory', 'heartbeat__cpu',
-        'heartbeat__disk', 'register_time', 'is_core_running', 'is_control',
-        'online', 'id', 'bind_project__id', 'version', 'except_running_status',
+        'server__ipaddresslist', 'events', 'server__hostname',
+        'heartbeat__memory', 'heartbeat__cpu', 'heartbeat__disk',
+        'register_time', 'is_core_running', 'is_control', 'online', 'id',
+        'bind_project__id', 'version', 'except_running_status',
         'actual_running_status', 'state_status').order_by('-latest_time')

dongtai_web/views/project_version_current.py

@@ -4,7 +4,8 @@
 
 # software: PyCharm
 # project: lingzhi-webapi
-import logging, time
+import logging
+import time
 from dongtai_common.endpoint import R
 from django.db.models import Q
 from dongtai_common.endpoint import UserEndPoint

dongtai_web/views/project_version_delete.py

@@ -3,7 +3,8 @@
 # author:sjh
 # software: PyCharm
 # project: lingzhi-webapi
-import logging, time
+import logging
+import time
 from dongtai_common.endpoint import R
 from dongtai_common.endpoint import UserEndPoint
 from dongtai_common.models.project_version import IastProjectVersion

dongtai_web/views/project_version_list.py

@@ -3,7 +3,8 @@
 # author:sjh
 # software: PyCharm
 # project: lingzhi-webapi
-import logging, time
+import logging
+import time
 from dongtai_common.endpoint import R
 from dongtai_common.endpoint import UserEndPoint
 from dongtai_common.models.project_version import IastProjectVersion

dongtai_web/views/vul_summary.py

@@ -20,7 +20,8 @@
 
 from dongtai_web.serializers.vul import VulSummaryResponseDataSerializer
 
-import copy,time
+import copy
+import time
 
 
 _ResponseSerializer = get_response_serializer(VulSummaryResponseDataSerializer())

test/apiserver/test_agent_register.py

@@ -19,7 +19,6 @@ def setUp(self):
 
     def test_rep_register(self):
         data1 = self.raw_register(name='rep_data')
-        self.agent_heartbeat()
         data2 = self.raw_register(name='rep_data')
         assert data1.status_code == 200 and data2.status_code == 200
         self.test_agent_id += [
@@ -31,3 +30,9 @@ def test_rep_register(self):
     def test_register(self):
         assert not IastAgent.objects.filter(pk=self.agent_id,
                                             project_version_id=0).exists()
+
+    def test_old_register(self):
+        dic = {
+            "network": "{\"name\":\"en0\",\"ip\":\"192.168.2.143\"}",
+        }
+        self.raw_register(**dic)