galaxy-tsg-olap-dos-detection-job/src/main/java/com/zdjizhi/function/DosDetectionFunction.java

package com.zdjizhi.function;

import cn.hutool.log.Log;
import cn.hutool.log.LogFactory;
import com.geedgenetworks.utils.DateUtils;
import com.geedgenetworks.utils.StringUtil;
import com.zdjizhi.common.*;
import com.zdjizhi.utils.Snowflakeld.SnowflakeId;
import com.zdjizhi.utils.Threshold.ParseBaselineThreshold;
import com.zdjizhi.utils.Threshold.ParseStaticThreshold;
import com.zdjizhi.utils.connections.http.HttpClientService;
import com.zdjizhi.utils.knowledgebase.IpLookupUtils;
import inet.ipaddr.IPAddress;
import inet.ipaddr.IPAddressString;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.shaded.guava18.com.google.common.collect.TreeRangeMap;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;

import java.math.BigDecimal;
import java.text.NumberFormat;
import java.util.*;

import static com.zdjizhi.conf.DosConfigs.*;

/**
 * @author wlh
 */
public class DosDetectionFunction extends ProcessFunction<DosSketchLog, DosEventLog> {
    private static final Log logger = LogFactory.get();
    private  Map<String, Map<String, DosBaselineThreshold>> baselineMap = new HashMap<>();
    private final  NumberFormat PERCENT_INSTANCE = NumberFormat.getPercentInstance();
    private HashMap<Integer, HashMap<String, TreeRangeMap<IPAddress, DosDetectionThreshold>>> thresholdRangeMap;
    private final int BASELINE_SIZE = 144;
    private final int STATIC_CONDITION_TYPE = 1;
    private final int BASELINE_CONDITION_TYPE = 2;
    private final int SENSITIVITY_CONDITION_TYPE = 3;
    private final String SESSIONS_TAG = "sessions";
    private final String PACKETS_TAG = "packets";
    private final String BITS_TAG = "bits";
    private final int OTHER_BASELINE_TYPE = 3;
    private SnowflakeId snowflakeId;
    private Configuration configuration;
    private HttpClientService httpClientService;
    private IpLookupUtils ipLookupUtils;
    private ParseBaselineThreshold parseBaselineThresholdld;
    private ParseStaticThreshold parseStaticThreshold;

    @Override
    public void open(Configuration parameters) {

        configuration = (Configuration) getRuntimeContext()
                .getExecutionConfig().getGlobalJobParameters();
        httpClientService = new HttpClientService(configuration);

        snowflakeId = new SnowflakeId(configuration.get(DATA_CENTER_ID_NUM), getRuntimeContext().getIndexOfThisSubtask());

        try {
            ipLookupUtils = new IpLookupUtils(configuration, httpClientService);
            ipLookupUtils.stuffKnowledgeMetaCache();
            Timer timer = new Timer();
            timer.schedule(new TimerTask() {
                @Override
                public void run() {
                    ipLookupUtils.stuffKnowledgeMetaCache();
                    logger.info("定位库定时调度成功");
                }
            }, configuration.get(KNOWLEDGE_BASE_SCHEDULE_MINUTES) * 60 * 1000, configuration.get(KNOWLEDGE_BASE_SCHEDULE_MINUTES) * 60 * 1000);
        } catch (Exception e) {
            logger.error("定位库加载失败,具体原因为" + e);
        }

        try {
            parseStaticThreshold = new ParseStaticThreshold(configuration, httpClientService);
            thresholdRangeMap = parseStaticThreshold.createStaticThreshold();
            Timer timer = new Timer();
            timer.schedule(new TimerTask() {
                @Override
                public void run() {
                    thresholdRangeMap = parseStaticThreshold.createStaticThreshold();
                    logger.info("基于静态阈值构建threshold RangeMap成功,Threshold RangeMap:" + thresholdRangeMap.toString());
                }
            }, configuration.get(STATIC_THRESHOLD_SCHEDULE_MINUTES) * 60 * 1000, configuration.get(STATIC_THRESHOLD_SCHEDULE_MINUTES) * 60 * 1000);

        } catch (Exception e) {
            logger.error("基于静态阈值构建threshold RangeMap失败,失败原因为:" + e);
        }

        try {
            parseBaselineThresholdld = new ParseBaselineThreshold(configuration);
            baselineMap = parseBaselineThresholdld.readFromHbase();
            Timer timer = new Timer();
            timer.schedule(new TimerTask() {
                @Override
                public void run() {
                    baselineMap = parseBaselineThresholdld.readFromHbase();
                    logger.info("从Hbase获取baselineMap成功,baselineMap:" + thresholdRangeMap.toString());
                }
            }, configuration.get(BASELINE_THRESHOLD_SCHEDULE_DAYS) * 24 * 60 * 60 * 1000, configuration.get(BASELINE_THRESHOLD_SCHEDULE_DAYS) * 24 * 60 * 60 * 1000);

        } catch (Exception e) {
            logger.error("从Hbase获取baselineMap失败,失败原因为:" + e);
        }

        PERCENT_INSTANCE.setMinimumFractionDigits(2);
    }

