From ed9c273be63fdd8f1f8e24a4c91c1a91db880acc Mon Sep 17 00:00:00 2001 From: HuangXin Date: Sat, 10 Oct 2020 18:07:14 +0800 Subject: [PATCH] =?UTF-8?q?OCT=20REM:=201.=20=E5=A2=9E=E5=8A=A0=E4=BA=BF?= =?UTF-8?q?=E9=98=B3=E8=AE=BE=E5=A4=87=E5=90=AF=E5=8A=A8=E3=80=81=E5=81=9C?= =?UTF-8?q?=E6=AD=A2=E5=A4=84=E7=BD=AE=E4=BB=BB=E5=8A=A1=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ability/impl/YiYangAbilityImpl.java | 136 +++++++++++++++++- .../impl/DeviceTaskManagerServiceImpl.java | 86 +++++++++++ .../java/com/dispose/setup/SystemInitial.java | 2 +- 3 files changed, 219 insertions(+), 5 deletions(-) diff --git a/src/main/java/com/dispose/ability/impl/YiYangAbilityImpl.java b/src/main/java/com/dispose/ability/impl/YiYangAbilityImpl.java index 3998f01e..f7b91bf5 100644 --- a/src/main/java/com/dispose/ability/impl/YiYangAbilityImpl.java +++ b/src/main/java/com/dispose/ability/impl/YiYangAbilityImpl.java @@ -12,6 +12,9 @@ import com.dispose.pojo.po.MulReturnType; import com.dispose.pojo.vo.DeviceFirewareInfo; import com.yiyang.dispose.common.YiYangDisposeAbilityRsp; import com.yiyang.dispose.common.YiYangLoginRsp; +import com.yiyang.dispose.common.YiYangStartTaskItem; +import com.yiyang.dispose.common.YiYangStartTaskRsp; +import com.yiyang.dispose.common.YiYangStopTaskRsp; import com.yiyang.dispose.protocol.YiYangInterface; import inet.ipaddr.IPAddress; import inet.ipaddr.IPAddressString; @@ -46,6 +49,11 @@ public class YiYangAbilityImpl implements DisposeAbility { */ private final YiYangInterface restfulInterface = new YiYangInterface(); + /** + * The Task req id. + */ + private Long taskReqId = System.currentTimeMillis(); + /** * The Url root path. */ @@ -72,6 +80,38 @@ public class YiYangAbilityImpl implements DisposeAbility { */ private String token; + /** + * The Yi yang dispose type. + */ + private final Integer[] yiYangDisposeType = new Integer[] {0, 4, 1, 3}; + + /** + * Gets yi yang task type. + * + * @param capType the cap type + * @return the yi yang task type + */ + private int getYiYangTaskType(DisposeCapacityType capType) { + return yiYangDisposeType[capType.getValue()]; + } + + /** + * Gets yi yang object type. + * + * @param objType the obj type + * @return the yi yang object type + */ + private int getYiYangObjectType(DisposeObjectType objType) { + switch (objType) { + case URL: + return 2; + case DOMAIN: + return 1; + default: + return 0; + } + } + /** * Init device env. * @@ -102,8 +142,57 @@ public class YiYangAbilityImpl implements DisposeAbility { @Override public MulReturnType runDispose(String ip, DisposeCapacityType capType, @Nullable NetflowDirection nfDirection, - @Nullable Integer attackType, @Nullable Long duration) { - return null; + @Nullable Integer attackType, + @Nullable Long duration) { + List reqItems = new ArrayList<>(); + + String url = urlRootPath + "dispose_device/task/start"; + + if (token == null || token.length() == 0) { + return new MulReturnType<>(ErrorCode.ERR_LOGOUT, ""); + } + + try { + log.info("++++Begging YiYang Start Cleanup Task: {}", ip); + + if (capType != DisposeCapacityType.BLACKHOOL) { + log.error("----Error YiYang don't support dispose capacity type: {}", capType); + return new MulReturnType<>(ErrorCode.ERR_UNSUPPORT, null); + } + + // 适配处置时间参数, -1为不限制处置时间 + if (duration == null || duration < 0) { + duration = -1L; + } + + reqItems.add(YiYangStartTaskItem.builder() + .type(getYiYangTaskType(capType)) + .disposeObject(ip) + .objectType(getYiYangObjectType(DisposeObjectType.IP)) + .disposeTime(duration) + .taskReqId(String.valueOf(taskReqId++)) + .build()); + + ProtocolRespDTO rspInfo = restfulInterface.startDisposeTask(url, token, reqItems); + + // 判断是否token过期 + if (rspInfo != null && rspInfo.getCode() == HttpServletResponse.SC_UNAUTHORIZED) { + // 重新登录获取 token + upgradeToken(); + rspInfo = restfulInterface.startDisposeTask(url, token, reqItems); + } + + if (rspInfo == null || rspInfo.getCode() != ErrorCode.ERR_OK.getHttpCode()) { + log.error("----Error YiYang start clean {} server return error", ip); + return new MulReturnType<>(ErrorCode.ERR_HAOHAN_ERROR, null); + } + + log.debug("----Finish YiYang Start Cleanup Task: {}", ip); + return new MulReturnType<>(ErrorCode.ERR_OK, rspInfo.getMsgContent().getItems().get(0).getTaskId()); + } catch (Exception ex) { + log.error("----Exception YiYang Start Cleanup Task: {}, {}, {}", ip, nfDirection, duration); + return new MulReturnType<>(ErrorCode.ERR_SYSTEMEXCEPTION, null); + } } /** @@ -119,8 +208,47 @@ public class YiYangAbilityImpl implements DisposeAbility { @Override public MulReturnType stopDispose(String ip, DisposeCapacityType capType, @Nullable NetflowDirection nfDirection, - @Nullable Integer attackType, @Nullable String taskId) { - return null; + @Nullable Integer attackType, + @Nullable String taskId) { + String url = urlRootPath + "dispose_device/task/stop"; + + if (token == null || token.length() == 0) { + return new MulReturnType<>(ErrorCode.ERR_LOGOUT, null); + } + + try { + log.debug("++++Begging YiYang Stop Cleanup Task: {}", taskId); + + if (capType != DisposeCapacityType.BLACKHOOL) { + log.error("----Error YiYang don't support dispose capacity type: {}", capType); + return new MulReturnType<>(ErrorCode.ERR_UNSUPPORT, null); + } + + if (taskId == null) { + return new MulReturnType<>(ErrorCode.ERR_PARAMS, null); + } + + ProtocolRespDTO rspInfo = restfulInterface.stopDisposeTask(url, token, + new String[]{taskId}); + + // 判断是否token过期 + if (rspInfo != null && rspInfo.getCode() == HttpServletResponse.SC_UNAUTHORIZED) { + // 重新登录获取 token + upgradeToken(); + rspInfo = restfulInterface.stopDisposeTask(url, token, new String[]{taskId}); + } + + if (rspInfo == null || rspInfo.getCode() != ErrorCode.ERR_OK.getHttpCode()) { + log.error("----Error YiYang stop task{} server return error", taskId); + return new MulReturnType<>(ErrorCode.ERR_HAOHAN_ERROR, null); + } + + log.debug("----Finish YiYang Stop Cleanup Task: {}", taskId); + return new MulReturnType<>(ErrorCode.ERR_OK, null); + } catch (Exception ex) { + log.error("----Exception YiYang Stop Cleanup Task: {}, {}, {}", ip, nfDirection, taskId); + return new MulReturnType<>(ErrorCode.ERR_SYSTEMEXCEPTION, null); + } } /** diff --git a/src/main/java/com/dispose/service/impl/DeviceTaskManagerServiceImpl.java b/src/main/java/com/dispose/service/impl/DeviceTaskManagerServiceImpl.java index 188620c6..87371e5c 100644 --- a/src/main/java/com/dispose/service/impl/DeviceTaskManagerServiceImpl.java +++ b/src/main/java/com/dispose/service/impl/DeviceTaskManagerServiceImpl.java @@ -31,6 +31,9 @@ import java.util.concurrent.ConcurrentHashMap; @Slf4j public class DeviceTaskManagerServiceImpl implements DeviceTaskManagerService { + /** + * The Task cache. + */ private final ConcurrentHashMap taskCache = new ConcurrentHashMap<>(); /** @@ -413,8 +416,87 @@ public class DeviceTaskManagerServiceImpl implements DeviceTaskManagerService { } } + /** + * Yi yang device task run. + * + * @param ai the ai + * @param deviceTask the device task + * @param disposeTask the dispose task + */ private void yiYangDeviceTaskRun(AbilityInfo ai, DeviceTask deviceTask, DisposeTask disposeTask) { + MulReturnType ret; + // 重试错误次数过多 + if (deviceTask.getErrRetry() >= DisposeConfigValue.CALL_ERROR_RETRY_TIMES) { + return; + } + + // 设置任务状态为启动中 + deviceTaskManager.changeDisposeDeviceTaskInfoStatus(deviceTask.getId(), DisposeTaskStatus.TASK_STARTING); + // 设置启动任务攻击类型状态 + deviceTaskManager.setExecAttackType(deviceTask.getId(), NetflowDirection.DIRECTION_BI, + deviceTask.getTaskAttackType()); + + ret = ai.getDb().runDispose(disposeTask.getDisposeObject(), disposeTask.getDisposeCapacity(), null, null, + (long) -(Helper.getTimestampDiffNow(disposeTask.getPlanEndTime()) / 60)); + + if (ret.getFirstParam() == ErrorCode.ERR_OK) { + // 设置攻击类型任务启动结果 + deviceTaskManager.setAttackTypeStatus(deviceTask.getId(), + disposeTask.getFlowDirection(), deviceTask.getTaskAttackType()); + // 更改处置任务状态为处置中 + deviceTaskManager.changeDisposeDeviceTaskInfoStatus(deviceTask.getId(), DisposeTaskStatus.TASK_STARTED); + + // 记录亿阳设备返回的任务ID + deviceTaskManager.setTaskExternId(deviceTask.getId(), ret.getSecondParam()); + + log.info("YIYANG_PLATFORM setup task succeed: {}, device taskId {}", deviceTask, ret.getSecondParam()); + + // 重置错误尝试次数 + deviceTaskManager.setTaskErrRetryTimes(deviceTask.getId(), 0); + } else if (deviceTask.getErrRetry() < DisposeConfigValue.CALL_ERROR_RETRY_TIMES) { + // 设置该任务为新任务,待下次重试启动 + deviceTaskManager.changeDisposeDeviceTaskInfoStatus(deviceTask.getId(), DisposeTaskStatus.TASK_NEW); + // 记录任务出错重试次数 + deviceTaskManager.setTaskErrRetryTimes(deviceTask.getId(), deviceTask.getErrRetry() + 1); + log.error("YIYANG_PLATFORM setup task times {} error {}: {}", deviceTask.getErrRetry(), + ret.getSecondParam(), deviceTask); + } else { + // 任务出错,不在重试,当做失败任务处理 + deviceTaskManager.setAttackTypeStatus(deviceTask.getId(), + disposeTask.getFlowDirection(), 0L); + log.error("YIYANG_PLATFORM setup task error {}: {}", ret.getFirstParam(), deviceTask); + } + } + + /** + * Yi yang device task stop. + * + * @param ai the ai + * @param deviceTask the device task + * @param disposeTask the dispose task + */ + private void yiYangDeviceTaskStop(AbilityInfo ai, DeviceTask deviceTask, DisposeTask disposeTask) { + MulReturnType ret; + + // 停止处置任务 + ret = ai.getDb().stopDispose(disposeTask.getDisposeObject(), disposeTask.getDisposeCapacity(), null, null, + deviceTask.getExternId()); + + if (ret.getFirstParam() == ErrorCode.ERR_OK) { + log.info("YIYANG_PLATFORM stop task succeed: {}, device taskId {}", deviceTask, ret.getSecondParam()); + // 设置任务状态为结束 + deviceTaskManager.changeDisposeDeviceTaskInfoStatus(deviceTask.getId(), DisposeTaskStatus.TASK_FINISHED); + deviceTaskManager.setTaskErrRetryTimes(deviceTask.getId(), 0); + deviceTaskManager.setAttackTypeStatus(deviceTask.getId(), disposeTask.getFlowDirection(), 0L); + } else if (deviceTask.getErrRetry() < DisposeConfigValue.CALL_ERROR_RETRY_TIMES) { + // 记录任务出错重试次数 + deviceTaskManager.setTaskErrRetryTimes(deviceTask.getId(), deviceTask.getErrRetry() + 1); + log.error("YIYANG_PLATFORM stop task times {} error {}: {}", deviceTask.getErrRetry(), + ret.getSecondParam(), deviceTask); + } else { + log.error("YIYANG_PLATFORM stop task error {}: {}", ret.getFirstParam(), deviceTask); + } } /** @@ -565,6 +647,10 @@ public class DeviceTaskManagerServiceImpl implements DeviceTaskManagerService { virtualDeviceTaskStop(ai, v, task); break; + case YIYANG_PLATFORM: + yiYangDeviceTaskStop(ai, v, task); + break; + default: log.error("Unknown dispose device type: {}", ai.getDev()); break; diff --git a/src/main/java/com/dispose/setup/SystemInitial.java b/src/main/java/com/dispose/setup/SystemInitial.java index 0d51d8c1..de7c8095 100644 --- a/src/main/java/com/dispose/setup/SystemInitial.java +++ b/src/main/java/com/dispose/setup/SystemInitial.java @@ -51,6 +51,6 @@ public class SystemInitial implements CommandLineRunner { // 初始化处置能力设备 setupAbilityDevice(); - RestfulInterface.initEnv(1000); + RestfulInterface.initEnv(5000); } }