Merge pull request #1397 from HXSecurity/develop

Develop

Author: Bidaya0

dongtai_protocol/views/agent_register.py

@@ -38,8 +38,9 @@ class AgentRegisterEndPoint(OpenApiEndPoint):
     @staticmethod
     def register_agent(token, version, language, project_name, user,
                        project_version):
-        project = IastProject.objects.values('id').filter(name=project_name,
-                                                          user=user).first()
+        department = user.get_using_department()
+        project = IastProject.objects.values('id').filter(
+            name=project_name, department=department).first()
         is_audit = AgentRegisterEndPoint.get_is_audit()
         if project:
             if project_version:
@@ -256,6 +257,7 @@ def post(self, request: Request):
                 'mode': '插桩模式',
                 'latest_time': int(time.time()),
                 'template_id': template.id if template else -1,
+                'user': user,
             }
 
             default_params.update(
@@ -329,11 +331,20 @@ def post(self, request: Request):
 
     @staticmethod
     def get_agent_id(token, project_name, user, current_project_version_id):
-        queryset = IastAgent.objects.values('id').filter(
-            token=token,
-            project_name=project_name,
-            user=user,
-            project_version_id=current_project_version_id)
+        project = IastProject.objects.filter(
+            department=user.get_using_department(), name=project_name).first()
+        if project:
+            queryset = IastAgent.objects.values('id').filter(
+                token=token,
+                bind_project=project,
+                department=user.get_using_department(),
+                project_version_id=current_project_version_id)
+        else:
+            queryset = IastAgent.objects.values('id').filter(
+                token=token,
+                project_name=project_name,
+                department=user.get_using_department(),
+                project_version_id=current_project_version_id)
         agent = queryset.first()
         if agent:
             queryset.update(is_core_running=1, online=1, is_running=1)
@@ -344,8 +355,9 @@ def get_agent_id(token, project_name, user, current_project_version_id):
     def __register_agent(exist_project, token, user, version, project_id,
                          project_name, project_version_id, language, is_audit):
         if exist_project:
-            IastAgent.objects.filter(token=token, online=1,
-                                     user=user).update(online=0)
+            IastAgent.objects.filter(
+                token=token, online=1,
+                department=user.get_using_department()).update(online=0)
         agent = IastAgent.objects.create(
             token=token,
             version=version,
@@ -406,13 +418,11 @@ def project_create(default_params, project_name, user, version_name, template):
     project_created = False
     obj = IastProject.objects.filter(
         name=project_name,
-        user=user,
         department=department,
     ).first()
     if not obj:
         obj, project_created = IastProject.objects.get_or_create(
             name=project_name,
-            user=user,
             department=department,
             defaults=default_params,
         )

test/apiserver/test_agent_register_with_grouptoken.py

@@ -0,0 +1,93 @@
+######################################################################
+# @author      : bidaya0 (bidaya0@$HOSTNAME)
+# @file        : test_agent_register
+# @created     : 星期五 12月 10, 2021 14:46:44 CST
+#
+# @description :
+######################################################################
+
+from test.apiserver.test_agent_base import AgentTestCase
+from dongtai_common.models.agent import IastAgent
+from dongtai_common.models.department import Department
+from dongtai_web.projecttemplate.update_department_data import update_department_data
+import json
+import uuid
+
+
+class AgentNewRegisterGroupTokenTestCase(AgentTestCase):
+
+    def setUp(self):
+        super().setUp()
+        update_department_data()
+        self.client.force_authenticate(user=None)
+        self.test_agent_id = []
+
+    def test_rep_register(self):
+        department = Department.objects.first()
+        self.client.credentials(HTTP_AUTHORIZATION="Token GROUP" +
+                                department.token)
+        resp1 = self.raw_register(projectVersion='V1.1', projectName='PNAME')
+        data1 = json.loads(resp1.content)
+        resp2 = self.raw_register(name='newtoken',
+                                  projectVersion='V1.1',
+                                  projectName='PNAME')
+        data2 = json.loads(resp2.content)
+        self.assertNotEqual(data1['data']['id'], data2['data']['id'])
+
+    def test_rep_register_2(self):
+        department = Department.objects.first()
+        self.client.credentials(HTTP_AUTHORIZATION="Token GROUP" +
+                                department.token)
+        resp1 = self.raw_register(projectTemplateId=1,
+                                  projectVersion='V1.1',
+                                  projectName='PNAME')
+        data1 = json.loads(resp1.content)
+        self.assertEqual(data1['status'], 201)
+
+    def test_rep_register_3(self):
+        department = Department.objects.first()
+        self.client.credentials(HTTP_AUTHORIZATION="Token GROUP" +
+                                department.token)
+        resp1 = self.raw_register(token=uuid.uuid4().hex)
+        data1 = json.loads(resp1.content)
+        self.assertEqual(data1['status'], 201)
+        agent_id1 = data1['data']['id']
+        agent1 = IastAgent.objects.filter(pk=agent_id1).first()
+        agent1.bind_project.department_id = 2
+        agent1.bind_project.save()
+        resp2 = self.raw_register(token=uuid.uuid4().hex)
+        data2 = json.loads(resp2.content)
+        self.assertEqual(data2['status'], 201)
+        agent_id2 = data2['data']['id']
+        agent2 = IastAgent.objects.filter(pk=agent_id2).first()
+        self.assertNotEqual(agent1.bind_project_id, agent2.bind_project_id)
+        self.assertNotEqual(agent1.project_version_id,
+                            agent2.project_version_id)
+
+    def test_rep_register_4(self):
+        department = Department.objects.first()
+        self.client.credentials(HTTP_AUTHORIZATION="Token GROUP" +
+                                department.token)
+        resp1 = self.raw_register(token=uuid.uuid4().hex)
+        data1 = json.loads(resp1.content)
+        self.assertEqual(data1['status'], 201)
+        agent_id1 = data1['data']['id']
+        agent1 = IastAgent.objects.filter(pk=agent_id1).first()
+        agent1.bind_project.department_id = 2
+        agent1.bind_project.save()
+        resp2 = self.raw_register(token=uuid.uuid4().hex)
+        data2 = json.loads(resp2.content)
+        self.assertEqual(data2['status'], 201)
+        agent_id2 = data2['data']['id']
+        agent2 = IastAgent.objects.filter(pk=agent_id2).first()
+        agent2.bind_project_id = agent1.bind_project_id
+        agent2.save()
+        self.client.credentials(HTTP_AUTHORIZATION="Token GROUP" +
+                                department.token)
+        resp3 = self.raw_register(token=uuid.uuid4().hex)
+        data3 = json.loads(resp3.content)
+        self.assertEqual(data3['status'], 201)
+        agent_id3 = data3['data']['id']
+        agent3 = IastAgent.objects.filter(pk=agent_id3).first()
+        self.assertNotEqual(agent_id3, agent_id2)
+        self.assertNotEqual(agent3.bind_project_id, agent2.bind_project_id)