galaxy-tsg-olap-log-stream-doublewrite/src/main/java/com/zdjizhi/tools/general/IpLookupUtils.java

package com.zdjizhi.tools.general;

import cn.hutool.crypto.digest.DigestUtil;
import cn.hutool.log.Log;
import cn.hutool.log.LogFactory;
import com.alibaba.fastjson2.*;
import com.geedgenetworks.utils.IpLookupV2;
import com.geedgenetworks.utils.StringUtil;
import com.google.common.base.Joiner;
import com.zdjizhi.common.CommonConfig;
import com.zdjizhi.common.FlowWriteConfig;
import com.zdjizhi.common.pojo.KnowlegeBaseMeta;
import com.zdjizhi.tools.connections.http.HttpClientService;
import org.apache.http.client.utils.URIBuilder;

import java.io.ByteArrayInputStream;
import java.net.URISyntaxException;
import java.util.HashMap;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;


/**
 * @author wangchengcheng
 * @version 2023/11/10 15:23
 */
public class IpLookupUtils {
    private static final Log logger = LogFactory.get();
    private static final String ipBuiltInName = "ip_builtin.mmdb";
    private static final String ipUserDefinedName = "ip_user_defined.mmdb";

    private static final String asnName = "asn_builtin.mmdb";

    /**
     * ip定位库
     */
    private static IpLookupV2 ipLookup;

    /**
     * 定位库默认分隔符
     */
    private static final String LOCATION_SEPARATOR = ".";

    /**
     * 最大重试次数
     */
    private static final int TRY_TIMES = 5;

    /**
     * http connections
     */
    private static final HttpClientService httpClientService;

    /**
     * 定位库元数据缓存
     */
    private static final HashMap<String, KnowlegeBaseMeta> knowledgeMetaCache = new HashMap<>(16);

    private static String currentSha256IpUserDefined = "";

    private static String currentSha256IpBuiltin = "";


    private static String currentSha256AsnBuiltin = "";

    static {
        httpClientService = new HttpClientService();
        try {
            stuffKnowledgeMetaCache();
            Timer timer = new Timer();
            timer.schedule(new TimerTask() {
                @Override
                public void run() {
                    stuffKnowledgeMetaCache();

                }
            }, 0, FlowWriteConfig.KNOWLEDGE_EXECUTION_MINUTES * 1000 * 60);
        } catch (Exception e) {
            logger.error("知识库加载失败,失败原因为:" + e);
        }
    }


    private static void stuffKnowledgeMetaCache() {
        final KnowlegeBaseMeta ipBuiltinknowlegeBaseMeta = getKnowlegeBaseMeta(FlowWriteConfig.IP_BUILTIN_KD_ID);
        if (!currentSha256IpBuiltin.equals(ipBuiltinknowlegeBaseMeta.getSha256())) {
            String fileName = Joiner.on(LOCATION_SEPARATOR).useForNull("").join(ipBuiltinknowlegeBaseMeta.getName(), ipBuiltinknowlegeBaseMeta.getFormat());
            knowledgeMetaCache.put(fileName, ipBuiltinknowlegeBaseMeta);
        }
        final KnowlegeBaseMeta ipUserDefinedknowlegeBaseMeta = getKnowlegeBaseMeta(FlowWriteConfig.IP_USER_DEFINED_KD_ID);
        if (!currentSha256IpUserDefined.equals(ipUserDefinedknowlegeBaseMeta.getSha256())) {
            String fileName = Joiner.on(LOCATION_SEPARATOR).useForNull("").join(ipUserDefinedknowlegeBaseMeta.getName(), ipUserDefinedknowlegeBaseMeta.getFormat());
            knowledgeMetaCache.put(fileName, ipUserDefinedknowlegeBaseMeta);
        }

        final KnowlegeBaseMeta asnBuiltinknowlegeBaseMeta = getKnowlegeBaseMeta(FlowWriteConfig.ASN_BUILTIN_KD_ID);
        if (!currentSha256AsnBuiltin.equals(asnBuiltinknowlegeBaseMeta.getSha256())) {
            String fileName = Joiner.on(LOCATION_SEPARATOR).useForNull("").join(asnBuiltinknowlegeBaseMeta.getName(), asnBuiltinknowlegeBaseMeta.getFormat());
            knowledgeMetaCache.put(fileName, asnBuiltinknowlegeBaseMeta);
        }

        if (!currentSha256IpUserDefined.equals(ipUserDefinedknowlegeBaseMeta.getSha256()) || !currentSha256IpBuiltin.equals(ipBuiltinknowlegeBaseMeta.getSha256()) || !currentSha256AsnBuiltin.equals(asnBuiltinknowlegeBaseMeta.getSha256())) {
            currentSha256IpBuiltin = ipBuiltinknowlegeBaseMeta.getSha256();
            currentSha256IpUserDefined = ipUserDefinedknowlegeBaseMeta.getSha256();
            currentSha256AsnBuiltin = asnBuiltinknowlegeBaseMeta.getSha256();
            reloadIpLookup();
            logger.info("知识库加载成功.");
        }
    }