    @Override
    public void processElement(DosSketchLog value, Context ctx, Collector<DosEventLog> out) throws Exception {
        DosEventLog finalResult = null;
        try {
            String destinationIp = value.getDestination_ip();
            int vsysId = value.getVsys_id();
            String key = destinationIp + "-" + vsysId;
            String attackType = value.getAttack_type();
            IPAddress destinationIpAddress = new IPAddressString(destinationIp).getAddress();

            DosDetectionThreshold threshold = null;
            if (thresholdRangeMap.containsKey(vsysId)) {
                threshold = thresholdRangeMap.get(vsysId).getOrDefault(attackType, TreeRangeMap.create()).get(destinationIpAddress);
            }

            logger.debug("当前判断IP：{}, 类型: {}", key, attackType);
            if (threshold == null && baselineMap.containsKey(key)) {
                finalResult = getDosEventLogByBaseline(value, key);
            } else if (threshold == null && !baselineMap.containsKey(key)) {
                finalResult = getDosEventLogBySensitivityThreshold(value);
            } else if (threshold != null) {
                finalResult = getDosEventLogByStaticThreshold(value, threshold);
            } else {
                logger.debug("未获取到当前server IP：{} 类型 {} 静态阈值 和 baseline", key, attackType);
            }

        } catch (Exception e) {
            logger.error("判定失败\n {} \n{}", value, e);
        }
        if (finalResult != null) {
            out.collect(finalResult);
        }
    }


    private DosEventLog getDosEventLogBySensitivityThreshold(DosSketchLog value) {
        long sketchSessions = value.getSketch_sessions();
        Integer staticSensitivityThreshold = configuration.get(STATIC_SENSITIVITY_THRESHOLD);
        long diff = sketchSessions - staticSensitivityThreshold;
        return getDosEventLog(value, staticSensitivityThreshold, diff, 0, SENSITIVITY_CONDITION_TYPE, SESSIONS_TAG);
    }

    private DosEventLog getDosEventLogByBaseline(DosSketchLog value, String key) {
        String attackType = value.getAttack_type();
        long sketchSessions = value.getSketch_sessions();
        DosBaselineThreshold dosBaselineThreshold = baselineMap.get(key).get(attackType);
        Integer base = getBaseValue(dosBaselineThreshold, value);
        long diff = sketchSessions - base;
        return getDosEventLog(value, base, diff, 0, BASELINE_CONDITION_TYPE, SESSIONS_TAG);
    }

    private DosEventLog getDosEventLogByStaticThreshold(DosSketchLog value, DosDetectionThreshold threshold) throws CloneNotSupportedException {
        long sessionBase = threshold.getSessions_per_sec();
        long pktBase = threshold.getPackets_per_sec();
        long bitBase = threshold.getBits_per_sec();

        long diffSession = value.getSketch_sessions() - sessionBase;
        long diffPkt = value.getSketch_packets() - pktBase;
        long diffByte = value.getSketch_bytes() - bitBase;

        double diffSessionPercent = 0.0;
        double diffPktPercent = 0.0;
        double diffBitPercent = 0.0;

        if (sessionBase > 0) {
            diffSessionPercent = getDiffPercent(diffSession, sessionBase) * 100;
        }
        if (pktBase > 0) {
            diffPktPercent = getDiffPercent(diffPkt, pktBase) * 100;
        }
        if (bitBase > 0) {
            diffBitPercent = getDiffPercent(diffByte, bitBase) * 100;
        }

        long profileId = 0;
        DosEventLog result = null;

        if (diffSessionPercent >= diffPktPercent && diffSessionPercent >= diffBitPercent) {
            profileId = threshold.getId();
            result = getDosEventLog(value, sessionBase, diffSession, profileId, STATIC_CONDITION_TYPE, SESSIONS_TAG);
        } else if (diffPktPercent >= diffSessionPercent && diffPktPercent >= diffBitPercent) {
            profileId = threshold.getId();
            result = getDosEventLog(value, pktBase, diffPkt, profileId, STATIC_CONDITION_TYPE, PACKETS_TAG);
        } else if (diffBitPercent >= diffPktPercent && diffBitPercent >= diffSessionPercent) {
            profileId = threshold.getId();
            result = getDosEventLog(value, bitBase, diffByte, profileId, STATIC_CONDITION_TYPE, BITS_TAG);
        }
        return result;
    }

