1、更新策略模板增删改查和alertmessage的模板处理，增加is_full_flow和description字段

7 files changed, 195 insertions, 56 deletions

diff --git a/src/main/java/com/realtime/protection/configuration/entity/alert/AlertMessage.java b/src/main/java/com/realtime/protection/configuration/entity/alert/AlertMessage.java
index 06b5e26..f24d89c 100644
--- a/src/main/java/com/realtime/protection/configuration/entity/alert/AlertMessage.java
+++ b/src/main/java/com/realtime/protection/configuration/entity/alert/AlertMessage.java
@@ -35,4 +35,6 @@ public class AlertMessage {
     @JsonProperty("alert_message_uuid")
     private String alertMessageUUID;
 
+    @JsonProperty("protect_object_is_src_dst")
+    private int protectIsSrcOrDst;
 }
diff --git a/src/main/java/com/realtime/protection/configuration/entity/defense/template/ProtectLevel.java b/src/main/java/com/realtime/protection/configuration/entity/defense/template/ProtectLevel.java
index a34f019..8e954b0 100644
--- a/src/main/java/com/realtime/protection/configuration/entity/defense/template/ProtectLevel.java
+++ b/src/main/java/com/realtime/protection/configuration/entity/defense/template/ProtectLevel.java
@@ -28,4 +28,7 @@ public class ProtectLevel {
 
     @Schema(description = "该防护等级是否需要提取DNS")
     private Boolean hasDNS = false;
+
+    @Schema(description = "该防护等级是处置防护对象的全流量or单向流量")
+    private Boolean isFullFlow = false;
 }
