域名查询增加查询模式参数query.mode=2

#update 2 第三方api优先，并更新数据库
#local 1 只查db
#general 0 db优先，不存在则查第三方api

4 files changed, 140 insertions, 84 deletions

diff --git a/src/main/java/com/mesasoft/cn/sketch/service/impl/DomainServiceImpl.java b/src/main/java/com/mesasoft/cn/sketch/service/impl/DomainServiceImpl.java
index 2388e96..3d610f7 100644
--- a/src/main/java/com/mesasoft/cn/sketch/service/impl/DomainServiceImpl.java
+++ b/src/main/java/com/mesasoft/cn/sketch/service/impl/DomainServiceImpl.java
@@ -63,8 +63,11 @@ public class DomainServiceImpl implements DomainService {
     @Value("${query.readin.batch}")
     private int queryReadinBatch;
 
-    @Value("${query.isLocal}")
-    private boolean isLocal;
+    @Value("${query.mode}")
+    private int queryMode;
+
+//    private static ThreadLocal<Integer> bcSuccessCounter = ThreadLocal.withInitial(() -> 0);
+//    private static ThreadLocal<Integer> bcFailureCounter = ThreadLocal.withInitial(() -> 0);
 
     @Override
     public List<DomainInfo> getCategoryInfo(List<String> domains, String username) throws Exception {
@@ -73,7 +76,80 @@ public class DomainServiceImpl implements DomainService {
         logger.info("domain category query distinct total : {}", domains.size());
         int queryNum = domains.size();
         List<DomainInfo> results = new ArrayList<DomainInfo>();
-        //1. 查询本地数据库
+        if (queryMode == 1) {
+            //1. 查询本地数据库
+            getDomainsByDB(domains, results);
+        } else if (queryMode == 2) {
+            //2. 查第三方api ，并更新数据库
+            getDomainByBrightCloud(domains, username, queryNum, results, 0);
+            getDomainsByDB(domains, results);
+        } else {
+            //3. 查数据库，再查第三方api
+            int domainsByDBSize = getDomainsByDB(domains, results);
+            getDomainByBrightCloud(domains, username, queryNum, results, domainsByDBSize);
+        }
+        return results;
+    }
+
+    /**
+     * @param domains
+     * @param username
+     * @param queryNum
+     * @param results
+     * @param dbResultNum
+     * @throws IOException
+     */
+    private void getDomainByBrightCloud(List<String> domains, String username, int queryNum, List<DomainInfo> results, int dbResultNum) throws IOException {
+        //2. 调用api查询
+        int apiResultNum = 0;
+        int failedQueryNum = 0;
+        List<DomainCategory> bcResults = new ArrayList<>();
+        //批量查询API
+        BrightCloud brightCloud = new BrightCloud();
+        List<List<String>> apiPartitions = Lists.partition(domains, bcConfig.getMaximumQueryNum());
+        List<String> bcDomainRecords = new ArrayList<>();
+        for (List<String> partition : apiPartitions) {
+            List<DomainInfo> bcDomainInfos = new ArrayList<>();
+            List<DomainCategory> recordsFromBcApi = brightCloud.getBrightCloudDomainCategory(partition);
+            for (DomainCategory record : recordsFromBcApi) {
+                //查询成功的结果
+                if (record.getQuery_success().equals(true)) {
+                    record.setSubmit_user(username);
+                    record.setCreate_time(new Date());
+                    record.setUpdate_time(new Date());
+
+                    bcResults.add(record);
+
+                    bcDomainInfos.add(DomainInfo.builder()
+                            .domain(record.getFqdn())
+                            .result(Lists.newArrayList(record))
+                            .build());
+                    apiResultNum += 1;
+//                    bcSuccessCounter.set(bcSuccessCounter.get() + 1);
+                    bcDomainRecords.add(record.getFqdn());
+                } else {
+                    failedQueryNum += 1;
+//                    bcFailureCounter.set(bcFailureCounter.get() + 1);
+                }
+            }
+            if (bcResults.size() > 0) {
+                insertBatchDomains(bcResults);
+                results.addAll(bcDomainInfos);
+                bcResults.clear();
+                bcDomainInfos.clear();
+            }
+            processLog("category", queryNum, apiResultNum, failedQueryNum, dbResultNum);
+        }
+        //移除查询到的域名
+        domains.removeAll(bcDomainRecords);
+//        // 记录api调用次数
+        FileWriter fileWriter = new FileWriter(bcConfig.getUsereportFilePath(), true);
+        fileWriter.write(DateUtil.date() + "," + " domain category query, brightCloud, success records " + apiResultNum + ", fail records " + failedQueryNum + "\n");
+        fileWriter.close();
+        logger.info("domain category query brightCloud success records : " + apiResultNum + ",fail records : " + failedQueryNum);
+    }
+
+    private int getDomainsByDB(List<String> domains, List<DomainInfo> results) {
         List<List<String>> partitionDomains = Lists.partition(domains, appConfig.getDbQueryBatchSize());
         List<String> dBDomains = new ArrayList<>();
         for (List<String> partitionDomain : partitionDomains) {
@@ -84,54 +160,7 @@ public class DomainServiceImpl implements DomainService {
         }
         domains.removeAll(dBDomains);
         logger.info("domain category query DB record {}", results.size());
-
-        //2. 调用api查询
-        int apiResultNum = 0;
-        int failedQueryNum = 0;
-        int dbResultNum = dBDomains.size();
-
-        if (!isLocal && domains.size() > 0) {
-            List<DomainCategory> bcResults = new ArrayList<>();
-            //批量查询API
-            BrightCloud brightCloud = new BrightCloud();
-            List<List<String>> apiPartitions = Lists.partition(domains, bcConfig.getMaximumQueryNum());
-            for (List<String> partition : apiPartitions) {
-                List<DomainInfo> bcDomainInfos = new ArrayList<>();
-                List<DomainCategory> recordsFromBcApi = brightCloud.getBrightCloudDomainCategory(partition);
-                for (DomainCategory record : recordsFromBcApi) {
-                    //查询成功的结果
-                    if (record.getQuery_success().equals(true)) {
-                        record.setSubmit_user(username);
-                        record.setCreate_time(new Date());
-                        record.setUpdate_time(new Date());
-
-                        bcResults.add(record);
-                        bcDomainInfos.add(DomainInfo.builder()
-                                .domain(record.getFqdn())
-                                .result(Lists.newArrayList(record))
-                                .build());
-                        apiResultNum += 1;
-                    } else {
-                        failedQueryNum += 1;
-                    }
-                }
-                if (bcResults.size() > 0) {
-                    insertBatchDomains(bcResults);
-                    results.addAll(bcDomainInfos);
-                    bcResults.clear();
-                    bcDomainInfos.clear();
-                }
-                processLog("category", queryNum, apiResultNum, failedQueryNum, dbResultNum);
-            }
-
-            // 记录api调用次数
-            FileWriter fileWriter = new FileWriter(bcConfig.getUsereportFilePath(), true);
-            fileWriter.write(DateUtil.date() + "," + " domain category query brightCloud success record " + apiResultNum + ",fail query record " + failedQueryNum + "\n");
-            fileWriter.close();
-            logger.info("domain category query brightCloud success record : " + apiResultNum + ",fail query record : " + failedQueryNum);
-        }
-
-        return results;
+        return dBDomains.size();
     }
 
     @Override
@@ -162,12 +191,12 @@ public class DomainServiceImpl implements DomainService {
         domains.removeAll(dBDomains);
         logger.info("domain whois query DB record {}", results.size());
 
-        //2. 调用api
-        int apiResultNum = 0;
-        int failedQueryNum = 0;
-        int dbResultNum = dBDomains.size();
+        if (queryMode == 1 && domains.size() > 0) {
+            //2. 调用api
+            int apiResultNum = 0;
+            int failedQueryNum = 0;
+            int dbResultNum = dBDomains.size();
 
-        if (!isLocal && domains.size() > 0) {
             List<DomainWhois> chinazResults = new ArrayList<>();
             ChinaZ chinaz = new ChinaZ();
             List<List<String>> apiPartitions = Lists.partition(domains, chinazConfig.getMaximumQueryNum());
@@ -201,9 +230,9 @@ public class DomainServiceImpl implements DomainService {
             // 记录api调用次数
             try {
                 FileWriter fileWriter = new FileWriter(chinazConfig.getUsereportFilePath(), true);
-                fileWriter.write(DateUtil.date() + "," + " domain whois query chinaz success record " + apiResultNum + ",fail query record " + failedQueryNum + "\n");
+                fileWriter.write(DateUtil.date() + "," + " domain whois query, chinaz, success records " + apiResultNum + ", fail records " + failedQueryNum + "\n");
                 fileWriter.close();
-                logger.info("domain whois query chinaz success record " + apiResultNum + ",fail query record " + failedQueryNum);
+                logger.info("domain whois query chinaz success records " + apiResultNum + ",fail records " + failedQueryNum);
             } catch (IOException e) {
                 logger.error(e);
             }
@@ -229,7 +258,7 @@ public class DomainServiceImpl implements DomainService {
         domains.removeAll(dBDomains);
         logger.info("domain ICP query DB record {}", results.size());
 
-        if (!isLocal && domains.size() > 0) {
+        if (queryMode == 1 && domains.size() > 0) {
             //2. 调用api
             int apiResultNum = 0;
             int failedQueryNum = 0;
@@ -267,9 +296,9 @@ public class DomainServiceImpl implements DomainService {
             // 记录api调用次数
             try {
                 FileWriter fileWriter = new FileWriter(chinazConfig.getUsereportFilePath(), true);
-                fileWriter.write(DateUtil.date() + "," + " domain ICP query chinaz success record " + apiResultNum + ",fail query record " + failedQueryNum + "\n");
+                fileWriter.write(DateUtil.date() + "," + " domain ICP query, chinaz, success records " + apiResultNum + ", fail records " + failedQueryNum + "\n");
                 fileWriter.close();
-                logger.info("domain ICP query chinaz success record " + apiResultNum + ",fail query record " + failedQueryNum);
+                logger.info("domain ICP query chinaz success records " + apiResultNum + ",fail records " + failedQueryNum);
             } catch (IOException e) {
                 logger.error(e);
             }
@@ -301,12 +330,9 @@ public class DomainServiceImpl implements DomainService {
 
     @Override
     public Map getDomainInfoByFile(String requestUser, String srcFile, String queryType, Boolean export) throws Exception {
-        //输出文件
+        //查询的源文件
         String fileName = String.join("-", FileUtil.getPrefix(srcFile), queryType, DateUtil.format(new Date(), "yyyyMMdd-HHmmss") + ".json");
-        File resultFile = FileUtil.touch(queryOutputDir + File.separator + fileName);
-
         Map<String, Object> resultMsg = new HashMap<>();
-        resultMsg.put("result", resultFile.toString());
         resultMsg.put("export", export);
 
         // 文件读取
@@ -316,7 +342,7 @@ public class DomainServiceImpl implements DomainService {
         int resultSize = 0;
         int querySize = 0;
         List<DomainInfo> queryResults = new ArrayList<>();
-        cn.hutool.core.io.file.FileWriter fileWriter = new cn.hutool.core.io.file.FileWriter(resultFile);
+
         for (List<String> domains : partition) {
             switch (queryType) {
                 case "category":
@@ -329,24 +355,30 @@ public class DomainServiceImpl implements DomainService {
                     queryResults = getICPInfo(domains);
                     break;
             }
-            querySize +=domains.size();
+            querySize += domains.size();
             resultSize += queryResults.size();
             if (ObjectUtil.isNotEmpty(queryResults) && export) {
-                List<String> prettyResults = queryResults.stream().map(x -> JSON.toJSONString(x)).collect(Collectors.toList());
-                //List<String> prettyResults = queryResults.stream().map(x -> JSON.toJSONString(x, true)).collect(Collectors.toList());
-                fileWriter.appendLines(prettyResults);
-                logger.info("[File query]-" + srcFile +" result: " + queryResults.size() + " Results saved in " + resultFile.getAbsolutePath());
+                exportFile(srcFile, fileName, resultMsg, querySize, queryResults);
             }
             Double queryTotal = Double.valueOf(readUtf8Lines.size());
-            logger.info("[File query]- " + queryType + " query result size: " + resultSize  +", query size: " + querySize +", process: "+ NumberUtil.decimalFormat("#.##%", querySize/queryTotal));
+            logger.info("[File query]- " + queryType + " query result records: " + resultSize + ", query records: " + querySize + ", process: " + NumberUtil.decimalFormat("#.##%", querySize / queryTotal));
             queryResults.clear();
-
         }
-        resultMsg.put("total_export", querySize);
-
         return resultMsg;
     }
 
+    private void exportFile(String srcFile, String fileName, Map<String, Object> resultMsg, int querySize, List<DomainInfo> queryResults) {
+        File resultFile = FileUtil.touch(queryOutputDir + File.separator + fileName);
+        cn.hutool.core.io.file.FileWriter fileWriter = new cn.hutool.core.io.file.FileWriter(resultFile);
+        List<String> prettyResults = queryResults.stream().map(x -> JSON.toJSONString(x)).collect(Collectors.toList());
+        //List<String> prettyResults = queryResults.stream().map(x -> JSON.toJSONString(x, true)).collect(Collectors.toList());
+        fileWriter.appendLines(prettyResults);
+        resultMsg.put("filePath", resultFile.toString());
+        resultMsg.put("fileSize", resultFile.length());
+        resultMsg.put("records", querySize);
+        logger.info("[File query]-" + srcFile + " result: " + queryResults.size() + " Results saved in " + resultFile.getAbsolutePath());
+    }
+
     /**
      * 域名入库
      *
@@ -386,7 +418,7 @@ public class DomainServiceImpl implements DomainService {
     private void processLog(String logType, int queryNum, int apiResultNum, int failedQueryNum, int dbResultNum) {
         logger.info("[domain " + logType + "]-"
                 + "Query result: submit " + queryNum + " valid objects, "
-                + dbResultNum + " (" + new DecimalFormat("##.0%").format((float) dbResultNum / queryNum) + ")" + " results from database,"
+                + dbResultNum + " (" + new DecimalFormat("##.0%").format((float) dbResultNum / queryNum) + ")" + " results from database. "
                 + apiResultNum + " (" + new DecimalFormat("##.0%").format((float) apiResultNum / queryNum) + ")" + " results from api. "
                 + failedQueryNum + " (" + new DecimalFormat("##.0%").format((float) failedQueryNum / queryNum) + ")" + " failed queries,");
     }
diff --git a/src/main/resources/application-sketch.properties b/src/main/resources/application-sketch.properties
index 0dacc23..3567192 100644
--- a/src/main/resources/application-sketch.properties
+++ b/src/main/resources/application-sketch.properties
@@ -1,5 +1,6 @@
 #datasource
-spring.datasource.url=jdbc:mysql://192.168.44.12:3306/web_sketch_v2?useUnicode=true&characterEncoding=utf-8&useSSL=true
+#spring.datasource.url=jdbc:mysql://api.geedge.net:3306/web_sketch_v2?useUnicode=true&characterEncoding=utf-8&useSSL=true&autoReconnect=true&failOverReadOnly=false
+spring.datasource.url=jdbc:mysql://192.168.44.12:3306/web_sketch_v2?useUnicode=true&characterEncoding=utf-8&useSSL=true&autoReconnect=true&failOverReadOnly=false
 spring.datasource.username=root
 spring.datasource.password=galaxy2019
 
@@ -17,7 +18,10 @@ database = web_sketch_v2
 db.query.batch.size = 10000
 
 ###################### api #########################
-query.isLocal=true
+#update 2 第三方api优先，并更新数据库
+#local 1 只查db
+#general 0 db优先，不存在则查第三方api
+query.mode=2
 ##### bright cloud #######
 bright-cloud.oemid = GeedgeNet
 bright-cloud.deviceid = TSG-Dev
@@ -32,8 +36,8 @@ bright-cloud.isReputation = 1
 bright-cloud.isxml = 0
 # api单次查询url长度限制 API最高限制
 bright-cloud.maximum-query-num = 100
-bright-cloud.cateinfo-filepath = categoryinfo.json
-bright-cloud.usereport-filepath = bright_cloud_query_count.csv
+bright-cloud.cateinfo-filepath = C:\\code\\cn\\webskt-query-agent\\src\\main\\resources\\categoryinfo.json
+bright-cloud.usereport-filepath = C:\\code\\cn\\webskt-query-agent\\src\\main\\resources\\bright_cloud_query_count.csv
 
 ######### chinaz #########
 chinaz.url-single = https://apidatav2.chinaz.com/single
diff --git a/src/main/resources/mapper/DomainDao.xml b/src/main/resources/mapper/DomainDao.xml
index e326092..9c31247 100644
--- a/src/main/resources/mapper/DomainDao.xml
+++ b/src/main/resources/mapper/DomainDao.xml
@@ -170,6 +170,22 @@
                 #{submit_user,jdbcType=VARCHAR},
                 #{create_time,jdbcType=TIMESTAMP},
                 #{update_time,jdbcType=TIMESTAMP})
+        ON DUPLICATE KEY UPDATE
+                `source` = VALUES(`source`),
+                query_success = VALUES(query_success),
+                match_pattern = VALUES(match_pattern),
+                reputation_score = VALUES(reputation_score),
+                reputation_level = VALUES(reputation_level),
+                category_id = VALUES(category_id),
+                category_name = VALUES(category_name),
+                category_group = VALUES(category_group),
+                category_conf = VALUES(category_conf),
+                is_a1_cat = VALUES(is_a1_cat),
+                status_code = VALUES(status_code),
+                submit_user = VALUES(submit_user),
+                create_time = VALUES(create_time),
+                update_time = VALUES(update_time);
+
     </insert>
 
     <insert id="insertWhois" keyColumn="id" keyProperty="id"
diff --git a/src/test/java/com/mesasoft/cn/WebSketchApplicationTest.java b/src/test/java/com/mesasoft/cn/WebSketchApplicationTest.java
index dc5f089..6c6e1bd 100644
--- a/src/test/java/com/mesasoft/cn/WebSketchApplicationTest.java
+++ b/src/test/java/com/mesasoft/cn/WebSketchApplicationTest.java
@@ -1,7 +1,9 @@
 package com.mesasoft.cn;
 
 import cn.hutool.core.date.DateUtil;
+import com.mesasoft.cn.sketch.config.BrightCloudConfig;
 import com.mesasoft.cn.sketch.config.ChinazConfig;
+import com.mesasoft.cn.sketch.util.FileUtils;
 import com.zhazhapan.config.JsonParser;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -9,6 +11,7 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.test.context.junit4.SpringRunner;
 
+import javax.annotation.Resource;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
@@ -24,13 +27,14 @@ public class WebSketchApplicationTest {
             e.printStackTrace();
         }
     }
+
+    @Resource
+    BrightCloudConfig brightCloudConfig;
     @Test
     public void contextLoads() {
 
-        List<String> domain = new ArrayList<>();
-        domain.add("a");
-        domain.remove("a");
-        System.err.println(DateUtil.date());
+        String s = FileUtils.readJsonFile(brightCloudConfig.getCateinfoFilepath());
+        System.err.println(s);
     }
 
 }