    private DosEventLog getDosEventLog(DosSketchLog value, long base, long diff, long profileId, int type, String tag) {
        DosEventLog result = null;
        String destinationIp = value.getDestination_ip();
        String attackType = value.getAttack_type();
        if (diff > 0 && base != 0) {
            double percent = getDiffPercent(diff, base);
            Severity severity = judgeSeverity(percent);
            Integer staticSensitivityThreshold = configuration.get(STATIC_SENSITIVITY_THRESHOLD);
            if (severity != Severity.NORMAL) {
                if (type == BASELINE_CONDITION_TYPE && percent < configuration.get(BASELINE_SENSITIVITY_THRESHOLD)) {
                    logger.debug("当前server IP:{},类型:{},基线值{}百分比{}未超过基线敏感阈值,日志详情\n{}", destinationIp, attackType, base, percent, value);
                } else if ((type == BASELINE_CONDITION_TYPE || type == SENSITIVITY_CONDITION_TYPE) && value.getSketch_sessions() < staticSensitivityThreshold) {
                    logger.debug("当前server IP:{},类型:{},基线值{}百分比{}未超过静态敏感阈值,日志详情\n{}", destinationIp, attackType, base, percent, value);
                } else {
                    result = getResult(value, base, profileId, severity, percent, type, tag);
                    if (type == SENSITIVITY_CONDITION_TYPE) {
                        result.setSeverity(Severity.MAJOR.severity);
                    }
                    logger.info("检测到当前server IP {} 存在 {} 异常,超出基线{} {}倍,基于{}:{}检测,日志详情\n {}", destinationIp, attackType, base, percent, type, tag, result);
                }
            } else {
                logger.debug("当前server IP:{} 未出现 {} 异常,日志详情 {}", destinationIp, attackType, value);
            }
        }
        return result;
    }

    private DosEventLog getResult(DosSketchLog value, long base, long profileId, Severity severity, double percent, int type, String tag) {
        DosEventLog dosEventLog = new DosEventLog();
        dosEventLog.setRecv_time(value.getCommon_recv_time());
        dosEventLog.setLog_id(snowflakeId.nextId());
        dosEventLog.setVsys_id(value.getVsys_id());
        dosEventLog.setStart_time(value.getSketch_start_time());
        dosEventLog.setEnd_time(value.getSketch_start_time() + value.getSketch_duration());
        dosEventLog.setProfile_id(profileId);
        dosEventLog.setAttack_type(value.getAttack_type());
        dosEventLog.setSeverity(severity.severity);
        dosEventLog.setConditions(getConditions(PERCENT_INSTANCE.format(percent), base, value.getSketch_sessions(), type, tag, dosEventLog));
        dosEventLog.setDestination_ip(value.getDestination_ip());
        dosEventLog.setDestination_country(ipLookupUtils.getCountryLookup(value.getDestination_ip()));
        String ipList = value.getSource_ip();
        dosEventLog.setSource_ip_list(ipList);
        dosEventLog.setSource_country_list(getSourceCountryList(ipList));
        dosEventLog.setSession_rate(value.getSketch_sessions());
        dosEventLog.setPacket_rate(value.getSketch_packets());
        dosEventLog.setBit_rate(value.getSketch_bytes());
        return dosEventLog;
    }

    private Integer getBaseValue(DosBaselineThreshold dosBaselineThreshold, DosSketchLog value) {
        Integer base = 0;
        try {
            if (dosBaselineThreshold != null) {
                ArrayList<Integer> baselines = dosBaselineThreshold.getSession_rate();
                Integer defaultVaule = dosBaselineThreshold.getSession_rate_default_value();
                Integer sessionRateBaselineType = dosBaselineThreshold.getSession_rate_baseline_type();
                if (baselines != null && baselines.size() == BASELINE_SIZE) {
                    int timeIndex = getCurrentTimeIndex(value.getSketch_start_time());
                    base = baselines.get(timeIndex);
                    if (base == 0) {
                        logger.debug("获取到当前IP: {},类型: {} baseline值为0,替换为P95观测值{}", value.getDestination_ip(), value.getAttack_type(), defaultVaule);
                        base = defaultVaule;
                    }
                    if (sessionRateBaselineType == OTHER_BASELINE_TYPE && base < configuration.get(STATIC_SENSITIVITY_THRESHOLD)) {
                        base = configuration.get(STATIC_SENSITIVITY_THRESHOLD);
                    }
                }
            }
        } catch (Exception e) {
            logger.error("解析baseline数据失败,返回默认值0", e);
        }
        return base;
    }

