1. 新增AuditStatusValidator类，用以作为审核状态机

2. 新建DataListener，用以读取excel文件 3. 完成防护对象配置所有接口 4. 添加SqlSessionWrapper类用以进行批处理 5. ProtectObject类添加更多校验（IP、大小等）
2024-01-05 21:42:19 +08:00
parent 776c7c0f6d
commit 0fb8dd87fe
15 changed files with 468 additions and 21 deletions
--- a/src/main/java/com/realtime/protection/configuration/entity/defense/object/ProtectObject.java
+++ b/src/main/java/com/realtime/protection/configuration/entity/defense/object/ProtectObject.java
@@ -1,7 +1,10 @@
 package com.realtime.protection.configuration.entity.defense.object;

 import com.fasterxml.jackson.annotation.JsonProperty;
+import jakarta.validation.constraints.Max;
+import jakarta.validation.constraints.Min;
 import jakarta.validation.constraints.NotNull;
+import jakarta.validation.constraints.Pattern;
 import lombok.Data;

@Data
@@ -17,10 +20,13 @@ public class ProtectObject {
    private String protectObjectSystemName;

    @JsonProperty("proobj_ip_address")
+    @Pattern(regexp = "^(\\d{1,3})\\.(\\d{1,3})\\.(\\d{1,3})\\.(\\d{1,3})$", message = "Invalid IPv4 Address")
    private String protectObjectIPAddress;

    @JsonProperty("proobj_port")
    @NotNull(message = "proobj_port should not be empty.")
+    @Max(value = 65535, message = "port should not be more than 65535")
+    @Min(value = 1, message = "port should not be less than 1")
    private Integer protectObjectPort;

    @JsonProperty("proobj_url")
--- a/src/main/java/com/realtime/protection/configuration/entity/task/Task.java
+++ b/src/main/java/com/realtime/protection/configuration/entity/task/Task.java
@@ -1,6 +1,7 @@
 package com.realtime.protection.configuration.entity.task;

 import com.fasterxml.jackson.annotation.JsonProperty;
+import jakarta.validation.constraints.Future;
 import jakarta.validation.constraints.NotNull;
 import lombok.Data;

@@ -13,15 +14,16 @@ public class Task {
    private Integer taskId;

    @JsonProperty("task_name")
-    @NotNull(message = "task_name should not be empty.")
+    @NotNull(message = "task_name should not be empty. ")
    private String taskName;

    @JsonProperty("task_start_time")
-    @NotNull(message = "task_start_time should not be empty.")
+    @NotNull(message = "task_start_time should not be empty. ")
    private LocalDateTime taskStartTime;

    @JsonProperty("task_end_time")
-    @NotNull(message = "task_end_time should not be empty.")
+    @NotNull(message = "task_end_time should not be empty. ")
+    @Future(message = "task_end_time should be a future time. ")
    private LocalDateTime taskEndTime;

    @JsonProperty("task_create_time")
@@ -31,11 +33,11 @@ public class Task {
    private LocalDateTime taskModifyTime;

    @JsonProperty("task_type")
-    @NotNull(message = "task_type should not be empty.")
+    @NotNull(message = "task_type should not be empty. ")
    private String taskType;

    @JsonProperty("task_act")
-    @NotNull(message = "task_act should not be empty.")
+    @NotNull(message = "task_act should not be empty. ")
    private String taskAct;

    @JsonProperty("task_create_username")
--- a/src/main/java/com/realtime/protection/configuration/exception/GlobalExceptionHandler.java
+++ b/src/main/java/com/realtime/protection/configuration/exception/GlobalExceptionHandler.java
@@ -48,7 +48,7 @@ public class GlobalExceptionHandler {
    }

    @Order(2)
-    @ExceptionHandler(value = HandlerMethodValidationException.class)
+    @ExceptionHandler(value = {HandlerMethodValidationException.class, IllegalArgumentException.class})
    public ResponseResult handleHandlerMethodValidationException(HandlerMethodValidationException e) {
        return ResponseResult.invalid().setMessage(e.getMessage());
    }
--- a/src/main/java/com/realtime/protection/configuration/response/ResponseResult.java
+++ b/src/main/java/com/realtime/protection/configuration/response/ResponseResult.java
@@ -4,13 +4,14 @@ import lombok.Data;

 import java.io.Serializable;
 import java.util.LinkedHashMap;
+import java.util.Map;

@Data
 public class ResponseResult implements Serializable {

    private int code;
    private String message;
-    private LinkedHashMap<String, Object> data;
+    private Map<String, Object> data;

    public ResponseResult(int code, String message, LinkedHashMap<String, Object> data) {
        this.code = code;
@@ -45,6 +46,10 @@ public class ResponseResult implements Serializable {
        return new ResponseResult(400, "invalid request");
    }

+    public static ResponseResult invalid(String message) {
+        return new ResponseResult(400, message);
+    }
+
    public static ResponseResult unAuthorized() {
        return new ResponseResult(401, "UnAuthorized User");
    }
@@ -68,7 +73,7 @@ public class ResponseResult implements Serializable {
        return this;
    }

-    public ResponseResult setDataMap(LinkedHashMap<String, Object> data) {
+    public ResponseResult setDataMap(Map<String, Object> data) {
        this.data = data;
        return this;
    }
--- a/src/main/java/com/realtime/protection/configuration/utils/AuditStatusValidator.java
+++ b/src/main/java/com/realtime/protection/configuration/utils/AuditStatusValidator.java
@@ -0,0 +1,22 @@
+package com.realtime.protection.configuration.utils;
+
+public class AuditStatusValidator {
+
+    private final Integer auditStatusOriginal;
+
+    public AuditStatusValidator(Integer auditStatusOriginal) {
+        this.auditStatusOriginal = auditStatusOriginal;
+    }
+
+    public static AuditStatusValidator setOriginal(Integer auditStatusOriginal) {
+        return new AuditStatusValidator(auditStatusOriginal);
+    }
+
+    public Boolean checkValidate(Integer auditStatusNow) {
+        switch (auditStatusNow) {
+            case 0, 1 -> {return auditStatusOriginal != 2;}
+            case 2 -> {return auditStatusOriginal != 1;}
+            default -> {return false;}
+        }
+    }
+}
--- a/src/main/java/com/realtime/protection/configuration/utils/SqlSessionWrapper.java
+++ b/src/main/java/com/realtime/protection/configuration/utils/SqlSessionWrapper.java
@@ -0,0 +1,37 @@
+package com.realtime.protection.configuration.utils;
+
+import org.apache.ibatis.session.ExecutorType;
+import org.apache.ibatis.session.SqlSession;
+import org.apache.ibatis.session.SqlSessionFactory;
+import org.springframework.stereotype.Component;
+
+import java.util.function.Function;
+
+@Component
+public class SqlSessionWrapper {
+
+    private final SqlSessionFactory sqlSessionFactory;
+
+    public SqlSessionWrapper(SqlSessionFactory sqlSessionFactory) {
+        this.sqlSessionFactory = sqlSessionFactory;
+    }
+
+    public <M, I, O> O startBatchSession(Class<M> mapperClass,
+                                         Function<M, Function<I, O>> wrappedFunction,
+                                         I arguments) {
+        SqlSession sqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH, false);
+        M mapper = sqlSession.getMapper(mapperClass);
+        try {
+            O result = wrappedFunction.apply(mapper).apply(arguments);
+            sqlSession.commit();
+            sqlSession.clearCache();
+
+            return result;
+        } catch (Exception e) {
+            sqlSession.rollback();
+            throw e;
+        } finally {
+            sqlSession.close();
+        }
+    }
+}
--- a/src/main/java/com/realtime/protection/server/defense/object/ProjectObjectDataListener.java
+++ b/src/main/java/com/realtime/protection/server/defense/object/ProjectObjectDataListener.java
@@ -0,0 +1,42 @@
+package com.realtime.protection.server.defense.object;
+
+
+import com.alibaba.excel.context.AnalysisContext;
+import com.alibaba.excel.read.listener.ReadListener;
+import com.alibaba.excel.util.ListUtils;
+import com.realtime.protection.configuration.entity.defense.object.ProtectObject;
+import org.springframework.beans.factory.annotation.Autowired;
+
+import java.util.List;
+
+public class ProjectObjectDataListener implements ReadListener<ProtectObject> {
+
+    private final ProtectObjectService protectObjectService;
+    private final List<ProtectObject> cachedDataList = ListUtils.newArrayListWithExpectedSize(batchCount);
+    private static final int batchCount = 100;
+
+    public ProjectObjectDataListener(ProtectObjectService protectObjectService) {
+        this.protectObjectService = protectObjectService;
+    }
+
+    @Override
+    public void invoke(ProtectObject protectObject, AnalysisContext analysisContext) {
+        cachedDataList.add(protectObject);
+        if (cachedDataList.size() > 1000) {
+            saveData();
+            cachedDataList.clear();
+        }
+    }
+
+    @Override
+    public void doAfterAllAnalysed(AnalysisContext analysisContext) {
+        saveData();
+    }
+
+    private void saveData() {
+        Boolean success = protectObjectService.newProtectObjects(cachedDataList);
+        if (!success) {
+            throw new RuntimeException("Error reading data in newProtectObjects");
+        }
+    }
+}
--- a/src/main/java/com/realtime/protection/server/defense/object/ProtectObjectController.java
+++ b/src/main/java/com/realtime/protection/server/defense/object/ProtectObjectController.java
@@ -1,4 +1,104 @@
 package com.realtime.protection.server.defense.object;

+import com.alibaba.excel.EasyExcel;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.realtime.protection.configuration.entity.defense.object.ProtectObject;
+import com.realtime.protection.configuration.response.ResponseResult;
+import jakarta.validation.Valid;
+import jakarta.validation.constraints.Min;
+import org.springframework.web.bind.annotation.*;
+import org.springframework.web.multipart.MultipartFile;
+
+import java.io.IOException;
+import java.util.List;
+
+@RestController
+@RequestMapping("/proobj")
 public class ProtectObjectController {
+
+    private final ProtectObjectService protectObjectService;
+
+    public ProtectObjectController(ProtectObjectService protectObjectService) {
+        this.protectObjectService = protectObjectService;
+    }
+
+    @PostMapping("/new")
+    public ResponseResult newProtectObject(@RequestBody @Valid ProtectObject protectObject) {
+        Integer protectObjectId = protectObjectService.newProtectObject(protectObject);
+        if (protectObjectId == 0) {
+            return ResponseResult.error().setMessage("failed to create a protect object")
+                    .setData("proobj_name", protectObject.getProtectObjectName())
+                    .setData("proobj_id", protectObjectId)
+                    .setData("success", false);
+        }
+
+        return ResponseResult.ok()
+                .setData("proobj_name", protectObject.getProtectObjectName())
+                .setData("proobj_id", protectObject.getProtectObjectId())
+                .setData("success", true);
+    }
+
+    @PutMapping("/new")
+    public ResponseResult newProtectObjectFromFile(MultipartFile updateFile) throws IOException {
+        EasyExcel.read(updateFile.getInputStream(), ProtectObject.class,
+                new ProjectObjectDataListener(protectObjectService)).sheet().doRead();
+        return ResponseResult.ok();
+    }
+
+    @GetMapping("/query")
+    public ResponseResult queryProtectObjects(@RequestParam(value = "proobj_name", required = false)
+                                              String protectObjectName,
+                                              @RequestParam(value = "proobj_id", required = false) @Min(1)
+                                              Integer protectObjectId,
+                                              @RequestParam("page") @Min(1) Integer page,
+                                              @RequestParam("page_size") @Min(1) Integer pageSize) {
+        return ResponseResult.ok()
+                .setData("proobj_list", protectObjectService.queryProtectObjects(protectObjectName,
+                        protectObjectId, page, pageSize));
+
+    }
+
+    @GetMapping("/{id}/query")
+    public ResponseResult queryProtectObject(@PathVariable("id") Integer protectObjectId) {
+        ProtectObject protectObject = protectObjectService.queryProtectObject(protectObjectId);
+        return ResponseResult.ok()
+                .setData("proobj_id", protectObject.getProtectObjectId())
+                .setData("proobj_name", protectObject.getProtectObjectName())
+                .setData("proobj_system_name", protectObject.getProtectObjectSystemName())
+                .setData("proobj_ip_address", protectObject.getProtectObjectIPAddress())
+                .setData("proobj_port", protectObject.getProtectObjectPort())
+                .setData("proobj_url", protectObject.getProtectObjectURL())
+                .setData("proobj_protocol", protectObject.getProtectObjectProtocol());
+    }
+
+    @PostMapping("/{id}/update")
+    public ResponseResult updateProtectObject(@PathVariable("id") Integer protectObjectId,
+                                              @RequestBody @Valid ProtectObject protectObject) {
+        protectObject.setProtectObjectId(protectObjectId);
+        return ResponseResult.ok()
+                .setData("proobj_id", protectObject.getProtectObjectId())
+                .setData("success", protectObjectService.updateProtectObject(protectObject));
+    }
+
+    @DeleteMapping("/{id}/delete")
+    public ResponseResult deleteProtectObject(@PathVariable("id") Integer protectObjectId) {
+        return ResponseResult.ok()
+                .setData("proobj_id", protectObjectId)
+                .setData("success", protectObjectService.deleteProtectObject(protectObjectId));
+    }
+
+    @PostMapping("/delete")
+    public ResponseResult deleteProtectObject(@RequestBody @JsonProperty("proobj_ids") List<Integer> protectObjectIds) {
+        return ResponseResult.ok()
+                .setData("proobj_ids", protectObjectIds)
+                .setData("success", protectObjectService.deleteProtectObjects(protectObjectIds));
+    }
+
+    @PostMapping("/{id}/audit/{status}")
+    public ResponseResult changeProtectObjectAuditStatus(@PathVariable("id") Integer protectObjectId,
+                                                         @PathVariable("status") Integer auditStatus) {
+        return ResponseResult.ok()
+                .setDataMap(protectObjectService.changeProtectObjectAuditStatus(protectObjectId, auditStatus))
+                .setData("proobj_id", protectObjectId);
+    }
 }
--- a/src/main/java/com/realtime/protection/server/defense/object/ProtectObjectMapper.java
+++ b/src/main/java/com/realtime/protection/server/defense/object/ProtectObjectMapper.java
@@ -10,6 +10,8 @@ import java.util.List;
 public interface ProtectObjectMapper {
    void newProtectObject(@Param("proobj") ProtectObject protectObject);

+    void newProtectObjects(@Param("proobjs") List<ProtectObject> protectObject);
+
    List<ProtectObject> queryProtectObjects(@Param("proobj_name") String protectObjectName,
                                            @Param("proobj_id") Integer protectObjectId,
                                            @Param("page") Integer page,
@@ -21,5 +23,8 @@ public interface ProtectObjectMapper {

    Boolean deleteProtectObject(@Param("proobj_id") Integer protectObjectId);

-    Boolean changeProtectObjectAuditStatus(@Param("proobj_audit_status") Integer protectObjectAuditStatus);
+    Boolean deleteProtectObjects(@Param("proobj_ids") List<Integer> protectObjectIds);
+
+    Boolean changeProtectObjectAuditStatus(@Param("proobj_id") Integer protectObjectId,
+                                           @Param("proobj_audit_status") Integer protectObjectAuditStatus);
 }
--- a/src/main/java/com/realtime/protection/server/defense/object/ProtectObjectService.java
+++ b/src/main/java/com/realtime/protection/server/defense/object/ProtectObjectService.java
@@ -1,18 +1,26 @@
 package com.realtime.protection.server.defense.object;

+import com.alibaba.excel.util.ListUtils;
 import com.realtime.protection.configuration.entity.defense.object.ProtectObject;
+import com.realtime.protection.configuration.utils.AuditStatusValidator;
+import com.realtime.protection.configuration.utils.SqlSessionWrapper;
 import org.springframework.stereotype.Service;
 import org.springframework.transaction.annotation.Transactional;

+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
 import java.util.stream.Collectors;

@Service
 public class ProtectObjectService {
    private final ProtectObjectMapper protectObjectMapper;
+    private final SqlSessionWrapper sqlSessionWrapper;

-    public ProtectObjectService(ProtectObjectMapper protectObjectMapper) {
+    public ProtectObjectService(ProtectObjectMapper protectObjectMapper, SqlSessionWrapper sqlSessionWrapper) {
        this.protectObjectMapper = protectObjectMapper;
+        this.sqlSessionWrapper = sqlSessionWrapper;
    }

    public Integer newProtectObject(ProtectObject protectObject) {
@@ -24,10 +32,22 @@ public class ProtectObjectService {
        return protectObject.getProtectObjectId();
    }

-    @Transactional
-    public List<Integer> newProtectObjects(List<ProtectObject> protectObjectList) {
-        protectObjectList.forEach(protectObjectMapper::newProtectObject);
-        return protectObjectList.stream().map(ProtectObject::getProtectObjectId).collect(Collectors.toList());
+    public Boolean newProtectObjects(List<ProtectObject> protectObjectList) {
+        Function<ProtectObjectMapper, Function<List<ProtectObject>, Boolean>> newProtectObjectFunction = mapper -> list -> {
+            List<ProtectObject> protectObjectBatch = ListUtils.newArrayListWithExpectedSize(100);
+            for (ProtectObject protectObject : protectObjectList) {
+                protectObjectBatch.add(protectObject);
+                if (protectObjectBatch.size() < 1000) {
+                    continue;
+                }
+                mapper.newProtectObjects(protectObjectBatch);
+                protectObjectBatch.clear();
+            }
+            mapper.newProtectObjects(protectObjectBatch);
+            return true;
+        };
+
+        return sqlSessionWrapper.startBatchSession(ProtectObjectMapper.class, newProtectObjectFunction, protectObjectList);
    }

    public List<ProtectObject> queryProtectObjects(String protectObjectName, Integer protectObjectId, Integer page, Integer pageSize) {
@@ -46,8 +66,42 @@ public class ProtectObjectService {
        return protectObjectMapper.deleteProtectObject(protectObjectId);
    }

-    @Transactional
    public Boolean deleteProtectObjects(List<Integer> protectObjectIds) {
-        return protectObjectIds.stream().allMatch(protectObjectMapper::deleteProtectObject);
+        Function<ProtectObjectMapper, Function<List<Integer>, Void>> deleteProtectObjectFunction = mapper -> list -> {
+            if (list == null || list.isEmpty()) {
+                return null;
+            }
+
+            List<Integer> protectObjectBatch = ListUtils.newArrayListWithExpectedSize(100);
+            for (Integer protectObjectId : list) {
+                protectObjectBatch.add(protectObjectId);
+                if (protectObjectBatch.size() < 100) {
+                    continue;
+                }
+                mapper.deleteProtectObjects(protectObjectIds);
+                protectObjectBatch.clear();
+            }
+            mapper.deleteProtectObjects(protectObjectBatch);
+            return null;
+        };
+
+        sqlSessionWrapper.startBatchSession(ProtectObjectMapper.class, deleteProtectObjectFunction, protectObjectIds);
+        return true;
+    }
+
+    @Transactional
+    public Map<String, Object> changeProtectObjectAuditStatus(Integer protectObjectId, Integer auditStatus) {
+        Integer originalAuditStatus = protectObjectMapper.queryProtectObject(protectObjectId).getProtectObjectAuditStatus();
+        if (!AuditStatusValidator.setOriginal(originalAuditStatus).checkValidate(auditStatus)) {
+            throw new IllegalArgumentException("invalid audit status");
+        }
+        Boolean success = protectObjectMapper.changeProtectObjectAuditStatus(protectObjectId, auditStatus);
+
+        Integer auditStatusNow = protectObjectMapper.queryProtectObject(protectObjectId).getProtectObjectAuditStatus();
+        Map<String, Object> resultMap = new HashMap<>();
+        resultMap.put("success", success);
+        resultMap.put("audit_status", auditStatusNow);
+
+        return resultMap;
    }
 }
--- a/src/main/java/com/realtime/protection/server/task/TaskService.java
+++ b/src/main/java/com/realtime/protection/server/task/TaskService.java
@@ -82,11 +82,11 @@ public class TaskService {

        taskMapper.clearTaskProtectObjectConcat(task.getTaskId());
        taskMapper.clearTaskConnectedStaticRule(task.getTaskId());
-        taskMapper.clearTaskConnectedDynamicRule(task.getTaskId());
+//        taskMapper.clearTaskConnectedDynamicRule(task.getTaskId());

        task.getProtectObjectIds().forEach(
                proobjId -> taskMapper.newTaskProtectObjectConcat(task.getTaskId(), proobjId));
        taskMapper.newTaskStaticRuleConcat(task.getTaskId(), task.getStaticRuleIds());
-        taskMapper.newTaskDynamicRuleConcat(task.getTaskId(), task.getDynamicRuleIds());
+//        taskMapper.newTaskDynamicRuleConcat(task.getTaskId(), task.getDynamicRuleIds());
    }
 }