diff --git a/src/main/java/com/dispose/ability/impl/PengXinAbilityImpl.java b/src/main/java/com/dispose/ability/impl/PengXinAbilityImpl.java index c3c62725..c4d90511 100644 --- a/src/main/java/com/dispose/ability/impl/PengXinAbilityImpl.java +++ b/src/main/java/com/dispose/ability/impl/PengXinAbilityImpl.java @@ -7,12 +7,9 @@ import com.dispose.common.ErrorCode; import com.dispose.common.Helper; import com.dispose.common.IpAddrType; import com.dispose.common.NetflowDirection; -import com.dispose.common.DisposeDeviceType; import com.dispose.manager.DisposeTaskManager; import com.dispose.mapper.DisposeDeviceMapper; import com.dispose.pojo.dto.protocol.base.ProtocolRespDTO; -import com.dispose.pojo.entity.DeviceTask; -import com.dispose.pojo.entity.DisposeTask; import com.dispose.pojo.po.MulReturnType; import com.dispose.pojo.vo.DeviceFirewareInfo; import com.pengxin.dispose.common.PengXinDisposeAbilityRsp; @@ -33,7 +30,6 @@ import javax.annotation.Resource; import javax.servlet.http.HttpServletResponse; import java.util.ArrayList; import java.util.List; -import java.util.stream.Collectors; /** * The type Peng xin ability. @@ -42,18 +38,6 @@ import java.util.stream.Collectors; */ @Slf4j public class PengXinAbilityImpl implements DisposeAbility { - /** - * The Dispose task manager. - */ - @Resource - private DisposeTaskManager disposeTaskManager; - - /** - * The Dispose device mapper. - */ - @Resource - private DisposeDeviceMapper disposeDeviceMapper; - /** * The Protect ip v 4. */ @@ -62,48 +46,51 @@ public class PengXinAbilityImpl implements DisposeAbility { * The Protect ip v 6. */ private final List protectIpV6 = new ArrayList<>(); - /** * The Restful interface. */ private final PengXinInterface restfulInterface = new PengXinInterface(); - + /** + * The Peng xin dispose type. + */ + private final Integer[] pengXinDisposeType = new Integer[]{0, 4, 1, 3}; + /** + * The Dispose task manager. + */ + @Resource + private DisposeTaskManager disposeTaskManager; + /** + * The Dispose device mapper. + */ + @Resource + private DisposeDeviceMapper disposeDeviceMapper; /** * The Task req id. */ private Long taskReqId = System.currentTimeMillis(); - /** * The Url root path. */ @Getter @Setter private String urlRootPath; - /** * The Username. */ @Getter @Setter private String username; - /** * The Password. */ @Getter @Setter private String password; - /** * The Token. */ private String token; - /** - * The Peng xin dispose type. - */ - private final Integer[] pengXinDisposeType = new Integer[]{0, 4, 1, 3}; - /** * Gets peng xin task type. * @@ -160,7 +147,8 @@ public class PengXinAbilityImpl implements DisposeAbility { * @return the mul return type */ @Override - public MulReturnType runDispose(String disposeObject, DisposeObjectType objectType, DisposeCapacityType capType, + public MulReturnType runDispose(String disposeObject, DisposeObjectType objectType, + DisposeCapacityType capType, @Nullable NetflowDirection nfDirection, @Nullable Integer attackType, @Nullable Long duration) { @@ -186,12 +174,12 @@ public class PengXinAbilityImpl implements DisposeAbility { } reqItems.add(PengXinStartTaskItem.builder() - .type(getPengXinTaskType(capType)) - .disposeObject(disposeObject) - .objectType(getPengXinObjectType(objectType)) - .disposeTime(duration) - .taskReqId(String.valueOf(taskReqId++)) - .build()); + .type(getPengXinTaskType(capType)) + .disposeObject(disposeObject) + .objectType(getPengXinObjectType(objectType)) + .disposeTime(duration) + .taskReqId(String.valueOf(taskReqId++)) + .build()); ProtocolRespDTO rspInfo = restfulInterface.startDisposeTask(url, token, reqItems); @@ -249,7 +237,7 @@ public class PengXinAbilityImpl implements DisposeAbility { } ProtocolRespDTO rspInfo = restfulInterface.stopDisposeTask(url, token, - new String[]{taskId}); + new String[]{taskId}); // 判断是否token过期 if (rspInfo != null && rspInfo.getCode() == HttpServletResponse.SC_UNAUTHORIZED) { @@ -279,18 +267,18 @@ public class PengXinAbilityImpl implements DisposeAbility { @Override public MulReturnType getAbilityDeviceFireware() { return new MulReturnType<>(ErrorCode.ERR_OK, - DeviceFirewareInfo.builder() - .vendor("PengXin") - .model("Unknown") - .firmware("Unknown") - .os("Linux Server") - .kernel("Linux") - .arch("x86_64") - .version("Unknown") - .memory(-1) - .freeMemory(-1) - .cpuUsed(-1) - .build()); + DeviceFirewareInfo.builder() + .vendor("PengXin") + .model("Unknown") + .firmware("Unknown") + .os("Linux Server") + .kernel("Linux") + .arch("x86_64") + .version("Unknown") + .memory(-1) + .freeMemory(-1) + .cpuUsed(-1) + .build()); } /** @@ -441,25 +429,14 @@ public class PengXinAbilityImpl implements DisposeAbility { try { String url = urlRootPath + "dispose_device/task/get"; - DisposeTask task = disposeTaskManager.getDisposeTaskById(taskId); - - if (task == null) { - return new MulReturnType<>(ErrorCode.ERR_NOSUCHTASK, -1L); - } - - //根据设备ID获取设备信息，判断任务列表中的设备为鹏信，获取externId - List dev = task.getDeviceTask().stream(). - filter(v -> disposeDeviceMapper.getDeviceById(v.getDeviceId()).getDeviceType() == DisposeDeviceType.PENGXIN_PLATFORM).collect(Collectors.toList()); - ProtocolRespDTO rspInfo = restfulInterface.getDeviceTaskStatus(url, token, - new String[]{String.valueOf(dev.get(0).getExternId())}); + new String[]{taskId}); // 判断是否token过期 if (rspInfo != null && rspInfo.getCode() == HttpServletResponse.SC_UNAUTHORIZED) { // 重新登录获取 token upgradeToken(); - rspInfo = restfulInterface.getDeviceTaskStatus(url, token, - new String[]{String.valueOf(dev.get(0).getExternId())}); + rspInfo = restfulInterface.getDeviceTaskStatus(url, token, new String[]{taskId}); } if (rspInfo != null && rspInfo.getCode() == ErrorCode.ERR_OK.getHttpCode()) { @@ -469,7 +446,7 @@ public class PengXinAbilityImpl implements DisposeAbility { } return new MulReturnType<>(ErrorCode.ERR_OK, - rspInfo.getMsgContent().getItems().get(0).getTaskStatus()); + rspInfo.getMsgContent().getItems().get(0).getTaskStatus()); } else { return new MulReturnType<>(ErrorCode.ERR_CALLDEVICE, -1L); } diff --git a/src/main/java/com/dispose/mapper/DeviceTaskMapper.java b/src/main/java/com/dispose/mapper/DeviceTaskMapper.java index a357836a..dc7e5c6a 100644 --- a/src/main/java/com/dispose/mapper/DeviceTaskMapper.java +++ b/src/main/java/com/dispose/mapper/DeviceTaskMapper.java @@ -88,12 +88,12 @@ public interface DeviceTaskMapper { /** * Change task dev status int. * - * @param id the id - * @param status the status + * @param externId the extern id + * @param devStatus the dev status * @return the int */ - int changeTaskDevStatus(@Param("id") String id, - @Param("status") Long status); + int changeTaskDevStatus(@Param("externId") String externId, + @Param("devStatus") Long devStatus); /** * Change exec attack type in value int. diff --git a/src/main/java/com/dispose/service/impl/DeviceTaskManagerServiceImpl.java b/src/main/java/com/dispose/service/impl/DeviceTaskManagerServiceImpl.java index a4179df5..ae1ca66e 100644 --- a/src/main/java/com/dispose/service/impl/DeviceTaskManagerServiceImpl.java +++ b/src/main/java/com/dispose/service/impl/DeviceTaskManagerServiceImpl.java @@ -715,7 +715,7 @@ public class DeviceTaskManagerServiceImpl implements DeviceTaskManagerService { // 任务预计结束5分钟内， if (Helper.getTimestampDiffNow(task.getPlanEndTime()) < 5 * 60) { - if (v.getEndTime() == null && Helper.getTimestampDiffNow(v.getEndTime()) < 5 * 60) { + if (v.getEndTime() == null || Helper.getTimestampDiffNow(v.getEndTime()) < 5 * 60) { if (v.getExternId() != null && v.getExternId().length() > 0) { MulReturnType ret = ai.getDb().taskStatus(v.getExternId()); diff --git a/src/main/resources/mappers/DeviceTask.xml b/src/main/resources/mappers/DeviceTask.xml index 63325b59..7071efb5 100644 --- a/src/main/resources/mappers/DeviceTask.xml +++ b/src/main/resources/mappers/DeviceTask.xml @@ -91,7 +91,7 @@ UPDATE device_task SET devStatus = #{devStatus} - WHERE id = #{id} + WHERE externId = #{externId} AND devStatus != #{devStatus}