    /**
     * 从HDFS下载文件更新IpLookup
     */
    private static void reloadIpLookup() {
        IpLookupV2.Builder builder = new IpLookupV2.Builder(false);
        for (String fileName : knowledgeMetaCache.keySet()) {
            int retryNum = 0;
            KnowlegeBaseMeta knowlegeBaseMeta = knowledgeMetaCache.get(fileName);
            String metaSha256 = knowlegeBaseMeta.getSha256();
            while (retryNum < TRY_TIMES) {
                System.out.println("download file ：" + fileName + ",HOS path :" + knowlegeBaseMeta.getPath());
                Long startTime = System.currentTimeMillis();
                byte[] httpGetByte = httpClientService.httpGetByte(knowlegeBaseMeta.getPath(), FlowWriteConfig.HTTP_SOCKET_TIMEOUT);
                if (httpGetByte != null && httpGetByte.length > 0) {
                    String downloadFileSha256 = DigestUtil.sha256Hex(httpGetByte);
                    if (metaSha256.equals(downloadFileSha256)) {
                        ByteArrayInputStream inputStream = new ByteArrayInputStream(httpGetByte);
                        switch (fileName) {
                            case ipBuiltInName:
                                builder.loadDataFile(inputStream);
                                break;
                            case ipUserDefinedName:
                                builder.loadDataFilePrivate(inputStream);
                                break;
                            case asnName:
                                builder.loadAsnDataFile(inputStream);
                                break;
                            default:
                        }
                        System.out.println("update " + fileName + " finished, speed :" + (System.currentTimeMillis() - startTime) + "ms");
                        retryNum = TRY_TIMES;
                    } else {
                        logger.error("通过HOS下载{}的sha256为:{} ,网关内记录为:{} ,sha256不相等 开始第{}次重试下载文件", fileName, downloadFileSha256, metaSha256, retryNum);
                        retryNum++;
                    }
                } else {
                    logger.error("通过HOS下载{}的流为空 ,开始第{}次重试下载文件", fileName, retryNum);
                    retryNum++;
                }
            }
        }
        ipLookup = builder.build();
    }

    public static IpLookupV2 getIpLookup() {
        return ipLookup;
    }

    /**
     * 根据配置组合生成知识库元数据过滤参数
     *
     * @return 过滤参数
     */
    private static String getFilterParameter() {

        String expr = "[?(@.version=='latest')][?(@.name in ('ip_builtin','ip_user_defined','asn_builtin'))]";

        return expr;
    }

    public static String getCountryLookup(String ip) {
        return ipLookup.countryLookup(ip);
    }

    private static KnowlegeBaseMeta getKnowlegeBaseMeta(String kd_id) {
        KnowlegeBaseMeta knowlegeBaseMeta = null;
        String knowledgeInfo = null;
        try {
            URIBuilder uriBuilder = new URIBuilder(FlowWriteConfig.KNOWLEDGE_BASE_URL);
            HashMap<String, Object> parms = new HashMap<>();
            parms.put("kb_id", kd_id);
            httpClientService.setUrlWithParams(uriBuilder, FlowWriteConfig.KNOWLEDGE_BASE_PATH, parms);
            knowledgeInfo = httpClientService.httpGet(uriBuilder.build(), FlowWriteConfig.HTTP_SOCKET_TIMEOUT);
            if (knowledgeInfo.contains("200")) {
                final Map<String, Object> jsonObject = JSONObject.parseObject(knowledgeInfo, Map.class);
                JSONPath jsonPath = JSONPath.of(getFilterParameter());
                String extract = jsonPath.extract(JSONReader.of(jsonObject.get("data").toString())).toString();
                if (StringUtil.isNotBlank(extract)) {
                    JSONArray jsonArray = JSON.parseArray(extract);
                    if (jsonArray.size() > 0) {
                        for (int i = 0; i < jsonArray.size(); i++) {
                            knowlegeBaseMeta = JSONObject.parseObject(jsonArray.getString(i), KnowlegeBaseMeta.class);
                        }
                    }
                }
            } else {
                logger.error("获取knowledge_base失败,请求回执为" + knowledgeInfo);
            }
        } catch (URISyntaxException e) {
            logger.error("构造URI异常", e);
        } catch (Exception e) {
            logger.error("获取knowledge_base失败", e);
        }
        return knowlegeBaseMeta;
    }

    public static void main(String[] args) {
        final String countryLookup = IpLookupUtils.getIpLookup().asnLookup("10.64.10.7");
        System.out.println(countryLookup);
    }
}