feature: 云梯安全组添加白名单功能 (closed TencentBlueKing#1760)

apps/backend/components/collections/agent_new/configure_policy.py

@@ -36,7 +36,8 @@ def _execute(self, data, parent_data, common_data):
         for host in host_id_obj_map.values():
             ip_list.extend([host.outer_ip, host.login_ip])
         # 不同的安全组工厂添加策略后得到的输出可能是不同的，输出到outputs中，在schedule中由工厂对应的check_result方法来校验结果
-        data.outputs.add_ip_output = security_group_factory.add_ips_to_security_group(ip_list)
+        creator: str = common_data.subscription.creator
+        data.outputs.add_ip_output = security_group_factory.add_ips_to_security_group(ip_list, creator=creator)
         data.outputs.polling_time = 0
         return True
 

apps/backend/tests/components/collections/agent_new/test_configure_policy.py

@@ -26,6 +26,7 @@
 from apps.node_man.handlers.security_group import (
     SopsSecurityGroupFactory,
     TencentVpcSecurityGroupFactory,
+    YunTiSecurityGroupFactory,
 )
 from pipeline.component_framework.test import (
     ComponentTestCase,
@@ -204,3 +205,101 @@ def cases(self):
                 execute_call_assertion=None,
             ),
         ]
+
+
+class MockYunTiClient(mock_data_utils.BaseMockClient):
+    """mock云梯客户端"""
+
+    MOCK_RETUREN_DATA = {
+        "get_security_group_details": {
+            "pilicies": {
+                "SecurityGroupPolicySet": {
+                    "Ingress": [
+                        {
+                            "Port": "ALL",
+                            "CidrBlock": "0.0.0.1",
+                            "Ipv6CidrBlock": "",
+                            "SecurityGroupId": "",
+                            "Action": "ACCEPT",
+                            "Protocol": "ALL",
+                            "PolicyDescription": "",
+                        }
+                    ],
+                    "Version": 1,
+                }
+            },
+            "SecurityGroupId": "test_sid",
+        },
+        "operate_security_group": {},
+    }
+
+    def __init__(self, *args, **kwargs):
+
+        self.get_security_group_details = self.generate_magic_mock(
+            mock_data_utils.MockReturn(
+                return_type=mock_data_utils.MockReturnType.RETURN_VALUE.value,
+                return_obj=self.MOCK_RETUREN_DATA["get_security_group_details"],
+            )
+        )
+        self.operate_security_group = self.generate_magic_mock(
+            mock_data_utils.MockReturn(
+                return_type=mock_data_utils.MockReturnType.RETURN_VALUE.value,
+                return_obj=self.MOCK_RETUREN_DATA["operate_security_group"],
+            )
+        )
+
+
+class YunTiConfigurePolicyComponentBaseTest(ConfigurePolicyComponentBaseTest):
+    YUNTI_MOCK_CLIENT_PATH = "apps.node_man.handlers.security_group.YunTiApi"
+
+    def setUp(self):
+        models.GlobalSettings.set_config(
+            models.GlobalSettings.KeyEnum.SECURITY_GROUP_TYPE.value,
+            YunTiSecurityGroupFactory.SECURITY_GROUP_TYPE,
+        )
+        models.GlobalSettings.set_config(
+            models.GlobalSettings.KeyEnum.YUNTI_POLICY_CONFIGS.value,
+            [
+                {
+                    "dept_id": 0,
+                    "region": "ap-test",
+                    "sid": "test_sid",
+                    "group_name": "test_group_name",
+                    "port": "ALL",
+                    "action": "ACCEPT",
+                    "protocol": "ALL",
+                }
+            ],
+        )
+        mock.patch(self.YUNTI_MOCK_CLIENT_PATH, MockYunTiClient()).start()
+        super().setUp()
+
+    def cases(self):
+        ip = self.obj_factory.host_objs[0].inner_ip
+        return [
+            ComponentTestCase(
+                name="通过云梯配置策略",
+                inputs=self.common_inputs,
+                parent_data={},
+                execute_assertion=ExecuteAssertion(
+                    success=bool(self.common_inputs["subscription_instance_ids"]),
+                    outputs={
+                        "add_ip_output": {"ip_list": [ip, ip]},
+                        "polling_time": 0,
+                        "succeeded_subscription_instance_ids": self.common_inputs["subscription_instance_ids"],
+                    },
+                ),
+                schedule_assertion=[
+                    ScheduleAssertion(
+                        success=True,
+                        schedule_finished=True,
+                        outputs={
+                            "add_ip_output": {"ip_list": [ip, ip]},
+                            "polling_time": 0,
+                            "succeeded_subscription_instance_ids": self.common_inputs["subscription_instance_ids"],
+                        },
+                    ),
+                ],
+                execute_call_assertion=None,
+            ),
+        ]

