diff --git a/src/main/java/com/dispose/common/Http.java b/src/main/java/com/dispose/common/Http.java index 1f08a75a..1c0e51e1 100644 --- a/src/main/java/com/dispose/common/Http.java +++ b/src/main/java/com/dispose/common/Http.java @@ -2,6 +2,7 @@ package com.dispose.common; import cn.hutool.http.Header; import cn.hutool.http.HttpRequest; + import java.util.Map; /** diff --git a/src/main/java/com/dispose/controller/DisposeNodeInfoController.java b/src/main/java/com/dispose/controller/DisposeNodeInfoController.java index 719d73f2..9267a277 100644 --- a/src/main/java/com/dispose/controller/DisposeNodeInfoController.java +++ b/src/main/java/com/dispose/controller/DisposeNodeInfoController.java @@ -2,6 +2,7 @@ package com.dispose.controller; import com.dispose.common.ConstValue; import com.dispose.common.ErrorCode; +import com.dispose.common.Helper; import com.dispose.pojo.dto.ProtocolReqDTO; import com.dispose.pojo.dto.ProtocolRespDTO; import com.dispose.pojo.entity.DisposeDevice; @@ -36,7 +37,6 @@ import org.springframework.web.bind.annotation.RequestMapping; import org.springframework.web.bind.annotation.ResponseBody; import javax.annotation.Resource; -import java.sql.Timestamp; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -659,15 +659,13 @@ public class DisposeNodeInfoController { taskData.setMessage(ErrorCode.ERR_NOSUCHTASK.getMsg()); } else { runTasks.parallelStream().forEach(k -> { - Long tmLong = Timestamp.valueOf(k.getBeginTime()).toInstant().toEpochMilli() / 1000; - Long endTm = Timestamp.valueOf(k.getPlanEndTime()).toInstant().toEpochMilli() / 1000; - Long diff = endTm - tmLong; + int diff = Helper.getTimestampDiff(k.getBeginTime(), k.getPlanEndTime()); TaskInfoData taskInfoData = TaskInfoData.builder() .taskId(k.getId().toString()) .type(k.getType()) - .startTime(tmLong.intValue()) - .disposeTime(diff.intValue()) + .startTime(Helper.getTimestampSecond(k.getBeginTime())) + .disposeTime(diff) .disposeIp(k.getDisposeIp()) .flowAttack(0) .flowClean(0) @@ -740,15 +738,13 @@ public class DisposeNodeInfoController { taskData.setMessage(ErrorCode.ERR_NOSUCHTASK.getMsg()); } else { runTasks.parallelStream().forEach(k -> { - Long tmLong = Timestamp.valueOf(k.getBeginTime()).toInstant().toEpochMilli() / 1000; - Long endTm = Timestamp.valueOf(k.getPlanEndTime()).toInstant().toEpochMilli() / 1000; - Long diff = endTm - tmLong; + int diff = Helper.getTimestampDiff(k.getBeginTime(), k.getPlanEndTime()); TaskInfoData taskInfoData = TaskInfoData.builder() .taskId(k.getId().toString()) .type(k.getType()) - .startTime(tmLong.intValue()) - .disposeTime(diff.intValue()) + .startTime(Helper.getTimestampSecond(k.getBeginTime())) + .disposeTime(diff) .disposeIp(k.getDisposeIp()) .currentStatus(taskService.taskIsRunning(k) ? 1 : 0) .flowAttack(0) diff --git a/src/main/java/com/dispose/controller/DisposeTaskController.java b/src/main/java/com/dispose/controller/DisposeTaskController.java index 5c2815f1..7fe20d56 100644 --- a/src/main/java/com/dispose/controller/DisposeTaskController.java +++ b/src/main/java/com/dispose/controller/DisposeTaskController.java @@ -1,5 +1,6 @@ package com.dispose.controller; +import com.dispose.common.DeviceCapacity; import com.dispose.common.ErrorCode; import com.dispose.common.FlowDirection; import com.dispose.common.Helper; @@ -16,6 +17,7 @@ import com.dispose.pojo.vo.information.NodeTaskRsp; import com.dispose.pojo.vo.task.GetTaskReq; import com.dispose.pojo.vo.task.StartTaskReq; import com.dispose.pojo.vo.task.StartTaskRsp; +import com.dispose.pojo.vo.task.StopTaskData; import com.dispose.pojo.vo.task.StopTaskReq; import com.dispose.pojo.vo.task.TaskInfoRsp; import com.dispose.service.DisposeNodeManager; @@ -35,7 +37,6 @@ import org.springframework.web.bind.annotation.RequestMapping; import org.springframework.web.bind.annotation.ResponseBody; import javax.annotation.Resource; -import java.sql.Timestamp; import java.util.Arrays; import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; @@ -145,29 +146,34 @@ public class DisposeTaskController { return ProtocolRespDTO.result(ErrorCode.ERR_INPUTMISS); } - for (String v : reqInfo.getTaskId()) { - TaskInfoData taskData = TaskInfoData.builder().taskId(v).build(); - TaskInfoDetail task = taskCacheManager.getTaskById(Long.parseLong(v)); + Arrays.stream(reqInfo.getTaskId()).forEach(v -> { + TaskInfoData taskData = TaskInfoData.builder().type(DeviceCapacity.CLEANUP.getCode()).build(); - if (task == null) { - err = ErrorCode.ERR_TASKNOTRUNNING; + List> ret = taskService.stopTaskByTaskId(Long.parseLong(v)); + + if (ret.size() == 0) { + taskData.setId(v); + taskData.setStatus(ErrorCode.ERR_NOSUCHTASK.getCode()); + taskData.setMessage(ErrorCode.ERR_NOSUCHTASK.getMsg()); + rspInfo.getItems().add(taskData); } else { - err = taskService.stopTask(Long.parseLong(v)); + ret.forEach(k -> { + ErrorCode retError = k.getFirstParam(); + TaskInfoDetail taskInfo = k.getSecondParam(); + taskData.setId(String.valueOf(taskInfo.getDeviceId())); + taskData.setTaskId(String.valueOf(taskInfo.getId())); + taskData.setType(taskInfo.getType()); + taskData.setDisposeIp(taskInfo.getDisposeIp()); + taskData.setLeftTime(Math.abs(Helper.getTimestampDiffNow(taskInfo.getBeginTime()) + - Helper.getTimestampDiffNow(taskInfo.getPlanEndTime()))); - if (err == ErrorCode.ERR_OK) { - taskData.setTaskId(v); - taskData.setDisposeIp(task.getDisposeIp()); - taskData.setType(task.getType()); - taskData.setLeftTime(Math.abs(Helper.getTimestampDiffNow(task.getBeginTime()) - - Helper.getTimestampDiffNow(task.getPlanEndTime()))); - } + taskData.setStatus(retError.getCode()); + taskData.setMessage(retError.getMsg()); + + rspInfo.getItems().add(taskData); + }); } - - taskData.setStatus(err.getCode()); - taskData.setMessage(err.getMsg()); - - rspInfo.getItems().add(taskData); - } + }); return ProtocolRespDTO.result(ErrorCode.ERR_OK, rspInfo); } @@ -196,29 +202,32 @@ public class DisposeTaskController { TaskInfoRsp rspInfo = new TaskInfoRsp(); reqInfo.getItems().forEach(v -> { - ErrorCode retError; - TaskInfoData taskData = TaskInfoData.builder().type(v.getType()).build(); List> ret = taskService.stopTaskByDisposeIp(Long.parseLong(v.getId()), v.getType(), v.getDisposeIp()); if (ret.size() == 0) { - retError = ErrorCode.ERR_NOSUCHTASK; + taskData.setId(v.getId()); + taskData.setStatus(ErrorCode.ERR_NOSUCHTASK.getCode()); + taskData.setMessage(ErrorCode.ERR_NOSUCHTASK.getMsg()); + rspInfo.getItems().add(taskData); } else { - retError = ret.get(0).getFirstParam(); - TaskInfoDetail taskInfo = ret.get(0).getSecondParam(); - taskData.setTaskId(String.valueOf(taskInfo.getId())); - taskData.setType(taskInfo.getType()); - taskData.setDisposeIp(v.getDisposeIp()); - taskData.setLeftTime(Math.abs(Helper.getTimestampDiffNow(taskInfo.getBeginTime()) - - Helper.getTimestampDiffNow(taskInfo.getPlanEndTime()))); + ret.forEach(k -> { + ErrorCode retError = k.getFirstParam(); + TaskInfoDetail taskInfo = k.getSecondParam(); + taskData.setId(String.valueOf(taskInfo.getDeviceId())); + taskData.setTaskId(String.valueOf(taskInfo.getId())); + taskData.setType(taskInfo.getType()); + taskData.setDisposeIp(taskInfo.getDisposeIp()); + taskData.setLeftTime(Math.abs(Helper.getTimestampDiffNow(taskInfo.getBeginTime()) + - Helper.getTimestampDiffNow(taskInfo.getPlanEndTime()))); + + taskData.setStatus(retError.getCode()); + taskData.setMessage(retError.getMsg()); + + rspInfo.getItems().add(taskData); + }); } - - taskData.setId(v.getId()); - taskData.setStatus(retError.getCode()); - taskData.setMessage(retError.getMsg()); - - rspInfo.getItems().add(taskData); }); return ProtocolRespDTO.result(err, rspInfo); @@ -238,7 +247,45 @@ public class DisposeTaskController { public ProtocolRespDTO taskStopAllOfDisposeNode(@RequestBody ProtocolReqDTO mr, @RequestHeader HttpHeaders headers) throws JsonProcessingException { - return taskStopByIp(mr, headers); + ErrorCode err = mr.verifyRequest(headers); + + if (err != ErrorCode.ERR_OK) { + return ProtocolRespDTO.result(err); + } + + StopTaskReq reqInfo = mr.getRequestObject(StopTaskReq.class); + TaskInfoRsp rspInfo = new TaskInfoRsp(); + + reqInfo.getItems().forEach(v -> { + TaskInfoData taskData = TaskInfoData.builder().type(v.getType()).build(); + + List> ret = taskService.stopTaskByDisposeNode(Long.parseLong(v.getId()), v.getType()); + + if (ret.size() == 0) { + taskData.setId("-1"); + taskData.setStatus(ErrorCode.ERR_NOSUCHTASK.getCode()); + taskData.setMessage(ErrorCode.ERR_NOSUCHTASK.getMsg()); + rspInfo.getItems().add(taskData); + } else { + ret.forEach(k -> { + ErrorCode retError = k.getFirstParam(); + TaskInfoDetail taskInfo = k.getSecondParam(); + taskData.setId(String.valueOf(taskInfo.getDeviceId())); + taskData.setTaskId(String.valueOf(taskInfo.getId())); + taskData.setType(taskInfo.getType()); + taskData.setDisposeIp(taskInfo.getDisposeIp()); + taskData.setLeftTime(Math.abs(Helper.getTimestampDiffNow(taskInfo.getBeginTime()) + - Helper.getTimestampDiffNow(taskInfo.getPlanEndTime()))); + + taskData.setStatus(retError.getCode()); + taskData.setMessage(retError.getMsg()); + + rspInfo.getItems().add(taskData); + }); + } + }); + + return ProtocolRespDTO.result(err, rspInfo); } /** @@ -255,7 +302,43 @@ public class DisposeTaskController { public ProtocolRespDTO taskStopAll(@RequestBody ProtocolReqDTO mr, @RequestHeader HttpHeaders headers) throws JsonProcessingException { - return taskStopByIp(mr, headers); + + ErrorCode err = mr.verifyRequest(headers); + + if (err != ErrorCode.ERR_OK) { + return ProtocolRespDTO.result(err); + } + + StopTaskData reqInfo = mr.getRequestObject(StopTaskData.class); + TaskInfoRsp rspInfo = new TaskInfoRsp(); + + TaskInfoData taskData = TaskInfoData.builder().type(reqInfo.getType()).build(); + List> ret = taskService.stopTaskByType(reqInfo.getType()); + + if (ret.size() == 0) { + taskData.setId("-1"); + taskData.setStatus(ErrorCode.ERR_NOSUCHTASK.getCode()); + taskData.setMessage(ErrorCode.ERR_NOSUCHTASK.getMsg()); + rspInfo.getItems().add(taskData); + } else { + ret.forEach(k -> { + ErrorCode retError = k.getFirstParam(); + TaskInfoDetail taskInfo = k.getSecondParam(); + taskData.setId(String.valueOf(taskInfo.getDeviceId())); + taskData.setTaskId(String.valueOf(taskInfo.getId())); + taskData.setType(taskInfo.getType()); + taskData.setDisposeIp(taskInfo.getDisposeIp()); + taskData.setLeftTime(Math.abs(Helper.getTimestampDiffNow(taskInfo.getBeginTime()) + - Helper.getTimestampDiffNow(taskInfo.getPlanEndTime()))); + + taskData.setStatus(retError.getCode()); + taskData.setMessage(retError.getMsg()); + + rspInfo.getItems().add(taskData); + }); + } + + return ProtocolRespDTO.result(err, rspInfo); } /** @@ -294,14 +377,12 @@ public class DisposeTaskController { AtomicBoolean bHasData = new AtomicBoolean(false); taskService.getNodeAllTask(v.getId()).parallelStream().forEach(k -> { if (v.getType().equals(reqInfo.getType()) || reqInfo.getType() == -1) { - Long tmLong = Timestamp.valueOf(k.getBeginTime()).toInstant().toEpochMilli() / 1000; - Long endTm = Timestamp.valueOf(k.getPlanEndTime()).toInstant().toEpochMilli() / 1000; - Long diff = endTm - tmLong; + int diff = Helper.getTimestampDiff(k.getBeginTime(), k.getPlanEndTime()); TaskInfoData taskInfoData = TaskInfoData.builder() .taskId(k.getId().toString()) .type(k.getType()) - .startTime(tmLong.intValue()) - .disposeTime(diff.intValue()) + .startTime(Helper.getTimestampSecond(k.getBeginTime())) + .disposeTime(diff) .disposeIp(k.getDisposeIp()) .flowAttack(0) .flowClean(0) @@ -370,15 +451,13 @@ public class DisposeTaskController { .collect(Collectors.toList()); valuableData.parallelStream().forEach(v -> { - Long tmLong = Timestamp.valueOf(v.getBeginTime()).toInstant().toEpochMilli() / 1000; - Long endTm = Timestamp.valueOf(v.getPlanEndTime()).toInstant().toEpochMilli() / 1000; - Long diff = endTm - tmLong; + int diff = Helper.getTimestampDiff(v.getBeginTime(), v.getPlanEndTime()); TaskInfoData taskInfoData = TaskInfoData.builder() .taskId(v.getId().toString()) .type(v.getType()) - .startTime(tmLong.intValue()) - .disposeTime(diff.intValue()) + .startTime(Helper.getTimestampSecond(v.getBeginTime())) + .disposeTime(diff) .disposeIp(v.getDisposeIp()) .attackType(Helper.attackStringToArray(v.getAttackType())) .flowDirection(v.getFlowDirection()) diff --git a/src/main/java/com/dispose/service/TaskService.java b/src/main/java/com/dispose/service/TaskService.java index 4907f26c..ed6e6418 100644 --- a/src/main/java/com/dispose/service/TaskService.java +++ b/src/main/java/com/dispose/service/TaskService.java @@ -48,6 +48,31 @@ public interface TaskService { */ List> stopTaskByDisposeIp(Long devId, int type, String ipAddr); + /** + * Stop task by task id list. + * + * @param taskId the task id + * @return the list + */ + List> stopTaskByTaskId(Long taskId); + + /** + * Stop task by type list. + * + * @param type the type + * @return the list + */ + List> stopTaskByType(int type); + + /** + * Stop task by dispose node list. + * + * @param devId the dev id + * @param type the type + * @return the list + */ + List> stopTaskByDisposeNode(Long devId, int type); + /** * Finish task error code. * diff --git a/src/main/java/com/dispose/service/impl/TaskServiceImpl.java b/src/main/java/com/dispose/service/impl/TaskServiceImpl.java index d784db6b..2ee927d7 100644 --- a/src/main/java/com/dispose/service/impl/TaskServiceImpl.java +++ b/src/main/java/com/dispose/service/impl/TaskServiceImpl.java @@ -384,6 +384,25 @@ public class TaskServiceImpl implements TaskService { return taskCacheManager.getAllTask(); } + @Override + public List> stopTaskByTaskId(Long taskId) { + List> retList = new ArrayList<>(); + + // 根据处置IP，拿出所有正在处置的任务 + List taskList = taskCacheManager.getAllTask().parallelStream() + .filter(v -> (Objects.equals(v.getId(), taskId)) + && taskIsRunning(v)) + .collect(Collectors.toList()); + + // 对符合条件的任务调用停止命令 + taskList.parallelStream().forEach(v -> { + ErrorCode err = stopTask(v.getId()); + retList.add(new MReturnType<>(err, v)); + }); + + return retList; + } + /** * Task is exists boolean. * @@ -408,8 +427,8 @@ public class TaskServiceImpl implements TaskService { // 根据处置IP，拿出所有正在处置的任务 List taskList = taskCacheManager.getAllTask().parallelStream() - .filter(v -> ((devId == null || devId == -1L) || Objects.equals(v.getDeviceId(), devId) - && Objects.equals(v.getType(), type)) + .filter(v -> ((devId == null || devId == -1L) || Objects.equals(v.getDeviceId(), devId)) + && (Objects.equals(v.getType(), type) || type == -1) && Objects.equals(v.getDisposeIp(), ipAddr) && taskIsRunning(v)) .collect(Collectors.toList()); @@ -422,4 +441,43 @@ public class TaskServiceImpl implements TaskService { return retList; } + + @Override + public List> stopTaskByDisposeNode(Long devId, int type) { + List> retList = new ArrayList<>(); + + // 根据处置IP，拿出所有正在处置的任务 + List taskList = taskCacheManager.getAllTask().parallelStream() + .filter(v -> ((devId == null || devId == -1L) || Objects.equals(v.getDeviceId(), devId)) + && (Objects.equals(v.getType(), type) || type == -1) + && taskIsRunning(v)) + .collect(Collectors.toList()); + + // 对符合条件的任务调用停止命令 + taskList.parallelStream().forEach(v -> { + ErrorCode err = stopTask(v.getId()); + retList.add(new MReturnType<>(err, v)); + }); + + return retList; + } + + @Override + public List> stopTaskByType(int type) { + List> retList = new ArrayList<>(); + + // 根据处置IP，拿出所有正在处置的任务 + List taskList = taskCacheManager.getAllTask().parallelStream() + .filter(v -> (Objects.equals(v.getType(), type) || type == -1) + && taskIsRunning(v)) + .collect(Collectors.toList()); + + // 对符合条件的任务调用停止命令 + taskList.parallelStream().forEach(v -> { + ErrorCode err = stopTask(v.getId()); + retList.add(new MReturnType<>(err, v)); + }); + + return retList; + } } diff --git a/src/test/java/com/dispose/controller/AuthControllerTest.java b/src/test/java/com/dispose/controller/AuthControllerTest.java index fece7ba6..f413a43b 100644 --- a/src/test/java/com/dispose/controller/AuthControllerTest.java +++ b/src/test/java/com/dispose/controller/AuthControllerTest.java @@ -10,8 +10,10 @@ import com.fasterxml.jackson.databind.ObjectMapper; import javax.annotation.Resource; import lombok.extern.slf4j.Slf4j; import org.junit.Assert; +import org.junit.FixMethodOrder; import org.junit.Test; import org.junit.runner.RunWith; +import org.junit.runners.MethodSorters; import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMockMvc; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.http.MediaType; @@ -30,6 +32,7 @@ import static org.springframework.test.web.servlet.result.MockMvcResultMatchers. @RunWith(SpringRunner.class) @SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT) @DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_EACH_TEST_METHOD) +@FixMethodOrder(MethodSorters.NAME_ASCENDING) @Slf4j public class AuthControllerTest extends InitTestEnvironment { /** diff --git a/src/test/java/com/dispose/controller/DeviceNodeManagerControllerTest.java b/src/test/java/com/dispose/controller/DeviceNodeManagerControllerTest.java index d3b2eb26..999c0e87 100644 --- a/src/test/java/com/dispose/controller/DeviceNodeManagerControllerTest.java +++ b/src/test/java/com/dispose/controller/DeviceNodeManagerControllerTest.java @@ -22,8 +22,10 @@ import javax.annotation.Resource; import lombok.extern.slf4j.Slf4j; import org.junit.Assert; +import org.junit.FixMethodOrder; import org.junit.Test; import org.junit.runner.RunWith; +import org.junit.runners.MethodSorters; import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMockMvc; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.http.MediaType; @@ -43,6 +45,7 @@ import static org.springframework.test.web.servlet.result.MockMvcResultMatchers. @RunWith(SpringRunner.class) @SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT) @DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_EACH_TEST_METHOD) +@FixMethodOrder(MethodSorters.NAME_ASCENDING) @Slf4j public class DeviceNodeManagerControllerTest extends InitTestEnvironment { /** diff --git a/src/test/java/com/dispose/controller/TaskControllerTest.java b/src/test/java/com/dispose/controller/TaskControllerTest.java index 8189e2e8..526d983b 100644 --- a/src/test/java/com/dispose/controller/TaskControllerTest.java +++ b/src/test/java/com/dispose/controller/TaskControllerTest.java @@ -17,8 +17,11 @@ import com.fasterxml.jackson.databind.ObjectMapper; import java.util.ArrayList; import java.util.List; import javax.annotation.Resource; + +import org.junit.FixMethodOrder; import org.junit.Test; import org.junit.runner.RunWith; +import org.junit.runners.MethodSorters; import org.springframework.boot.test.autoconfigure.web.servlet.AutoConfigureMockMvc; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.http.MediaType; @@ -37,6 +40,7 @@ import static org.springframework.test.web.servlet.result.MockMvcResultMatchers. @RunWith(SpringRunner.class) @SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT) @DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_EACH_TEST_METHOD) +@FixMethodOrder(MethodSorters.NAME_ASCENDING) public class TaskControllerTest extends InitTestEnvironment { /** * The Mock mvc. @@ -133,7 +137,7 @@ public class TaskControllerTest extends InitTestEnvironment { @Test public void t2_stopTask() throws Exception { IDArrayReq reqData = IDArrayReq.builder() - .taskId(new String[]{"40", "2"}) + .taskId(new String[]{String.valueOf(getExistsTaskId()), "2"}) .build(); ProtocolReqDTO reqInfo = new ProtocolReqDTO(); @@ -263,17 +267,10 @@ public class TaskControllerTest extends InitTestEnvironment { */ @Test public void t5_stopAllTask() throws Exception { - StopTaskData itemData = StopTaskData.builder() - .disposeIp("192.168.1.1") + StopTaskData reqData = StopTaskData.builder() .type(DeviceCapacity.CLEANUP.getCode()) - .id("210") .build(); - StopTaskReq reqData = new StopTaskReq(); - - reqData.setItems(new ArrayList<>()); - reqData.getItems().add(itemData); - ProtocolReqDTO reqInfo = new ProtocolReqDTO(); reqInfo.setVer(ConstValue.Protocol.VERSION); reqInfo.setCryptoType(ConstValue.Protocol.CRYPTO_NONE);