    private String getConditions(String percent, long base, long sessions, int type, String tag, DosEventLog dosEventLog) {
        int condition = 0;
        if ("Minor".equals(dosEventLog.getSeverity())) {
            condition = 50;
        } else if ("Warning".equals(dosEventLog.getSeverity())) {
            condition = 100;
        } else if ("Major".equals(dosEventLog.getSeverity())) {
            condition = 250;
        } else if ("Severe".equals(dosEventLog.getSeverity())) {
            condition = 500;
        } else if ("Critical".equals(dosEventLog.getSeverity())) {
            condition = 800;
        }
        switch (type) {
            case STATIC_CONDITION_TYPE:
                return "Rate > " +
                        base + " " +
                        tag + "/s" + "(>" + condition + "%)";
            case BASELINE_CONDITION_TYPE:
                return tag + " > " +
                        percent + " of baseline";
            case SENSITIVITY_CONDITION_TYPE:
                return String.valueOf(sessions) + " " +
                        tag + "/s Unusually high " +
                        StringUtils.capitalize(tag);
            default:
                throw new IllegalArgumentException("Illegal Argument type:" + type + ", known types = [1,2,3]");
        }
    }

    private String getSourceCountryList(String sourceIpList) {
        if (StringUtil.isNotBlank(sourceIpList)) {
            String countryList;
            try {
                String[] ipArr = sourceIpList.split(",");
                HashSet<String> countrySet = new HashSet<>();
                for (String ip : ipArr) {
                    String country = ipLookupUtils.getCountryLookup(ip);
                    if (StringUtil.isNotBlank(country)) {
                        countrySet.add(country);
                    }
                }
                countryList = StringUtils.join(countrySet, ", ");
                return countryList;
            } catch (Exception e) {
                logger.error("{} source IP lists 获取国家失败", sourceIpList, e);
                return StringUtil.EMPTY;
            }
        } else {
            throw new IllegalArgumentException("Illegal Argument sourceIpList = null");
        }
    }

    private int getCurrentTimeIndex(long sketchStartTime) {
        int index = 0;
        try {
            long currentDayTime = DateUtils.getTimeFloor(new Date(sketchStartTime * 1000L), "P1D").getTime() / 1000;
            long indexLong = (sketchStartTime - currentDayTime) / (86400 / BASELINE_SIZE);
            index = Integer.parseInt(Long.toString(indexLong));
        } catch (Exception e) {
            logger.error("获取time index失败", e);
        }
        return index;
    }


    private Double getDiffPercent(long diff, long base) {
        try {
            return BigDecimal.valueOf((float) diff / base).setScale(4, BigDecimal.ROUND_HALF_UP).doubleValue();
        } catch (Exception e) {
            logger.info("当前阈值为0,进行下一阈值条件判断", e);
            return 0.0;
        }

    }

    private Severity judgeSeverity(double diffPercent) {
        if (diffPercent >= configuration.get(BASELINE_SESSIONS_MINOR_THRESHOLD) && diffPercent < configuration.get(BASELINE_SESSIONS_WARNING_THRESHOLD)) {
            return Severity.MINOR;
        } else if (diffPercent >= configuration.get(BASELINE_SESSIONS_WARNING_THRESHOLD) && diffPercent < configuration.get(BASELINE_SESSIONS_MAJOR_THRESHOLD)) {
            return Severity.WARNING;
        } else if (diffPercent >= configuration.get(BASELINE_SESSIONS_MAJOR_THRESHOLD) && diffPercent < configuration.get(BASELINE_SESSIONS_SEVERE_THRESHOLD)) {
            return Severity.MAJOR;
        } else if (diffPercent >= configuration.get(BASELINE_SESSIONS_SEVERE_THRESHOLD) && diffPercent < configuration.get(BASELINE_SESSIONS_CRITICAL_THRESHOLD)) {
            return Severity.SEVERE;
        } else if (diffPercent >= configuration.get(BASELINE_SESSIONS_CRITICAL_THRESHOLD)) {
            return Severity.CRITICAL;
        } else {
            return Severity.NORMAL;
        }
    }

    private enum Severity {
        /**
         * 判断严重程度枚举类型
         */
        CRITICAL("Critical"),
        SEVERE("Severe"),
        MAJOR("Major"),
        WARNING("Warning"),
        MINOR("Minor"),
        NORMAL("Normal");

        private final String severity;

        @Override
        public String toString() {
            return this.severity;
        }

        Severity(String severity) {
            this.severity = severity;
        }

    }

}