gfwleak/galaxy-tsg-olap-dos-detection-job
Archived
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

TSG-15167 新增知识库文件校验功能

Browse Source
This commit is contained in:
unknown
2023-05-23 10:38:15 +08:00
parent d8b0a7637b
commit 04ee45f77d
2 changed files with 225 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
src/main/java/com/zdjizhi/source/HttpSource.java
View File

@@ -2,6 +2,8 @@ package com.zdjizhi.source;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.io.file.FileReader;
import cn.hutool.crypto.digest.DigestUtil;
import cn.hutool.json.JSONObject;
import com.alibaba.nacos.api.NacosFactory;
import com.alibaba.nacos.api.PropertyKeyConst;
@@ -28,7 +30,7 @@ import java.util.concurrent.Executor;
public class HttpSource extends RichHttpSourceFunction<Map<String, byte[]>> {
private static final int TRY_TIMES = 3;
private static final Logger logger = LoggerFactory.getLogger(HttpSource.class);
private static final String EXPR = "$.[?(@.version=='latest' && @.name in ['ip_v4_built_in','ip_v6_built_in','ip_v4_user_defined','ip_v6_user_defined'])].['name','sha256','format','path']";
@@ -50,11 +52,19 @@ public class HttpSource extends RichHttpSourceFunction<Map<String, byte[]>> {
private ConfigService configService;
private static Header header;
//运行状态cancel时置为false
private boolean isRunning = true;
//是否下发，默认不发送
private boolean isSending = false;
private static HttpClientUtils2 httpClientUtils;
// private static JsonMapper jsonMapperInstance = JsonMapper.getInstance();
// private static JavaType listType = jsonMapperInstance.createCollectionType(List.class, KnowledgeLog.class);
private static Map<String, String> updateMap = new HashMap<>();
private static Map<String, String> knowledgeMetaCache = new HashMap<>();
private static HashMap<String, byte[]> knowledgeFileCache;
private boolean isRunning = true;
// private boolean isRunning = true;
public HttpSource(Properties nacosProperties, String NACOS_DATA_ID, String NACOS_GROUP, long NACOS_READ_TIMEOUT, String storePath) {
@@ -68,10 +78,12 @@ public class HttpSource extends RichHttpSourceFunction<Map<String, byte[]>> {
@Override
public void open(Configuration parameters) throws Exception {
super.open(parameters);
httpClientUtils = new HttpClientUtils2();
//初始化元数据缓存
updateMap = new HashMap<>(16);
knowledgeMetaCache = new HashMap<>(16);
//初始化定位库缓存
knowledgeFileCache = new HashMap<>(16);
header = new BasicHeader("token", CommonConfig.HOS_TOKEN);
logger.info("连接nacos：" + nacosProperties.getProperty(PropertyKeyConst.SERVER_ADDR));
configService = NacosFactory.createConfigService(nacosProperties);
}
@@ -86,6 +98,9 @@ public class HttpSource extends RichHttpSourceFunction<Map<String, byte[]>> {
if (StringUtil.isNotBlank(config)) {
ArrayList<Object> metaList = JsonPath.parse(config).read(EXPR);
loadKnowledge(metaList);
if (isSending){
ctx.collect(knowledgeFileCache);
}
}
@@ -108,13 +123,14 @@ public class HttpSource extends RichHttpSourceFunction<Map<String, byte[]>> {
knowledgeJson.getStr("format"));
String sha256 = knowledgeJson.getStr("sha256");
String filePath = knowledgeJson.getStr("path");
if (!sha256.equals(updateMap.get(fileName))) {
updateMap.put(fileName, sha256);
updateKnowledge(fileName, filePath);
if (!sha256.equals(knowledgeMetaCache.get(fileName))) {
knowledgeMetaCache.put(fileName, sha256);
updateKnowledge(fileName, filePath,sha256);
}
}
ctx.collect(knowledgeFileCache);
if (isSending){
ctx.collect(knowledgeFileCache);
}
}
}
@@ -126,7 +142,11 @@ public class HttpSource extends RichHttpSourceFunction<Map<String, byte[]>> {
});
while (isRunning) {
Thread.sleep(10000);
try {
Thread.sleep(10000);
}catch (InterruptedException e) {
e.printStackTrace();
}
}
}
@@ -141,38 +161,92 @@ public class HttpSource extends RichHttpSourceFunction<Map<String, byte[]>> {
knowledgeJson.getStr("format"));
String sha256 = knowledgeJson.getStr("sha256");
String filePath = knowledgeJson.getStr("path");
Header header = new BasicHeader("token", CommonConfig.HOS_TOKEN);
HttpClientUtils2 httpClientUtils = new HttpClientUtils2();
inputStream = httpClientUtils.httpGetInputStream(filePath, 3000, header);
updateMap.put(fileName, sha256);
knowledgeFileCache.put(fileName, IOUtils.toByteArray(inputStream));
byte[] localFileByte = getLocalFile(fileName);
String localFileSha256Hex = DigestUtil.sha256Hex(localFileByte);
if (sha256.equals(localFileSha256Hex)){
logger.info("本地文件{}的sha256为:{} ,Nacos内记录为:{} ,sha256相等", fileName, localFileSha256Hex, sha256);
knowledgeMetaCache.put(fileName, sha256);
// knowledgeFileCache.put(fileName, IOUtils.toByteArray(inputStream));
knowledgeFileCache.put(fileName, localFileByte);
}else {
logger.info("本地文件{}的sha256为:{} ,Nacos内记录为:{} ,sha256不相等，更新本地文件及缓存", fileName, localFileSha256Hex, sha256);
updateKnowledge(fileName,filePath,sha256);
}
}
}
// } catch (IOException ioException) {
} catch (Exception ioException) {
ioException.printStackTrace();
} finally {
IOUtils.closeQuietly(inputStream);
}
}
private void updateKnowledge(String fileName, String filePath,String sha256) {
InputStream inputStream = null;
// FileOutputStream outputStream = null;
int retryNum = 0;
try {
while (retryNum < TRY_TIMES){
inputStream = httpClientUtils.httpGetInputStream(filePath, 90000, header);
if (inputStream !=null){
byte[] downloadBytes = IOUtils.toByteArray(inputStream);
if (downloadBytes.length>0){
String downloadFileSha256Hex = DigestUtil.sha256Hex(downloadBytes);
if (sha256.equals(downloadFileSha256Hex)){
logger.info("通过HOS下载{}的sha256为:{} ,Nacos内记录为:{} ,sha256相等", fileName, downloadBytes, sha256);
// HdfsUtils.uploadFileByBytes(CommonConfig.HDFS_PATH + fileName, downloadBytes);
knowledgeMetaCache.put(fileName,sha256);
knowledgeFileCache.put(fileName, downloadBytes);
updateLocalFile(fileName);
retryNum = TRY_TIMES;
isSending = true;
}else {
logger.error("通过HOS下载{}的sha256为:{} ,Nacos内记录为:{} ,sha256不相等 开始第{}次重试下载文件", fileName, downloadFileSha256Hex, sha256, retryNum);
retryNum++;
}
}
}
}
} catch (IOException ioException) {
ioException.printStackTrace();
} finally {
IOUtils.closeQuietly(inputStream);
// IOUtils.closeQuietly(outputStream);
}
}
private void updateKnowledge(String fileName, String filePath) {
InputStream inputStream = null;
private void updateLocalFile(String fileName) {
FileOutputStream outputStream = null;
try {
Header header = new BasicHeader("token", CommonConfig.HOS_TOKEN);
HttpClientUtils2 httpClientUtils = new HttpClientUtils2();
inputStream = httpClientUtils.httpGetInputStream(filePath, 3000, header);
byte[] bytes = IOUtils.toByteArray(inputStream);
HdfsUtils.uploadFileByBytes(CommonConfig.HDFS_PATH + fileName, bytes);
knowledgeFileCache.put(fileName, bytes);
} catch (IOException ioException) {
ioException.printStackTrace();
HdfsUtils.uploadFileByBytes(CommonConfig.HDFS_PATH + fileName, knowledgeFileCache.get(fileName));
} catch (IOException ioe) {
logger.error("更新本地文件{}时发生IO异常,异常信息为:", fileName, ioe.getMessage());
ioe.printStackTrace();
} catch (RuntimeException e) {
logger.error("更新本地文件{}时发生异常,异常信息为:", fileName, e.getMessage());
e.printStackTrace();
} finally {
IOUtils.closeQuietly(inputStream);
IOUtils.closeQuietly(outputStream);
}
}
private static byte[] getLocalFile(String name) {
byte[] fileBytes = null;
try {
fileBytes = HdfsUtils.getFileBytes(CommonConfig.HDFS_PATH + name) ;
} catch (RuntimeException | IOException e) {
logger.error("IpLookupUtils download MMDB files error, message is:" + e.getMessage());
e.printStackTrace();
}
return fileBytes;
}
@Override
public void cancel() {
this.isRunning = false;