gfwleak/enderbyendera-realtime-protection
Archived
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/master'

Browse Source 
# Conflicts:
#	src/main/java/com/realtime/protection/server/rule/dynamicrule/DynamicRuleControllerApi.java
#	src/main/java/com/realtime/protection/server/rule/staticrule/StaticRuleControllerApi.java
This commit is contained in:
Hao Miao
2024-01-15 20:42:57 +08:00
parent 07d253ef8a 6cfe4bf5d3
commit 3a770192b3
32 changed files with 477 additions and 242 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/main/java/com/realtime/protection/server/command/CommandMapper.java
View File

@@ -1,6 +1,6 @@
package com.realtime.protection.server.command;
import com.realtime.protection.configuration.entity.task.Command;
import com.realtime.protection.configuration.entity.task.TaskCommandInfo;
import org.apache.ibatis.annotations.Mapper;
import org.apache.ibatis.annotations.Param;
@@ -8,13 +8,15 @@ import java.util.List;
@Mapper
public interface CommandMapper {
Boolean createCommand(@Param("command") Command command);
Boolean createCommand(@Param("info") TaskCommandInfo taskCommandInfo);
void createCommands(@Param("commands") List<Command> commands);
void createCommands(@Param("command_infos") List<TaskCommandInfo> taskCommandInfos);
Boolean stopCommandsByTaskId(@Param("task_id") Long taskId);
Boolean removeCommandsByTaskId(@Param("task_id") Long taskId);
Boolean startCommandsByTaskId(@Param("task_id") Long taskId);
List<TaskCommandInfo> queryCommandInfoByTaskId(@Param("task_id") Long taskId);
}

101
src/main/java/com/realtime/protection/server/command/CommandService.java
View File

