package com.zdjizhi.function;

import com.zdjizhi.common.DosSketchLog;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.api.java.tuple.Tuple4;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.util.HashMap;
import java.util.Map;

public class MetricsCalculate extends ProcessWindowFunction<
        DosSketchLog,                  // 输入类型
        DosSketchLog,  // 输出类型
        Tuple4<String, String, Integer, Integer>,                         // 键类型
        TimeWindow> {                   // 窗口类型
    private final Map<String, String> attackTypeMapping = new HashMap<>();

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        attackTypeMapping.put("TCP SYN","TCP SYN Flood");
        attackTypeMapping.put("DNS","UDP Flood");
        attackTypeMapping.put("ICMP","ICMP Flood");
        attackTypeMapping.put("UDP","DNS Flood");
        attackTypeMapping.put("NTP","NTP Flood");
        attackTypeMapping.put("","Custom Network Attack");
    }

    @Override
    public void process(Tuple4<String, String, Integer, Integer> key, ProcessWindowFunction<DosSketchLog, DosSketchLog, Tuple4<String, String, Integer,Integer>, TimeWindow>.Context context, Iterable<DosSketchLog> elements, Collector<DosSketchLog> out) throws Exception {

        for (DosSketchLog dosSketchLog: elements){
            dosSketchLog.setSession_rate(dosSketchLog.getSessions()/ (dosSketchLog.getDuration()/1000) );
            dosSketchLog.setPacket_rate(dosSketchLog.getPkts()/(dosSketchLog.getDuration()/1000));
            dosSketchLog.setBit_rate(dosSketchLog.getBytes()/(dosSketchLog.getDuration()/1000));
            dosSketchLog.setAttack_type(attackTypeMapping.getOrDefault(dosSketchLog.getDecoded_as(),""));
            out.collect(dosSketchLog);
        }
    }
}