cmhi
/
dispose
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

OCT 

REM:
1. 增加亿阳设备启动、停止处置任务功能
This commit is contained in:
HuangXin 2020-10-10 18:07:14 +08:00
parent 49aadd12a3
commit ed9c273be6
3 changed files with 219 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

136
src/main/java/com/dispose/ability/impl/YiYangAbilityImpl.java
View File

@ -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<ErrorCode, String> runDispose(String ip, DisposeCapacityType capType,
@Nullable NetflowDirection nfDirection,
@Nullable Integer attackType, @Nullable Long duration) {
return null;
@Nullable Integer attackType,
@Nullable Long duration) {
List<YiYangStartTaskItem> 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<YiYangStartTaskRsp> 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<ErrorCode, Long> 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<YiYangStopTaskRsp> 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);
}
}
/**

86
src/main/java/com/dispose/service/impl/DeviceTaskManagerServiceImpl.java
View File

@ -31,6 +31,9 @@ import java.util.concurrent.ConcurrentHashMap;
@Slf4j
public class DeviceTaskManagerServiceImpl implements DeviceTaskManagerService {
/**
* The Task cache.
*/
private final ConcurrentHashMap<String, Boolean> 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<ErrorCode, String> 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<ErrorCode, Long> 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;

2
src/main/java/com/dispose/setup/SystemInitial.java
View File

@ -51,6 +51,6 @@ public class SystemInitial implements CommandLineRunner {
// 初始化处置能力设备
setupAbilityDevice();
RestfulInterface.initEnv(1000);
RestfulInterface.initEnv(5000);
}
}