@@ -2,120 +2,69 @@ package com.realtime.protection.server.command;
import com.alibaba.excel.util.ListUtils;
import com.baomidou.dynamic.datasource.annotation.DS;
import com.realtime.protection.configuration.entity.task.Command;
import com.realtime.protection.configuration.entity.task.TaskCommandInfo;
import com.realtime.protection.configuration.exception.DorisStartException;
import com.realtime.protection.configuration.utils.SqlSessionWrapper;
import com.realtime.protection.configuration.utils.enums.StateEnum;
import com.realtime.protection.server.task.TaskService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import java.time.LocalDateTime;
import java.util.List;
import java.util.function.Function;
@Service
@Slf4j
@DS("doris")
public class CommandService {
private final CommandMapper commandMapper;
private final TaskService taskService;
private final SqlSessionWrapper sqlSessionWrapper;
private static final int BatchSize = 100;
private final Function<CommandMapper, Function<TaskCommandInfo, Void>> createCommandBatchFunction;
public CommandService(CommandMapper commandMapper, TaskService taskService, SqlSessionWrapper sqlSessionWrapper) {
public CommandService(CommandMapper commandMapper, SqlSessionWrapper sqlSessionWrapper) {
this.commandMapper = commandMapper;
this.taskService = taskService;
this.sqlSessionWrapper = sqlSessionWrapper;
this.createCommandBatchFunction = mapper -> info -> {
if (info.getFrequency() == null) {
Command command = Command.generateCommand(info, info.getStartTime());
mapper.createCommand(command);
}
}
List<Command> commandBatch = ListUtils.newArrayListWithExpectedSize(BatchSize);
LocalDateTime validTime = info.getStartTime();
public Boolean createCommand(TaskCommandInfo commandInfo) {
return commandMapper.createCommand(commandInfo);
}
while (validTime.isBefore(info.getEndTime())) {
Command command = Command.generateCommand(info, validTime);
commandBatch.add(command);
validTime = validTime.plusMinutes(info.getFrequency());
if (commandBatch.size() < BatchSize) {
public void createCommands(List<TaskCommandInfo> taskCommandInfos) {
Function<CommandMapper, Function<List<TaskCommandInfo>, Boolean>> function = mapper -> list -> {
List<TaskCommandInfo> taskCommandInfoBatch = ListUtils.newArrayListWithExpectedSize(BatchSize);
for (TaskCommandInfo info : list) {
taskCommandInfoBatch.add(info);
if (taskCommandInfoBatch.size() < BatchSize) {
continue;
}
mapper.createCommands(commandBatch);
commandBatch.clear();
commandMapper.createCommands(taskCommandInfoBatch);
taskCommandInfoBatch.clear();
}
if (!commandBatch.isEmpty()) {
mapper.createCommands(commandBatch);
commandBatch.clear();
if (!taskCommandInfoBatch.isEmpty()) {
commandMapper.createCommands(taskCommandInfoBatch);
taskCommandInfoBatch.clear();
}
log.debug(String.format("在task(%d)和rule(%d)中构建了全部指令",
info.getTaskId(), info.getRuleId()));
return null;
};
}
@Async
@DS("doris")
public void createCommand(TaskCommandInfo commandInfo) throws DorisStartException {
try {
sqlSessionWrapper.startBatchSession(CommandMapper.class, createCommandBatchFunction, commandInfo);
taskService.changeTaskStatus(commandInfo.getTaskId(), StateEnum.RUNNING.getStateNum());
} catch (Exception e) {
throw new DorisStartException(e, commandInfo.getTaskId());
}
}
@Async
@DS("doris")
public void createCommands(List<TaskCommandInfo> taskCommandInfos) throws DorisStartException {
Function<CommandMapper, Function<List<TaskCommandInfo>, Void>> function = mapper -> list -> {
if (list == null || list.isEmpty()) {
return null;
}
for (TaskCommandInfo info : list) {
createCommandBatchFunction.apply(mapper).apply(info);
}
taskService.changeTaskStatus(list.get(0).getTaskId(), StateEnum.RUNNING.getStateNum());
return null;
return true;
};
try {
sqlSessionWrapper.startBatchSession(CommandMapper.class, function, taskCommandInfos);
} catch (Exception e) {
TaskCommandInfo info = null;
if (taskCommandInfos != null) {
info = taskCommandInfos.get(0);
}
Long taskId = null;
if (info != null) {
taskId = info.getTaskId();
}
throw new DorisStartException(e, taskId);
}
sqlSessionWrapper.startBatchSession(CommandMapper.class, function, taskCommandInfos);
}
public List<TaskCommandInfo> queryCommandInfoByTaskId(Long taskId) {
return commandMapper.queryCommandInfoByTaskId(taskId);
}
@DS("doris")
public Boolean startCommandsByTaskId(Long taskId) {
return commandMapper.startCommandsByTaskId(taskId);
}
@DS("doris")
public Boolean stopCommandsByTaskId(Long taskId) {
return commandMapper.stopCommandsByTaskId(taskId);
}
@DS("doris")
public Boolean removeCommandsByTaskId(Long taskId) {
return commandMapper.removeCommandsByTaskId(taskId);
}

1
src/main/java/com/realtime/protection/server/defense/object/ProtectObjectService.java
View File

@@ -78,7 +78,6 @@ public class ProtectObjectService {
return false;
}
boolean success = true;
Integer result;
List<Integer> protectObjectBatch = ListUtils.newArrayListWithExpectedSize(batchSize);
for (Integer protectObjectId : list) {

37
src/main/java/com/realtime/protection/server/task/TaskController.java
View File

@@ -1,9 +1,11 @@
package com.realtime.protection.server.task;
import com.realtime.protection.configuration.entity.task.Task;
import com.realtime.protection.configuration.entity.task.TaskCommandInfo;
import com.realtime.protection.configuration.exception.DorisStartException;
import com.realtime.protection.configuration.response.ResponseResult;
import com.realtime.protection.configuration.utils.EntityUtils;
import com.realtime.protection.server.command.CommandService;
import com.realtime.protection.server.task.status.StateChangeService;
import jakarta.validation.Valid;
import jakarta.validation.constraints.Max;
@@ -18,10 +20,12 @@ import java.util.List;
public class TaskController implements TaskControllerApi {
private final TaskService taskService;
private final CommandService commandService;
private final StateChangeService stateChangeService;
public TaskController(TaskService taskService, StateChangeService stateChangeService) {
public TaskController(TaskService taskService, CommandService commandService, StateChangeService stateChangeService) {
this.taskService = taskService;
this.commandService = commandService;
this.stateChangeService = stateChangeService;
}
@@ -43,6 +47,24 @@ public class TaskController implements TaskControllerApi {
.setData("success", false);
}
// API推送Endpoint
@Override
@PostMapping("/api/new")
public ResponseResult newTaskWithAPI(@RequestBody @Valid TaskCommandInfo taskCommandInfo) {
Long taskId = taskService.newTaskUsingCommandInfo(taskCommandInfo);
if (taskId <= 0) {
return ResponseResult.invalid()
.setData("taskId", -1)
.setData("success", false);
}
commandService.createCommand(taskCommandInfo);
return ResponseResult.ok()
.setData("taskId", taskId)
.setData("success", true);
}
@Override
@GetMapping("/query")
public ResponseResult queryTasks(@RequestParam(value = "task_status", required = false) Integer taskStatus,
@@ -62,7 +84,7 @@ public class TaskController implements TaskControllerApi {
Task task = taskService.queryTask(id);
if (task == null) {
return ResponseResult.invalid().setMessage("Task ID is invalid");
return ResponseResult.invalid().setMessage("无效Task ID，也许该ID对应的任务不存在？");
}
return ResponseResult.ok()
@@ -103,7 +125,16 @@ public class TaskController implements TaskControllerApi {
@PathVariable @NotNull Long taskId) throws DorisStartException {
return ResponseResult.ok()
.setData("task_id", taskId)
.setData("success", stateChangeService.changeState(stateNum, taskId))
// 外部修改状态，需要进行状态检查
.setData("success", stateChangeService.changeState(stateNum, taskId, false))
.setData("status_now", taskService.queryTaskStatus(taskId));
}
@Override
@GetMapping("/{taskId}/commands")
public ResponseResult queryCommandInfoByTaskId(@PathVariable Long taskId) {
return ResponseResult.ok()
.setData("success", true)
.setData("commands", commandService.queryCommandInfoByTaskId(taskId));
}
}

38
src/main/java/com/realtime/protection/server/task/TaskControllerApi.java
View File

@@ -1,6 +1,7 @@
package com.realtime.protection.server.task;
import com.realtime.protection.configuration.entity.task.Task;
import com.realtime.protection.configuration.entity.task.TaskCommandInfo;
import com.realtime.protection.configuration.exception.DorisStartException;
import com.realtime.protection.configuration.response.ResponseResult;
import io.swagger.v3.oas.annotations.Operation;
@@ -34,6 +35,24 @@ public interface TaskControllerApi {
)
ResponseResult newTask(@RequestBody @Valid Task task);
// API推送Endpoint
@PostMapping("/api/new")
@Operation(
summary = "任务推送外部API",
description = "提供给外部的任务推送API",
responses = {
@ApiResponse(
description = "返回外部任务推送结果",
content = @Content(
mediaType = "application/json",
schema = @Schema(implementation = ResponseResult.class)
)
)
},
requestBody = @io.swagger.v3.oas.annotations.parameters.RequestBody(description = "任务推送信息")
)
ResponseResult newTaskWithAPI(@RequestBody @Valid TaskCommandInfo taskCommandInfo) throws DorisStartException;
@GetMapping("/query")
@Operation(
summary = "查询任务",
@@ -162,4 +181,23 @@ public interface TaskControllerApi {
)
ResponseResult changeTaskStatus(@PathVariable @NotNull @Min(0) @Max(6) Integer stateNum,
@PathVariable @NotNull Long taskId) throws DorisStartException;
@GetMapping("/{taskId}/commands")
@Operation(
summary = "获得任务已推送指令的相关数据",
description = "获得任务已推送指令的相关数据，包括最新下发时间、首次下发时间、下发次数、下发成功次数等",
responses = {
@ApiResponse(
description = "返回任务已推送指令的相关数据",
content = @Content(
mediaType = "application/json",
schema = @Schema(implementation = ResponseResult.class)
)
)
},
parameters = {
@Parameter(name = "taskId", description = "任务ID")
}
)
ResponseResult queryCommandInfoByTaskId(@PathVariable Long taskId);
}

12
src/main/java/com/realtime/protection/server/task/TaskMapper.java
View File

@@ -12,10 +12,12 @@ public interface TaskMapper {
void newTask(@Param("task") Task task);
void newTaskStaticRuleConcat(@Param("task_id") Long taskId,
@Param("rule_ids") List<Long> staticRuleIds);
@Param("rule_ids") List<Integer> staticRuleIds);
void newTaskDynamicRuleConcat(@Param("task_id") Long taskId,
@Param("rule_ids") List<Long> dynamicRuleIds);
@Param("rule_ids") List<Integer> dynamicRuleIds);
void newTaskUsingCommandInfo(@Param("info") TaskCommandInfo taskCommandInfo);
List<Task> queryTasks(@Param("task_status") Integer taskStatus, @Param("task_type") String task_type,
@Param("task_name") String taskName, @Param("task_creator") String taskCreator,
@@ -35,9 +37,13 @@ public interface TaskMapper {
Boolean changeTaskStatus(@Param("task_id") Long taskId, @Param("state") Integer stateNum);
List<TaskCommandInfo> getStaticCommands(@Param("task_id") Long taskId);
List<TaskCommandInfo> getStaticCommandInfos(@Param("task_id") Long taskId);
Integer queryTaskAuditStatus(@Param("task_id") Long taskId);
Integer queryTaskStatus(@Param("task_id") Long taskId);
List<Integer> queryDynamicRuleIdsFromTaskId(@Param("task_id") Long taskId);
List<Integer> queryStaticRuleIdsFromTaskId(@Param("task_id") Long taskId);
}

37
src/main/java/com/realtime/protection/server/task/TaskService.java
View File

@@ -10,6 +10,7 @@ import org.springframework.transaction.annotation.Transactional;
import java.util.List;
@Service
@DS("mysql")
public class TaskService {
private final TaskMapper taskMapper;
@@ -19,6 +20,10 @@ public class TaskService {
@Transactional
public Long newTask(Task task) {
task.setTaskCreateUserId(1);
task.setTaskCreateUsername("xxx");
task.setTaskCreateDepart("xxx");
taskMapper.newTask(task);
if (task.getStaticRuleIds() != null && !task.getStaticRuleIds().isEmpty())
@@ -30,14 +35,33 @@ public class TaskService {
return task.getTaskId();
}
@Transactional
public List<Task> queryTasks(Integer taskStatus,
String taskType, String taskName, String taskCreator,
Integer page, Integer pageSize) {
return taskMapper.queryTasks(taskStatus, taskType, taskName, taskCreator, page, pageSize);
List<Task> tasks = taskMapper.queryTasks(taskStatus, taskType, taskName, taskCreator, page, pageSize);
for (Task task : tasks) {
if (task == null) {
continue;
}
task.setStaticRuleIds(taskMapper.queryStaticRuleIdsFromTaskId(task.getTaskId()));
task.setDynamicRuleIds(taskMapper.queryDynamicRuleIdsFromTaskId(task.getTaskId()));
}
return tasks;
}
@Transactional
public Task queryTask(Long id) {
return taskMapper.queryTask(id);
Task task = taskMapper.queryTask(id);
if (task == null) {
return null;
}
task.setStaticRuleIds(taskMapper.queryStaticRuleIdsFromTaskId(task.getTaskId()));
task.setDynamicRuleIds(taskMapper.queryDynamicRuleIdsFromTaskId(task.getTaskId()));
return task;
}
@Transactional
@@ -74,13 +98,13 @@ public class TaskService {
return taskMapper.deleteTask(taskId);
}
@DS("mysql")
public Boolean changeTaskStatus(Long taskId, Integer stateNum) {
return taskMapper.changeTaskStatus(taskId, stateNum);
}
public List<TaskCommandInfo> getStaticCommandInfos(Long taskId) {
return taskMapper.getStaticCommands(taskId);
return taskMapper.getStaticCommandInfos(taskId);
}
public Integer queryTaskAuditStatus(Long taskId) {
@@ -90,4 +114,9 @@ public class TaskService {
public Integer queryTaskStatus(Long taskId) {
return taskMapper.queryTaskStatus(taskId);
}
public Long newTaskUsingCommandInfo(TaskCommandInfo taskCommandInfo) {
taskMapper.newTaskUsingCommandInfo(taskCommandInfo);
return taskCommandInfo.getTaskId();
}
}

35
src/main/java/com/realtime/protection/server/task/status/StateChangeService.java
View File

@@ -23,12 +23,7 @@ public class StateChangeService {
}
@DSTransactional
public Boolean changeState(Integer stateNum, Long taskId) throws DorisStartException {
if (Objects.equals(stateNum, StateEnum.GENERATING.getStateNum()) ||
Objects.equals(stateNum, StateEnum.FAILED.getStateNum())) {
throw new IllegalArgumentException("非法任务状态：" + StateEnum.getStateByNum(stateNum));
}
public Boolean changeState(Integer stateNum, Long taskId, Boolean inner) throws DorisStartException {
Integer originalStateNum = taskService.queryTaskStatus(taskId);
if (originalStateNum == null) {
throw new IllegalArgumentException("无法找到" + taskId + "的任务状态，也许任务ID不存在?");
@@ -38,8 +33,15 @@ public class StateChangeService {
State newState = StateEnum.getStateByNum(stateNum);
if (newState == null) {
return false;
if (!inner && !checkState(originalState, newState)) {
throw new IllegalArgumentException(
String.format("任务状态转换失败，原状态：%s，欲切换状态：%s",
originalState.getClass().getSimpleName(),
newState.getClass().getSimpleName()));
}
if (Objects.equals(originalState, newState)) {
return true;
}
if (!originalState.handle(newState, commandService, taskService, taskId)) {
@@ -54,4 +56,21 @@ public class StateChangeService {
// 这里一定是handle成功的状态，我们再进行task status的修改，如果handle失败，要么返回false，要么抛出异常，不会进入此处
return taskService.changeTaskStatus(taskId, stateNum);
}
private Boolean checkState(State originalState, State newState) {
if (originalState == null || newState == null) {
return false;
}
// FAILED、FINISHED状态以及GENERATING都只能在程序内部修改，外部接口不能修改
if (Objects.equals(newState, StateEnum.FAILED.getState())
|| Objects.equals(newState, StateEnum.FINISHED.getState())
|| Objects.equals(newState, StateEnum.GENERATING.getState())) {
return false;
}
// 在任务状态转换为GENERATING之后，我们需要在外部接口屏蔽掉所有状态
// 我们需要保证只有任务创建函数才能将GENERATING状态转换为RUNNING状态
return !Objects.equals(originalState, StateEnum.GENERATING.getState());
}
}

1
src/main/java/com/realtime/protection/server/task/status/states/FailedState.java
View File

@@ -13,7 +13,6 @@ public class FailedState extends StateHandler implements State {
return switch (StateEnum.getStateEnumByState(newState)) {
case RUNNING -> handleStart(taskService, commandService, taskId);
case STOP -> handleStop(commandService, taskId);
case FAILED -> true;
default -> throw new IllegalStateException("Unexpected value: " + StateEnum.getStateEnumByState(newState));
};
}

5
src/main/java/com/realtime/protection/server/task/status/states/GeneratingState.java
View File

@@ -1,6 +1,5 @@
package com.realtime.protection.server.task.status.states;
import com.realtime.protection.configuration.exception.DorisStartException;
import com.realtime.protection.configuration.utils.enums.StateEnum;
import com.realtime.protection.configuration.utils.status.State;
import com.realtime.protection.server.command.CommandService;
@@ -9,9 +8,9 @@ import com.realtime.protection.server.task.status.StateHandler;
public class GeneratingState extends StateHandler implements State {
@Override
public Boolean handle(State newState, CommandService commandService, TaskService taskService, Long taskId) throws DorisStartException {
public Boolean handle(State newState, CommandService commandService, TaskService taskService, Long taskId) {
return switch (StateEnum.getStateEnumByState(newState)) {
case RUNNING, GENERATING -> true;
case RUNNING -> true;
case FAILED -> handleFailed(commandService, taskId);
default -> throw new IllegalStateException("Unexpected value: " + StateEnum.getStateEnumByState(newState));
};

2
src/main/java/com/realtime/protection/server/task/status/states/PendingState.java
View File

@@ -11,8 +11,8 @@ public class PendingState extends StateHandler implements State {
@Override
public Boolean handle(State newState, CommandService commandService, TaskService taskService, Long taskId) throws DorisStartException {
return switch (StateEnum.getStateEnumByState(newState)) {
case GENERATING -> handleStart(taskService, commandService, taskId);
case FAILED -> handleFailed(commandService, taskId);
case RUNNING -> handleStart(taskService, commandService, taskId);
default -> throw new IllegalStateException("Unexpected value: " + StateEnum.getStateEnumByState(newState));
};
}

1
src/main/java/com/realtime/protection/server/task/status/states/RunningState.java
View File

@@ -10,7 +10,6 @@ public class RunningState extends StateHandler implements State {
@Override
public Boolean handle(State newState, CommandService commandService, TaskService taskService, Long taskId) {
return switch (StateEnum.getStateEnumByState(newState)) {
case RUNNING, GENERATING -> true;
case PAUSED -> handlePause(commandService, taskId);
case STOP -> handleStop(commandService, taskId);
case FINISHED -> handleFinish(commandService, taskId);

6
src/main/java/com/realtime/protection/server/whitelist/WhiteListService.java
View File

@@ -2,7 +2,7 @@ package com.realtime.protection.server.whitelist;
import com.alibaba.excel.util.ListUtils;
import com.realtime.protection.configuration.entity.rule.staticrule.StaticRuleObject;
import com.realtime.protection.configuration.entity.task.Command;
import com.realtime.protection.configuration.entity.task.TaskCommandInfo;
import com.realtime.protection.configuration.entity.whitelist.WhiteListObject;
import com.realtime.protection.configuration.utils.SqlSessionWrapper;
import com.realtime.protection.configuration.utils.status.AuditStatusValidator;
@@ -129,10 +129,10 @@ public class WhiteListService {
}
public List<WhiteListObject> whiteListCommandJudge(Command command) {
public List<WhiteListObject> whiteListCommandJudge(TaskCommandInfo taskCommandInfo) {
//参数应该是指令，不管动态静态
// 命中的whitelist列表：每一列包含ip port url
return whiteListMapper.whiteListCommandJudge(command.getFiveTupleWithMask());
return whiteListMapper.whiteListCommandJudge(taskCommandInfo.getFiveTupleWithMask());
}