1. application.yml修改为application-dev.yml和application-prod.yml

2. 添加更多Exception拦截器
3. 编写状态模式处理task状态的更改
4. 添加StateChangeService,用以处理所有任务状态转换相关的内容
5. 添加StateEnum, ProtocolEnum,TaskTypeEnum用以处理任务和协议相关的所有状态和类型
This commit is contained in:
EnderByEndera
2024-01-11 19:49:07 +08:00
parent 930ba8b5ac
commit 0f712618f2
70 changed files with 1209 additions and 400 deletions

View File

@@ -4,8 +4,6 @@ import com.fasterxml.jackson.annotation.JsonProperty;
import jakarta.validation.constraints.NotNull;
import lombok.Data;
import java.util.List;
@Data
public class Template {
@JsonProperty("template_id")
@@ -37,6 +35,12 @@ public class Template {
@NotNull(message = "protect_level_high should not be empty. ")
private ProtectLevel protectLevelHigh;
@JsonProperty("template_used_times")
private Integer usedTimes;
@JsonProperty("running_tasks")
private Integer runningTasks;
private Integer createUserId;
private String createUsername;

View File

@@ -26,7 +26,7 @@ public class DynamicRuleObject {
@JsonProperty("dynamic_rule_create_username")
private String dynamicRuleCreateUsername;
// @JsonProperty("dynamic_rule_audit_status")
// @JsonProperty("dynamic_rule_audit_status")
// private Integer dynamicRuleAuditStatus;
@JsonProperty("dynamic_rule_create_depart")
private String dynamicRuleCreateDepart;

View File

@@ -1,25 +1,32 @@
package com.realtime.protection.configuration.entity.task;
import com.realtime.protection.configuration.utils.enums.ProtocolEnum;
import lombok.Data;
import java.time.LocalDateTime;
@Data
public class Command {
private Integer id;
private FiveTupleWithMask fiveTupleWithMask;
private Long taskId;
private Integer type;
private String sourceIP;
private String sourcePort;
private String destinationIP;
private String destinationPort;
private Integer protocol;
private String operation;
private LocalDateTime validTime;
private LocalDateTime invalidTime;
private String maskSourceIP;
private String maskSourcePort;
private String maskDestinationIP;
private String maskDestinationPort;
private Integer direction;
public static Command generateCommand(TaskCommandInfo info, LocalDateTime validTime) {
Command command = new Command();
private LocalDateTime datetime;
FiveTupleWithMask fiveTupleWithMask = info.getFiveTupleWithMask();
if (fiveTupleWithMask.getProtocol() != null)
fiveTupleWithMask.setProtocolNum(ProtocolEnum.valueOf(fiveTupleWithMask.getProtocol()).getProtocolNumber());
command.setFiveTupleWithMask(fiveTupleWithMask);
command.setTaskId(info.getTaskId());
command.setOperation(info.getOperation());
command.setValidTime(validTime);
command.setInvalidTime(info.getEndTime());
return command;
}
}

View File

@@ -0,0 +1,19 @@
package com.realtime.protection.configuration.entity.task;
import lombok.Data;
@Data
public class FiveTupleWithMask {
private Integer addrType;
private String sourceIP;
private String sourcePort;
private String destinationIP;
private String destinationPort;
private String protocol;
private Integer protocolNum;
private String maskSourceIP;
private String maskSourcePort;
private String maskDestinationIP;
private String maskDestinationPort;
}

View File

@@ -11,7 +11,7 @@ import java.util.List;
@Data
public class Task {
@JsonProperty("task_id")
private Integer taskId;
private Long taskId;
@JsonProperty("task_name")
@NotNull(message = "task_name should not be empty. ")
@@ -19,6 +19,7 @@ public class Task {
@JsonProperty("task_start_time")
@NotNull(message = "task_start_time should not be empty. ")
@Future(message = "task_start_time should be a future time")
private LocalDateTime taskStartTime;
@JsonProperty("task_end_time")
@@ -34,7 +35,7 @@ public class Task {
@JsonProperty("task_type")
@NotNull(message = "task_type should not be empty. ")
private String taskType;
private Integer taskType;
@JsonProperty("task_act")
@NotNull(message = "task_act should not be empty. ")
@@ -47,13 +48,13 @@ public class Task {
private String taskCreateDepart;
@JsonProperty("task_create_userid")
private Integer taskCreateUserId;
private Long taskCreateUserId;
@JsonProperty("static_rule_ids")
private List<Integer> staticRuleIds;
private List<Long> staticRuleIds;
@JsonProperty("dynamic_rule_ids")
private List<Integer> dynamicRuleIds;
private List<Long> dynamicRuleIds;
@JsonProperty("task_status")
private Integer taskStatus;

View File

@@ -0,0 +1,19 @@
package com.realtime.protection.configuration.entity.task;
import lombok.Data;
import java.time.LocalDateTime;
@Data
public class TaskCommandInfo {
private FiveTupleWithMask fiveTupleWithMask;
private Long taskId;
private Long ruleId;
// 额外字段
private String operation;
private Integer frequency;
private LocalDateTime startTime;
private LocalDateTime endTime;
}

View File

@@ -0,0 +1,19 @@
package com.realtime.protection.configuration.exception;
public class DorisStartException extends Exception {
public Long taskId;
public DorisStartException(Exception e, Long taskId) {
super(e.getMessage(), e.getCause());
this.taskId = taskId;
}
public DorisStartException(Exception e) {
super(e.getMessage(), e.getCause());
}
public DorisStartException(String message, Long taskId) {
super(message);
this.taskId = taskId;
}
}

View File

@@ -3,6 +3,8 @@ package com.realtime.protection.configuration.exception;
import cn.dev33.satoken.exception.NotLoginException;
import cn.dev33.satoken.exception.SaTokenException;
import com.realtime.protection.configuration.response.ResponseResult;
import com.realtime.protection.configuration.utils.enums.StateEnum;
import com.realtime.protection.server.task.status.StateChangeService;
import org.apache.ibatis.exceptions.PersistenceException;
import org.springframework.context.support.DefaultMessageSourceResolvable;
import org.springframework.core.annotation.Order;
@@ -16,20 +18,18 @@ import java.util.stream.Collectors;
@RestControllerAdvice
public class GlobalExceptionHandler {
private final StateChangeService stateChangeService;
public GlobalExceptionHandler(StateChangeService stateChangeService) {
this.stateChangeService = stateChangeService;
}
@Order(3)
@ExceptionHandler(value = Exception.class)
public ResponseResult handleGlobalException(Exception e) {
return ResponseResult.error().setMessage(e.getMessage());
}
@Order(2)
@ExceptionHandler(value = NotLoginException.class)
public ResponseResult handleNotLoginException(NotLoginException e) {
return new ResponseResult(
401,
e.getMessage()
);
}
@Order(2)
@ExceptionHandler(value = PersistenceException.class)
@@ -48,14 +48,42 @@ public class GlobalExceptionHandler {
}
@Order(2)
@ExceptionHandler(value = {HandlerMethodValidationException.class, IllegalArgumentException.class})
public ResponseResult handleHandlerMethodValidationException(HandlerMethodValidationException e) {
@ExceptionHandler(value = {
HandlerMethodValidationException.class,
IllegalArgumentException.class,
IllegalStateException.class
})
public ResponseResult handleHandlerMethodValidationException(Exception e) {
return ResponseResult.invalid().setMessage(e.getMessage());
}
@Order(2)
@ExceptionHandler(value = NotLoginException.class)
public ResponseResult handleNotLoginException(NotLoginException e) {
return new ResponseResult(
401,
e.getMessage()
);
}
@Order(2)
@ExceptionHandler(value = SaTokenException.class)
public ResponseResult handleSaTokenException(SaTokenException e) {
return ResponseResult.unAuthorized().setMessage(e.getMessage());
}
@Order(2)
@ExceptionHandler(value = DorisStartException.class)
public ResponseResult handleDorisStartException(DorisStartException e) {
ResponseResult responseResult = ResponseResult.error()
.setMessage("Doris command creation meets error: " + e.getMessage());
try {
stateChangeService.changeState(StateEnum.FAILED.getStateNum(), e.taskId);
} catch (Exception another) {
responseResult.setAnother(ResponseResult.error().setMessage(e.getMessage()));
}
return responseResult;
}
}

View File

@@ -12,6 +12,7 @@ public class ResponseResult implements Serializable {
private int code;
private String message;
private Map<String, Object> data;
private ResponseResult another;
public ResponseResult(int code, String message, LinkedHashMap<String, Object> data) {
this.code = code;

View File

@@ -0,0 +1,5 @@
package com.realtime.protection.configuration.satoken;
public interface Nameable {
String name();
}

View File

@@ -1,11 +1,9 @@
package com.realtime.protection.configuration.satoken;
import cn.dev33.satoken.interceptor.SaInterceptor;
import cn.dev33.satoken.router.SaRouter;
import cn.dev33.satoken.stp.StpUtil;
import com.realtime.protection.configuration.satoken.permission.Permission;
import com.realtime.protection.configuration.satoken.permission.SystemConfiguration;
import com.realtime.protection.configuration.satoken.permission.WhiteList;
import com.realtime.protection.configuration.satoken.role.Role;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
@@ -15,13 +13,15 @@ public class SaTokenConfigure implements WebMvcConfigurer {
@Override
public void addInterceptors(InterceptorRegistry registry) {
registry.addInterceptor(new SaInterceptor(handler -> {
SaRouter.match("/whiteobj/new", r ->
this.checkPermissions(SystemConfiguration.NEW, WhiteList.NEW));
SaRouter.match("/whiteobj/update", r ->
this.checkPermissions(SystemConfiguration.UPDATE, WhiteList.UPDATE));
}))
.addPathPatterns("/**")
.excludePathPatterns("/user/doLogin");
// SaRouter.match("/**")
// .notMatch("/user/doLogin")
// .check(r -> StpUtil.checkLogin());
// SaRouter.match("/whiteobj/new", r ->
// this.checkPermissions(WhiteList.NEW));
// SaRouter.match("/whiteobj/{id}/update", r ->
// this.checkPermissions(WhiteList.UPDATE));
})).addPathPatterns("/**");
}
void checkPermissions(Permission... permissions) {
@@ -30,5 +30,10 @@ public class SaTokenConfigure implements WebMvcConfigurer {
}
}
void checkRole(Role... roles) {
for (Role role : roles) {
StpUtil.checkRole(role.getName());
}
}
}

View File

@@ -1,5 +0,0 @@
package com.realtime.protection.configuration.satoken.permission;
public interface Nameable {
String name();
}

View File

@@ -1,5 +1,7 @@
package com.realtime.protection.configuration.satoken.permission;
import com.realtime.protection.configuration.satoken.Nameable;
public interface Permission extends Nameable {
default String getName() {
return this.getClass().getSimpleName() + ":" + this.name();

View File

@@ -0,0 +1,5 @@
package com.realtime.protection.configuration.satoken.role;
public enum Admin implements Role {
ADMIN
}

View File

@@ -0,0 +1,10 @@
package com.realtime.protection.configuration.satoken.role;
import com.realtime.protection.configuration.satoken.Nameable;
public interface Role extends Nameable {
default String getName() {
return this.getClass().getSimpleName() + ":" + this.name();
}
}

View File

@@ -0,0 +1,18 @@
package com.realtime.protection.configuration.utils;
import java.lang.reflect.Field;
import java.util.HashMap;
import java.util.Map;
public class EntityUtils {
public static Map<String, Object> entityToMap(Object object) throws IllegalAccessException {
Map<String, Object> resultMap = new HashMap<>();
for (Field field : object.getClass().getDeclaredFields()) {
field.setAccessible(true);
Object o = field.get(object);
resultMap.put(field.getName(), o);
}
return resultMap;
}
}

View File

@@ -16,13 +16,15 @@ public class SqlSessionWrapper {
this.sqlSessionFactory = sqlSessionFactory;
}
/** 启动批量SQL会话
* @param mapperClass MyBatis Mapper类型
/**
* 启动批量SQL会话
*
* @param mapperClass MyBatis Mapper类型
* @param batchFunction 批量函数(批量添加、批量删除、批量更新等)
* @param arguments 函数附带的所有参数可以使用Map进行包装
* @param <M> Mapper class
* @param <I> Function input
* @param <O> Function output
* @param arguments 函数附带的所有参数可以使用Map进行包装
* @param <M> Mapper class
* @param <I> Function input
* @param <O> Function output
* @return 被包装的批量函数返回值
*/
public <M, I, O> O startBatchSession(Class<M> mapperClass,

View File

@@ -0,0 +1,30 @@
package com.realtime.protection.configuration.utils.enums;
import java.util.HashMap;
import java.util.Map;
public enum ProtocolEnum {
TCP(6),
UDP(17);
private final Integer number;
private static final Map<Integer, ProtocolEnum> map = new HashMap<>();
static {
for (ProtocolEnum protocol : ProtocolEnum.values()) {
map.put(protocol.getProtocolNumber(), protocol);
}
}
ProtocolEnum(int protocolNumber) {
this.number = protocolNumber;
}
public Integer getProtocolNumber() {
return this.number;
}
public static ProtocolEnum getProtocolEnumByNumber(Integer protocolNum) {
return map.get(protocolNum);
}
}

View File

@@ -0,0 +1,51 @@
package com.realtime.protection.configuration.utils.enums;
import com.realtime.protection.configuration.utils.status.State;
import com.realtime.protection.server.task.status.states.*;
import lombok.Getter;
import java.util.HashMap;
import java.util.Map;
@Getter
public enum StateEnum {
// 仅需修改此处即可将任务状态以及对应的State和Num进行对应
PENDING(0, new PendingState()),
RUNNING(1, new RunningState()),
PAUSED(2, new PauseState()),
STOP(3, new StopState()),
FINISHED(4, new FinishedState()),
FAILED(5, new FailedState());
// ----------------------------------------------
private final State state;
private final Integer stateNum;
private static final Map<Integer, State> numToStateMap = new HashMap<>();
private static final Map<State, Integer> stateToNumMap = new HashMap<>();
private static final Map<State, StateEnum> stateToStateEnumMap = new HashMap<>();
static {
for (StateEnum stateEnum : StateEnum.values()) {
numToStateMap.put(stateEnum.getStateNum(), stateEnum.getState());
stateToNumMap.put(stateEnum.getState(), stateEnum.getStateNum());
stateToStateEnumMap.put(stateEnum.getState(), stateEnum);
}
}
StateEnum(int stateNum, State state) {
this.stateNum = stateNum;
this.state = state;
}
public static State getStateByNum(Integer stateNum) {
return numToStateMap.get(stateNum);
}
public static Integer getNumByState(State state) {
return stateToNumMap.get(state);
}
public static StateEnum getStateEnumByState(State state) {
return stateToStateEnumMap.get(state);
}
}

View File

@@ -0,0 +1,33 @@
package com.realtime.protection.configuration.utils.enums;
import lombok.Getter;
import java.util.HashMap;
import java.util.Map;
@Getter
public enum TaskTypeEnum {
STATIC(1),
DYNAMIC(2),
JUDGED(3);
private final int taskType;
private static final Map<Integer, TaskTypeEnum> map = new HashMap<>();
static {
for (TaskTypeEnum taskType : TaskTypeEnum.values()) {
map.put(taskType.getTaskType(), taskType);
}
}
TaskTypeEnum(int taskType) {
this.taskType = taskType;
}
public static TaskTypeEnum getTaskTypeByNum(Integer taskType) {
if (taskType == null) {
return null;
}
return map.get(taskType);
}
}

View File

@@ -0,0 +1,17 @@
package com.realtime.protection.configuration.utils.status;
import lombok.Getter;
@Getter
public enum AuditStatus {
PENDING(0),
UNAUDITED(1),
AUDITED(2);
private final int auditStatus;
AuditStatus(int auditStatus) {
this.auditStatus = auditStatus;
}
}

View File

@@ -14,9 +14,15 @@ public class AuditStatusValidator {
public Boolean checkValidate(Integer auditStatusNow) {
switch (auditStatusNow) {
case 0, 1 -> {return auditStatusOriginal != 2;}
case 2 -> {return auditStatusOriginal != 1;}
default -> {return false;}
case 0, 1 -> {
return auditStatusOriginal != 2;
}
case 2 -> {
return auditStatusOriginal != 1;
}
default -> {
return false;
}
}
}
}

View File

@@ -0,0 +1,9 @@
package com.realtime.protection.configuration.utils.status;
import com.realtime.protection.configuration.exception.DorisStartException;
import com.realtime.protection.server.command.CommandService;
import com.realtime.protection.server.task.TaskService;
public interface State {
Boolean handle(State newState, CommandService commandService, TaskService taskService, Long taskId) throws DorisStartException;
}

View File

@@ -0,0 +1,20 @@
package com.realtime.protection.configuration.utils.status;
import lombok.Getter;
@Getter
public enum StateNum {
PENDING(0),
RUNNING(1),
PAUSED(2),
STOPPED(3),
FAILED(4),
FINISHED(5);
private final int stateNum;
StateNum(int stateNum) {
this.stateNum = stateNum;
}
}