apps/backend/utils/dataclass.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making 蓝鲸智云-节点管理(BlueKing-BK-NODEMAN) available.
+Copyright (C) 2017-2022 THL A29 Limited, a Tencent company. All rights reserved.
+Licensed under the MIT License (the "License"); you may not use this file except in compliance with the License.
+You may obtain a copy of the License at https://opensource.org/licenses/MIT
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+"""
+import copy
+from dataclasses import _is_dataclass_instance, fields
+
+
+def _asdict_inner(obj, dict_factory):
+    if _is_dataclass_instance(obj):
+        result = []
+        for f in fields(obj):
+            value = _asdict_inner(getattr(obj, f.name), dict_factory)
+            # 过滤掉为空或者None的字段
+            if not value:
+                continue
+            result.append((f.name, value))
+        return dict_factory(result)
+    elif isinstance(obj, (list, tuple)):
+        return type(obj)(_asdict_inner(v, dict_factory) for v in obj)
+    elif isinstance(obj, dict):
+        return type(obj)((_asdict_inner(k, dict_factory), _asdict_inner(v, dict_factory)) for k, v in obj.items())
+    else:
+        return copy.deepcopy(obj)
+
+
+def asdict(obj, *, dict_factory=dict):
+    if not _is_dataclass_instance(obj):
+        raise TypeError("asdict() should be called on dataclass instances")
+    return _asdict_inner(obj, dict_factory)

apps/node_man/exceptions.py

@@ -214,3 +214,9 @@ class ServiceInstanceNotFoundError(NodeManBaseException):
     MESSAGE = _("服务实例不存在")
     MESSAGE_TPL = _("服务实例 -> [{id}] 不存在")
     ERROR_CODE = 42
+
+
+class YunTiPolicyConfigNotExistsError(NodeManBaseException):
+    MESSAGE = _("云梯策略配置不存在")
+    MESSAGE_TPL = _("云梯策略配置不存在")
+    ERROR_CODE = 43

apps/node_man/handlers/security_group.py

@@ -1,12 +1,41 @@
 # -*- coding: utf-8 -*-
 import abc
-from typing import Dict, List, Optional
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
 
 from django.conf import settings
 
-from apps.node_man.exceptions import ConfigurationPolicyError
+from apps.backend.utils.dataclass import asdict
+from apps.node_man.exceptions import (
+    ConfigurationPolicyError,
+    YunTiPolicyConfigNotExistsError,
+)
+from apps.node_man.models import GlobalSettings
 from apps.node_man.policy.tencent_vpc_client import VpcClient
-from common.api import SopsApi
+from apps.utils.batch_request import request_multi_thread
+from common.api import SopsApi, YunTiApi
+from common.log import logger
+
+
+@dataclass
+class YunTiPolicyData:
+    Protocol: str
+    CidrBlock: str
+    Port: str
+    Action: str
+    PolicyDescription: Optional[str] = None
+    Ipv6CidrBlock: Optional[str] = None
+
+
+@dataclass
+class YunTiPolicyConfig:
+    dept_id: int
+    region: str
+    sid: str
+    group_name: str
+    port: str
+    action: str
+    protocol: str
 
 
 class BaseSecurityGroupFactory(abc.ABC):
@@ -16,7 +45,7 @@ def describe_security_group_address(self) -> List[str]:
         """获取安全组IP地址"""
         raise NotImplementedError()
 
-    def add_ips_to_security_group(self, ip_list: List[str]) -> Dict:
+    def add_ips_to_security_group(self, ip_list: List[str], creator: str = None) -> Dict:
         """添加IP到安全组中，输出的字典用作check_result的入参"""
         raise NotImplementedError()
 
@@ -31,7 +60,7 @@ class SopsSecurityGroupFactory(BaseSecurityGroupFactory):
     def describe_security_group_address(self):
         pass
 
-    def add_ips_to_security_group(self, ip_list: List[str]) -> Dict:
+    def add_ips_to_security_group(self, ip_list: List[str], creator: str = None) -> Dict:
         task_id = SopsApi.create_task(
             {
                 "name": "NodeMan Configure SecurityGroup",
@@ -76,7 +105,7 @@ def describe_security_group_address(self) -> List:
                 ip_set = ip_set & set(client.describe_address_templates(template))
         return list(ip_set)
 
-    def add_ips_to_security_group(self, ip_list: List[str]):
+    def add_ips_to_security_group(self, ip_list: List[str], creator: str = None):
         client = VpcClient()
         for template in client.ip_templates:
             using_ip_list = client.describe_address_templates(template)
@@ -94,6 +123,142 @@ def check_result(self, add_ip_output: Dict) -> bool:
         return set(add_ip_output["ip_list"]).issubset(set(current_ip_list))
 
 
+class YunTiSecurityGroupFactory(BaseSecurityGroupFactory):
+    SECURITY_GROUP_TYPE: str = "YUNTI"
+
+    def __init__(self) -> None:
+        """
+        policies_config: example
+        [
+            {
+                "dept_id": 0,
+                "region": "ap-xxx",
+                "sid": "xxxx",
+                "group_name": "xxx",
+                "port": "ALL",
+                "action": "ACCEPT",
+                "protocol": "ALL",
+                ""
+            }
+        ]
+        """
+        self.policy_configs: List[Dict[str, Any]] = GlobalSettings.get_config(
+            key=GlobalSettings.KeyEnum.YUNTI_POLICY_CONFIGS.value, default=[]
+        )
+        if not self.policy_configs:
+            raise YunTiPolicyConfigNotExistsError()
+
+    def describe_security_group_address(self) -> Dict:
+        # 批量获取当前安全组策略详情
+        params_list: List = []
+        for policy_config in self.policy_configs:
+            config = YunTiPolicyConfig(**policy_config)
+            params_list.append(
+                {
+                    "params": {
+                        "method": "get-security-group-policies",
+                        "params": {
+                            "deptId": config.dept_id,
+                            "region": config.region,
+                            "sid": config.sid,
+                        },
+                        "no_request": True,
+                    },
+                }
+            )
+        # 批量请求
+        result: List[Dict[str, Any]] = request_multi_thread(
+            func=YunTiApi.get_security_group_details,
+            params_list=params_list,
+            get_data=lambda x: [x],
+        )
+        return {sid_info["SecurityGroupId"]: sid_info for sid_info in result}
+
+    def add_ips_to_security_group(self, ip_list: List[str], creator: str = None):
+        result: Dict[str, Dict[str, Any]] = self.describe_security_group_address()
+        params_list: List[Dict[str, Any]] = []
+
+        for policy_config in self.policy_configs:
+            config = YunTiPolicyConfig(**policy_config)
+            # 新策略列表
+            new_in_gress: Dict[str, Dict[str, Any]] = {}
+            for ip in ip_list:
+                new_in_gress[ip] = asdict(
+                    YunTiPolicyData(
+                        Protocol=config.protocol,
+                        CidrBlock=ip,
+                        Port=config.port,
+                        Action=config.action,
+                        PolicyDescription="",
+                        Ipv6CidrBlock="",
+                    )
+                )
+
+            version: str = result[config.sid]["pilicies"]["SecurityGroupPolicySet"]["Version"]
+            current_policies: List[Dict[str, Any]] = result[config.sid]["pilicies"]["SecurityGroupPolicySet"]["Ingress"]
+
+            # 已有策略列表
+            in_gress: Dict[str, Dict[str, Any]] = {}
+            for policy in current_policies:
+                in_gress[policy["CidrBlock"]] = asdict(
+                    YunTiPolicyData(
+                        Protocol=policy["Protocol"],
+                        CidrBlock=policy["CidrBlock"],
+                        Port=policy["Port"],
+                        Action=policy["Action"],
+                        PolicyDescription=policy["PolicyDescription"],
+                        Ipv6CidrBlock=policy["Ipv6CidrBlock"],
+                    )
+                )
+            # 增加新IP
+            in_gress.update(new_in_gress)
+
+            params_list.append(
+                {
+                    "params": {
+                        "method": "createSecurityGroupForm",
+                        "params": {
+                            "deptId": config.dept_id,
+                            "region": config.region,
+                            "sid": config.sid,
+                            "type": "modify",
+                            "mark": "Add proxy whitelist to Shangyun security group security.",
+                            "groupName": config.group_name,
+                            "groupDesc": "Proxy whitelist for Shangyun",
+                            "creator": creator,
+                            "ext": {
+                                "Version": version,
+                                "Egress": [],
+                                "Ingress": list(in_gress.values()),
+                            },
+                        },
+                        "no_request": True,
+                    },
+                }
+            )
+        # 批量请求
+        logger.info(f"Add proxy whitelist to Shangyun security group security. params: {params_list}")
+        request_multi_thread(func=YunTiApi.operate_security_group, params_list=params_list)
+        return {"ip_list": ip_list}
+
+    def check_result(self, add_ip_output: Dict) -> bool:
+        """检查IP列表是否已添加到安全组中"""
+        result: Dict[str, Dict[str, Any]] = self.describe_security_group_address()
+        is_success: bool = True
+        for policy_config in self.policy_configs:
+            config = YunTiPolicyConfig(**policy_config)
+            current_policies: List[Dict[str, Any]] = result[config.sid]["pilicies"]["SecurityGroupPolicySet"]["Ingress"]
+            current_ip_list = [policy["CidrBlock"] for policy in current_policies]
+            logger.info(
+                f"check_result: Add proxy whitelist to Shangyun security group security. "
+                f"sid: {config.sid} ip_list: {add_ip_output['ip_list']}"
+            )
+            # 需添加的IP列表是已有IP的子集，则认为已添加成功
+            is_success: bool = is_success and set(add_ip_output["ip_list"]).issubset(set(current_ip_list))
+
+        return is_success
+
+
 def get_security_group_factory(security_group_type: Optional[str]) -> BaseSecurityGroupFactory:
     """获取安全组工厂，返回None表示无需配置安全组"""
     factory_map = {factory.SECURITY_GROUP_TYPE: factory for factory in BaseSecurityGroupFactory.__subclasses__()}

apps/node_man/models.py

@@ -136,6 +136,8 @@ class KeyEnum(Enum):
         CLEAN_SUBSCRIPTION_DATA_MAP = "CLEAN_SUBSCRIPTION_DATA_MAP"
         # 是否开启Agent包管理
         ENABLE_AGENT_PKG_MANAGE = "ENABLE_AGENT_PKG_MANAGE"
+        # 云梯策略相关配置
+        YUNTI_POLICY_CONFIGS = "YUNTI_POLICY_CONFIGS"
 
     key = models.CharField(_("键"), max_length=255, db_index=True, primary_key=True)
     v_json = JSONField(_("值"))