diff --git a/src/main/java/com/realtime/protection/configuration/entity/defense/template/Template.java b/src/main/java/com/realtime/protection/configuration/entity/defense/template/Template.java
index 5505daf..890fd07 100644
--- a/src/main/java/com/realtime/protection/configuration/entity/defense/template/Template.java
+++ b/src/main/java/com/realtime/protection/configuration/entity/defense/template/Template.java
@@ -13,7 +13,7 @@ public class Template {
 
     @JsonProperty("template_name")
     @NotNull(message = "template_name字段不能为空。")
-    @Schema(description = "防御策略模板名称", example = "自定义模板")
+    @Schema(description = "防御策略模板名称(事件类型)", example = "APT攻击事件")
     private String templateName;
 
     @JsonProperty("source_system")
@@ -21,6 +21,11 @@ public class Template {
     @Schema(description = "防御策略模板数据来源系统", example = "BW系统")
     private String sourceSystem;
 
+    @JsonProperty("description")
+    @NotNull(message = "source_system字段不能为空。")
+    @Schema(description = "对策略模板的文字描述。方便用户使用", example = "zd防护对象的全流量")
+    private String description;
+
     @JsonProperty("protect_level_low")
     @NotNull(message = "protect_level_low字段不能为空。")
     @Schema(description = "防御策略模板日常态字段提取选项")
diff --git a/src/main/java/com/realtime/protection/configuration/entity/task/TaskCommandInfo.java b/src/main/java/com/realtime/protection/configuration/entity/task/TaskCommandInfo.java
index 9a15837..803d8a2 100644
--- a/src/main/java/com/realtime/protection/configuration/entity/task/TaskCommandInfo.java
+++ b/src/main/java/com/realtime/protection/configuration/entity/task/TaskCommandInfo.java
@@ -103,4 +103,31 @@ public class TaskCommandInfo {
     @Schema(description = "指令所属任务的运行状态", accessMode = Schema.AccessMode.READ_ONLY)
     private Integer taskStatus;
 
+    // 复制构造函数
+    public void copyTaskCommandInfo(TaskCommandInfo original) {
+        this.UUID = original.UUID;
+        this.taskId = original.taskId;
+        this.ruleId = original.ruleId;
+        this.taskCreateUsername = original.taskCreateUsername;
+        this.taskCreateDepart = original.taskCreateDepart;
+        this.taskCreateUserId = original.taskCreateUserId;
+        this.taskName = original.taskName;
+        this.taskType = original.taskType;
+        this.taskAct = original.taskAct;
+        this.frequency = original.frequency;
+        this.startTime = original.startTime;
+        this.endTime = original.endTime;
+        this.isValid = original.isValid;
+        this.isJudged = original.isJudged;
+        this.fiveTupleWithMask = original.fiveTupleWithMask;
+        this.commandSentTimes = original.commandSentTimes;
+        this.commandSuccessTimes = original.commandSuccessTimes;
+        this.earliestSendTime = original.earliestSendTime;
+        this.latestSendTime = original.latestSendTime;
+        this.templateId = original.templateId;
+        this.protectLevel = original.protectLevel;
+        this.taskStatus = original.taskStatus;
+
+    }
+
 }
diff --git a/src/main/java/com/realtime/protection/server/alertmessage/AlertMessageService.java b/src/main/java/com/realtime/protection/server/alertmessage/AlertMessageService.java
index c905322..5987f1f 100644
--- a/src/main/java/com/realtime/protection/server/alertmessage/AlertMessageService.java
+++ b/src/main/java/com/realtime/protection/server/alertmessage/AlertMessageService.java
@@ -8,8 +8,10 @@ import com.realtime.protection.configuration.entity.task.TaskCommandInfo;
 import com.realtime.protection.configuration.utils.enums.StateEnum;
 import com.realtime.protection.configuration.utils.enums.TaskTypeEnum;
 import com.realtime.protection.server.command.CommandService;
+import lombok.Data;
 import org.springframework.stereotype.Service;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.UUID;
 
@@ -27,45 +29,46 @@ public class AlertMessageService {
     }
 
     @DSTransactional
-    public String processAlertMessage(AlertMessage alertMessage) {
-        TaskCommandInfo dynamicTaskCommandInfo = generateDynamicCommand(alertMessage);
+    public void processAlertMessage(AlertMessage alertMessage) {
+        //根据告警信息——>生成指令
+        List<TaskCommandInfo> dynamicTaskCommandInfoList = generateDynamicCommand(alertMessage);
+        //获取任务状态，设置指令的isValid字段，且是否生成指令入库（除了RUNING\PAUSED状态，其他都不入库）。
+        Integer taskStatus = dynamicTaskCommandInfoList.get(0).getTaskStatus();
+        //获取任务类型，设置指令的isJudged字段。
+        Integer taskType = dynamicTaskCommandInfoList.get(0).getTaskType();
 
-        Integer taskStatus = dynamicTaskCommandInfo.getTaskStatus();
-        Integer taskType = dynamicTaskCommandInfo.getTaskType();
-        String alertMessageUUID = null;
 
         if (taskType == TaskTypeEnum.DYNAMIC.getTaskType())//实时
             switch (StateEnum.getStateEnumByNum(taskStatus)) {
                 case RUNNING:
-                    alertMessageUUID = insertCommandAndAlertMessage(dynamicTaskCommandInfo, true, true, alertMessage);
+                    insertCommandAndAlertMessage(dynamicTaskCommandInfoList, true, true, alertMessage);
                     break;
                 case PAUSED:
-                    alertMessageUUID = insertCommandAndAlertMessage(dynamicTaskCommandInfo, false, true, alertMessage);
+                    insertCommandAndAlertMessage(dynamicTaskCommandInfoList, false, true, alertMessage);
                     break;
                 default://主要是stop
                     //command不入库
                     //alertmessage入库
-                    alertMessageUUID = insertAlertMessageOnly(alertMessage);
+                    insertAlertMessageOnly(alertMessage);
                     break;
             }
         else if (taskType == TaskTypeEnum.JUDGED.getTaskType())//研判后
             switch (StateEnum.getStateEnumByNum(taskStatus)) {
                 case RUNNING:
-                    alertMessageUUID = insertCommandAndAlertMessage(dynamicTaskCommandInfo, true, false, alertMessage);
+                    insertCommandAndAlertMessage(dynamicTaskCommandInfoList, true, false, alertMessage);
                     break;
                 case PAUSED:
-                    alertMessageUUID = insertCommandAndAlertMessage(dynamicTaskCommandInfo, false, false, alertMessage);
+                    insertCommandAndAlertMessage(dynamicTaskCommandInfoList, false, false, alertMessage);
                     break;
                 default://主要是stop
                     //command不入库
                     //alertmessage入库
-                    alertMessageUUID = insertAlertMessageOnly(alertMessage);
+                    insertAlertMessageOnly(alertMessage);
             }
-        return alertMessageUUID;
     }
 
 
-    private TaskCommandInfo generateDynamicCommand(AlertMessage alertMessage){
+    private List<TaskCommandInfo> generateDynamicCommand(AlertMessage alertMessage){
         Long taskId = alertMessage.getTaskId();
         Integer DynamicRuleId = alertMessage.getDynamicRuleId();
 //        查task信息
@@ -79,32 +82,44 @@ public class AlertMessageService {
         ProtectLevel templateProtectLevel = alertMessageMapper.queryTemplateProtectLevel(
                 dynamicCommandInfo.getTemplateId(),
                 dynamicCommandInfo.getProtectLevel());
+
         //根据策略模板和alertMessage中的FiveTupleWithMask生成要下发五元组信息
-        FiveTupleWithMask fiveTupleWithMaskNew =  updateFiveTupleWithMask(alertMessage.getFiveTupleWithMask(),
-                templateProtectLevel);
-        //指令加入策略模板筛选后的fiveTupleWithMaskNew
-        dynamicCommandInfo.setFiveTupleWithMask(fiveTupleWithMaskNew);
+        //根据策略模板的is_full_flow字段，如果是双向流量会生成两个fiveTuple，所以返回List
+        List<FiveTupleWithMask> fiveTupleWithMaskNew =  updateFiveTupleWithMask(alertMessage.getFiveTupleWithMask(),
+                alertMessage.getProtectIsSrcOrDst(), templateProtectLevel);
+
+        //根据fiveTuple生成动态指令信息
+        List<TaskCommandInfo> dynamicCommandInfoList = new ArrayList<TaskCommandInfo>();
+        if (fiveTupleWithMaskNew.size() == 2){
+            TaskCommandInfo dynamicCommandInfo_bi = new TaskCommandInfo();
+            dynamicCommandInfo_bi.copyTaskCommandInfo(dynamicCommandInfo);
+            dynamicCommandInfo_bi.setFiveTupleWithMask(fiveTupleWithMaskNew.get(1));
+            dynamicCommandInfoList.add(dynamicCommandInfo_bi);
+        }
+        dynamicCommandInfo.setFiveTupleWithMask(fiveTupleWithMaskNew.get(0));
+        dynamicCommandInfoList.add(dynamicCommandInfo);
 
-        return dynamicCommandInfo;
+        return dynamicCommandInfoList;
     }
 
     @DSTransactional
-    private String insertCommandAndAlertMessage(TaskCommandInfo dynamicTaskCommandInfo,
+    private void insertCommandAndAlertMessage(List<TaskCommandInfo> dynamicTaskCommandInfoList,
                                               Boolean isValid,
                                               Boolean isJudged,
                                               AlertMessage alertMessage){
-        //command入库
-        dynamicTaskCommandInfo.setIsValid(isValid);
-        dynamicTaskCommandInfo.setIsJudged(isJudged);
-        String commandUUID = commandService.createCommand(dynamicTaskCommandInfo);
-
-        //alertmessage入库
-        alertMessage.setCommandUUID(commandUUID);
-        String alertMessageUUID = UUID.randomUUID().toString();
-        alertMessage.setAlertMessageUUID(alertMessageUUID);
-        alertMessageMapper.insertAlertMessage(alertMessage);
+        for (TaskCommandInfo dynamicTaskCommandInfo : dynamicTaskCommandInfoList ){
+            //command入库
+            dynamicTaskCommandInfo.setIsValid(isValid);
+            dynamicTaskCommandInfo.setIsJudged(isJudged);
+            String commandUUID = commandService.createCommand(dynamicTaskCommandInfo);
+
+            //alertmessage入库
+            alertMessage.setCommandUUID(commandUUID);
+            String alertMessageUUID = UUID.randomUUID().toString();
+            alertMessage.setAlertMessageUUID(alertMessageUUID);
+            alertMessageMapper.insertAlertMessage(alertMessage);
+            }
 
-        return alertMessageUUID;
     }
     private String insertAlertMessageOnly(AlertMessage alertMessage){
         //alertmessage入库
@@ -119,31 +134,110 @@ public class AlertMessageService {
 
 
 
-    private FiveTupleWithMask updateFiveTupleWithMask(FiveTupleWithMask fiveTupleWithMask, ProtectLevel templateProtectLevel) {
-
-        FiveTupleWithMask newFiveTupleWithMask = new FiveTupleWithMask();
-        newFiveTupleWithMask.copyFiveTupleWithMask(fiveTupleWithMask);
-
-        if(!templateProtectLevel.getHasProtectObjectIP()){
-            newFiveTupleWithMask.setDestinationIP(null);
-            newFiveTupleWithMask.setMaskDestinationIP(null);
+    private List<FiveTupleWithMask> updateFiveTupleWithMask(FiveTupleWithMask fiveTupleWithMask,
+                                                      int protectIsSrcOrDst,
+                                                      ProtectLevel templateProtectLevel) {
+        //参数是告警信息的FiveTupleWithMask、防护对象是src还是dst、某个安全等级下的安全事件策略模板templateProtectLevel
+        //首先先从告警信息中获取protectObject和peer
+        @Data
+        class CommunicateObject {
+            private String IP;
+            private String maskIP;
+            private String Port;
+            private String maskPort;
+
+            public CommunicateObject(String IP,
+                                     String maskIP,
+                                     String Port,
+                                     String maskPort) {
+                this.IP = IP;
+                this.maskIP = maskIP;
+                this.Port = Port;
+                this.maskPort = maskPort;
+            }
+        }
+        CommunicateObject protectObject;
+        CommunicateObject peer;
+        if (protectIsSrcOrDst == 0) {
+            protectObject = new CommunicateObject(
+                    fiveTupleWithMask.getSourceIP(),
+                    fiveTupleWithMask.getMaskSourceIP(),
+                    fiveTupleWithMask.getSourcePort(),
+                    fiveTupleWithMask.getMaskSourcePort()
+            );
+            peer = new CommunicateObject(
+                    fiveTupleWithMask.getDestinationIP(),
+                    fiveTupleWithMask.getMaskDestinationIP(),
+                    fiveTupleWithMask.getDestinationPort(),
+                    fiveTupleWithMask.getMaskDestinationPort()
+            );
+        } else {
+            protectObject = new CommunicateObject(
+                    fiveTupleWithMask.getDestinationIP(),
+                    fiveTupleWithMask.getMaskDestinationIP(),
+                    fiveTupleWithMask.getDestinationPort(),
+                    fiveTupleWithMask.getMaskDestinationPort()
+            );
+            peer = new CommunicateObject(
+                    fiveTupleWithMask.getSourceIP(),
+                    fiveTupleWithMask.getMaskSourceIP(),
+                    fiveTupleWithMask.getSourcePort(),
+                    fiveTupleWithMask.getMaskSourcePort()
+            );
         }
-        if(!templateProtectLevel.getHasProtectObjectPort()){
-            newFiveTupleWithMask.setDestinationPort(null);
-            newFiveTupleWithMask.setMaskDestinationPort(null);
+        //根据模板抽取防护对象和对端需要的字段
+        if (!templateProtectLevel.getHasProtectObjectIP()) {
+            protectObject.setIP(null);
+            protectObject.setMaskIP(null);
         }
-        if(!templateProtectLevel.getHasPeerIP()){
-            newFiveTupleWithMask.setSourceIP(null);
-            newFiveTupleWithMask.setMaskSourceIP(null);
+        if (!templateProtectLevel.getHasProtectObjectPort()) {
+            protectObject.setPort(null);
+            protectObject.setMaskPort(null);
         }
-        if(!templateProtectLevel.getHasPeerPort()){
-            newFiveTupleWithMask.setSourcePort(null);
-            newFiveTupleWithMask.setMaskSourcePort(null);
+        if (!templateProtectLevel.getHasPeerIP()) {
+            peer.setIP(null);
+            peer.setMaskIP(null);
         }
-        if (!templateProtectLevel.getHasProtocol()) {
-            newFiveTupleWithMask.setProtocol(null);
-            newFiveTupleWithMask.setMaskProtocol(null);
+        if (!templateProtectLevel.getHasPeerPort()) {
+            peer.setPort(null);
+            peer.setMaskPort(null);
         }
+        List<FiveTupleWithMask> newFiveTupleWithMask = new ArrayList<FiveTupleWithMask>();
+        //生成指令
+        FiveTupleWithMask command1 = new FiveTupleWithMask();
+        command1.setSourceIP(peer.getIP());
+        command1.setMaskSourceIP(peer.getMaskIP());
+        command1.setSourcePort(peer.getPort());
+        command1.setMaskSourcePort(peer.getMaskPort());
+        command1.setDestinationIP(protectObject.getIP());
+        command1.setMaskDestinationIP(protectObject.getMaskIP());
+        command1.setSourcePort(protectObject.getPort());
+        command1.setMaskSourcePort(protectObject.getMaskPort());
+        if (templateProtectLevel.getHasProtocol()){
+            command1.setProtocol(fiveTupleWithMask.getProtocol());
+            command1.setProtocol(fiveTupleWithMask.getMaskProtocol());
+        }
+        newFiveTupleWithMask.add(command1);
+        //若需要处置全方向流量，再生成防护对象为源的规则
+        if(templateProtectLevel.getIsFullFlow()){
+            FiveTupleWithMask command2 = new FiveTupleWithMask();
+
+            command2.setSourceIP(protectObject.getIP());
+            command2.setMaskSourceIP(protectObject.getMaskIP());
+            command2.setSourcePort(protectObject.getPort());
+            command2.setMaskSourcePort(protectObject.getMaskPort());
+
+            command2.setDestinationIP(peer.getIP());
+            command2.setMaskDestinationIP(peer.getMaskIP());
+            command2.setSourcePort(peer.getPort());
+            command2.setMaskSourcePort(peer.getMaskPort());
+            if (templateProtectLevel.getHasProtocol()){
+                command2.setProtocol(fiveTupleWithMask.getProtocol());
+                command2.setProtocol(fiveTupleWithMask.getMaskProtocol());
+            }
+            newFiveTupleWithMask.add(command2);
+        }
+
         //目前告警信息还只是五元组，没有url、dns
         return newFiveTupleWithMask;
     }
diff --git a/src/main/resources/mappers/TemplateMapper.xml b/src/main/resources/mappers/TemplateMapper.xml
index d1fe826..cb36121 100644
--- a/src/main/resources/mappers/TemplateMapper.xml
+++ b/src/main/resources/mappers/TemplateMapper.xml
@@ -8,20 +8,22 @@
                                         strategy_template_low_level_id, strategy_template_medium_level_id,
                                         strategy_template_high_level_id,
                                         strategy_template_create_user_id, strategy_template_create_username,
-                                        strategy_template_create_depart)
+                                        strategy_template_create_depart,
+                                        strategy_template_description)
             VALUE (#{template.templateName}, #{template.sourceSystem},
                    #{template.protectLevelLow.protectLevelId}, #{template.protectLevelMedium.protectLevelId},
                    #{template.protectLevelHigh.protectLevelId},
-                   #{template.createUserId}, #{template.createUsername}, #{template.createDepart})
+                   #{template.createUserId}, #{template.createUsername}, #{template.createDepart},
+                    #{template.description})
     </insert>
 
     <insert id="newProtectLevel" useGeneratedKeys="true" keyProperty="protectLevelId">
         INSERT INTO t_protect_level(has_protect_object_ip, has_protect_object_port,
                                     has_protocol, has_url, has_dns,
-                                    has_peer_ip, has_peer_port)
+                                    has_peer_ip, has_peer_port, is_full_flow)
             VALUE (#{level.hasProtectObjectIP}, #{level.hasProtectObjectPort},
                    #{level.hasProtocol}, #{level.hasURL}, #{level.hasDNS},
-                   #{level.hasPeerIP}, #{level.hasPeerPort})
+                   #{level.hasPeerIP}, #{level.hasPeerPort}, #{level.isFullFlow})
     </insert>
 
     <resultMap id="templateMap" type="com.realtime.protection.configuration.entity.defense.template.Template">
@@ -33,6 +35,7 @@
 
         <result column="strategy_template_used_times" property="usedTimes"/>
         <result column="strategy_template_running_tasks" property="runningTasks"/>
+        <result column="strategy_template_description" property="description"/>
 
         <association property="protectLevelLow"
                      javaType="com.realtime.protection.configuration.entity.defense.template.ProtectLevel">
@@ -44,6 +47,7 @@
             <result column="low_has_protocol" property="hasProtocol"/>
             <result column="low_has_url" property="hasURL"/>
             <result column="low_has_dns" property="hasDNS"/>
+            <result column="low_is_full_flow" property="isFullFlow"/>
         </association>
 
         <association property="protectLevelMedium"
@@ -56,6 +60,7 @@
             <result column="medium_has_protocol" property="hasProtocol"/>
             <result column="medium_has_url" property="hasURL"/>
             <result column="medium_has_dns" property="hasDNS"/>
+            <result column="medium_is_full_flow" property="isFullFlow"/>
         </association>
 
         <association property="protectLevelHigh"
@@ -69,6 +74,7 @@
             <result column="high_has_protocol" property="hasProtocol"/>
             <result column="high_has_url" property="hasURL"/>
             <result column="high_has_dns" property="hasDNS"/>
+            <result column="high_is_full_flow" property="isFullFlow"/>
         </association>
     </resultMap>
 
@@ -81,6 +87,7 @@
         <result column="has_protocol" property="hasProtocol"/>
         <result column="has_url" property="hasURL"/>
         <result column="has_dns" property="hasDNS"/>
+        <result column="is_full_flow" property="isFullFlow"/>
     </resultMap>
 
     <select id="queryTemplates" resultMap="templateMap">
diff --git a/src/test/java/com/realtime/protection/server/defense/template/TemplateServiceTest.java b/src/test/java/com/realtime/protection/server/defense/template/TemplateServiceTest.java
index d25e4f5..cce2ea4 100644
--- a/src/test/java/com/realtime/protection/server/defense/template/TemplateServiceTest.java
+++ b/src/test/java/com/realtime/protection/server/defense/template/TemplateServiceTest.java
@@ -44,6 +44,7 @@ class TemplateServiceTest extends ProtectionApplicationTests {
 
         ProtectLevel protectLevelHigh = new ProtectLevel();
         protectLevelHigh.setHasDNS(true);
+        protectLevelHigh.setIsFullFlow(true);
 
         template.setProtectLevelLow(protectLevelLow);
         template.setProtectLevelMedium(protectLevelMedium);
@@ -67,7 +68,7 @@ class TemplateServiceTest extends ProtectionApplicationTests {
             System.out.println(e.getMessage());
         }
 
-        for (int i = 0; i < 100; i++) {
+        for (int i = 0; i < 10; i++) {
             assertThrows(DuplicateKeyException.class, () -> {
                 Integer templateId = templateService.newTemplate(template);
                 assertTrue(templateId > 0);