/**
 * Copyright (C) 2016 Hurence (support@hurence.com)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.hurence.logisland.processor.netflow;

import com.hurence.logisland.annotation.documentation.CapabilityDescription;
import com.hurence.logisland.annotation.documentation.ExtraDetailFile;
import com.hurence.logisland.annotation.documentation.Tags;
import com.hurence.logisland.component.InitializationException;
import com.hurence.logisland.component.PropertyDescriptor;
import com.hurence.logisland.processor.AbstractProcessor;
import com.hurence.logisland.processor.ProcessContext;
import com.hurence.logisland.record.*;
import com.hurence.logisland.validator.StandardValidators;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.*;

/**
 * Netflow (http://www.cisco.com/c/en/us/td/docs/ios/solutions_docs/netflow/nfwhite.html) processor
 */
@Tags({"netflow", "security"})
@CapabilityDescription(
        "The `Netflow V5 <http://www.cisco.com/c/en/us/td/docs/ios/solutions_docs/netflow/nfwhite.html>`_ processor "
                + "is the Logisland entry point to  process Netflow (V5) events. "
                + "NetFlow is a feature introduced on Cisco routers that provides the ability to collect IP network "
                + "traffic.We can distinguish 2 components:\n\n"
                + "\t- Flow exporter: aggregates packets into flows and exports flow records (binary format) towards one "
                + "or more flow collectors\n\n"
                + "\t- Flow collector: responsible for reception, storage and pre-processing of flow data received from "
                + "a flow exporter\n\n"
                + "The collected data are then available for analysis purpose (intrusion detection, traffic analysis...)\n"
                + "Netflow are sent to kafka in order to be processed by logisland.\n"
                + "In the tutorial we will simulate Netflow traffic using `nfgen <https://github.com/pazdera/NetFlow-Exporter-Simulator>`_. "
                + "this traffic will be sent to port 2055. The we rely on nifi to listen of that port for  "
                + " incoming netflow (V5) traffic and send them to a kafka topic. "
                + "The Netflow processor could thus treat these events and generate corresponding logisland records. The following processors "
                + "in the stream can then process the Netflow records generated by this processor.")
@ExtraDetailFile("./details/ParseNetflowEvent-Detail.rst")

public class ParseNetflowEvent extends AbstractProcessor {

    private static Logger logger = LoggerFactory.getLogger(ParseNetflowEvent.class);

    final int FIN = 1;
    final int SYN = 2;
    final int RST = 4;
    final int PSH = 8;
    final int ACK = 16;
    final int URG = 32;

    private boolean debug = false;
    
    private static final String KEY_DEBUG = "debug";
    
    public static final PropertyDescriptor DEBUG = new PropertyDescriptor.Builder()
            .name(KEY_DEBUG)
            .description("Enable debug. If enabled, the original JSON string is embedded in the record_value field of the record.")
            .addValidator(StandardValidators.BOOLEAN_VALIDATOR)
            .required(false)
            .defaultValue("false")
            .build();

    public static final PropertyDescriptor OUTPUT_RECORD_TYPE = new PropertyDescriptor.Builder()
            .name("output.record.type")
            .description("the output type of the record")
            .required(false)
            .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
            .defaultValue("netflowevent")
            .build();

    public static final PropertyDescriptor ENRICH_RECORD = new PropertyDescriptor.Builder()
            .name("enrich.record")
            .description("Enrich data. If enabledthe netflow record is enriched with inferred data")
            .required(false)
            .addValidator(StandardValidators.BOOLEAN_VALIDATOR)
            .defaultValue("false")
            .build();

    @Override
    public void init(final ProcessContext context) throws InitializationException
    {
        super.init(context);
        logger.debug("Initializing Netflow Processor");
    }
    
    @Override
    public List<PropertyDescriptor> getSupportedPropertyDescriptors() {
        
        final List<PropertyDescriptor> descriptors = new ArrayList<>();
        descriptors.add(DEBUG);
        descriptors.add(OUTPUT_RECORD_TYPE);
        descriptors.add(ENRICH_RECORD);
        return Collections.unmodifiableList(descriptors);
    }
  
    @Override
    public Collection<Record> process(ProcessContext context, Collection<Record> records)
    {
        Collection<Record> list = new ArrayList<>();
        final String outputRecordType = context.getPropertyValue(OUTPUT_RECORD_TYPE).asString();
        final boolean enrichRecord = context.getPropertyValue(ENRICH_RECORD).asBoolean();
        if (debug)
        {
            logger.debug("Netflow Processor records input: " + records);
        }

        /**
         * Get  Netflow event as a JSON string:
         */
        records.forEach(record -> {
            byte[] recordValue=null;
            recordValue = (byte[]) record.getField(FieldDictionary.RECORD_VALUE).getRawValue();
            if (debug) {
                logger.debug("record=" + Arrays.toString(recordValue));
            }


            Map header = getHeader(Arrays.copyOfRange(recordValue,0,24));

            for (int i=0; i < (int) header.get("count"); i++) {
                int offset = i*48;

                Map netflowRecord = getRecord(Arrays.copyOfRange(recordValue,24+offset,72+offset));

                Record evt = new StandardRecord(outputRecordType)
                        .setTime(new Date(((long) header.get("unix_secs"))*1000))
                        .setField(new Field("src_ip4", FieldType.STRING, netflowRecord.get("src_ip4")))
                        .setField(new Field("dst_ip4", FieldType.STRING, netflowRecord.get("dst_ip4")))
                        .setField(new Field("nexthop", FieldType.STRING, netflowRecord.get("nexthop")))
                        .setField(new Field("input", FieldType.INT, netflowRecord.get("input")))
                        .setField(new Field("output", FieldType.INT, netflowRecord.get("output")))
                        .setField(new Field("dPkts", FieldType.LONG, netflowRecord.get("dPkts")))
                        .setField(new Field("dOctets", FieldType.LONG, netflowRecord.get("dOctets")))
                        .setField(new Field("first", FieldType.LONG, netflowRecord.get("first")))
                        .setField(new Field("last", FieldType.LONG, netflowRecord.get("last")))
                        .setField(new Field("src_port", FieldType.INT, netflowRecord.get("src_port")))
                        .setField(new Field("dst_port", FieldType.INT, netflowRecord.get("dst_port")))
                        .setField(new Field("flags", FieldType.INT, netflowRecord.get("flags")))
                        .setField(new Field("nprot", FieldType.INT, netflowRecord.get("nprot")))
                        .setField(new Field("tos", FieldType.INT, netflowRecord.get("tos")))
                        .setField(new Field("src_as", FieldType.INT, netflowRecord.get("src_as")))
                        .setField(new Field("dst_as", FieldType.INT, netflowRecord.get("dst_as")))
                        .setField(new Field("src_mask", FieldType.INT, netflowRecord.get("src_mask")))
                        .setField(new Field("dst_mask", FieldType.INT, netflowRecord.get("dst_mask")));

                // enrich record if required
                if (enrichRecord) {
                    long duration = ((long) netflowRecord.get("last")) - ((long) netflowRecord.get("first"));
                    evt.setField(new Field("duration", FieldType.LONG, duration));
                    try{
                        String ipString = (String) netflowRecord.get("src_ip4");
                        InetAddress ia = InetAddress.getByName(ipString);
                        String host = ia.getCanonicalHostName();
                        if (host.compareTo(ipString) != 0) {
                            evt.setField(new Field("src_host", FieldType.STRING, host));

                            int index = host.indexOf('.');
                            String domain = (index >= 0) ? host.substring(index) : null;
                            if (domain != null)
                                evt.setField(new Field("src_domain", FieldType.STRING, domain));
                        }

                        ipString = (String) netflowRecord.get("dst_ip4");
                        ia = InetAddress.getByName(ipString);
                        host = ia.getCanonicalHostName();
                        if (host.compareTo(ipString) != 0) {
                            evt.setField(new Field("dst_host", FieldType.STRING, host));

                            int index = host.indexOf('.');
                            String domain = (index >= 0) ? host.substring(index) : null;
                            if (domain != null)
                                evt.setField(new Field("dst_domain", FieldType.STRING, domain));
                        }

                        // flags
                        short tcpflags = (short) netflowRecord.get("flags");
                        evt.setField(new Field("tcp_fin", FieldType.BOOLEAN, (tcpflags & FIN) == FIN));
                        evt.setField(new Field("tcp_syn", FieldType.BOOLEAN, (tcpflags & SYN) == SYN));
                        evt.setField(new Field("tcp_rst", FieldType.BOOLEAN, (tcpflags & RST) == RST));
                        evt.setField(new Field("tcp_psh", FieldType.BOOLEAN, (tcpflags & PSH) == PSH));
                        evt.setField(new Field("tcp_ack", FieldType.BOOLEAN, (tcpflags & ACK) == ACK));
                        evt.setField(new Field("tcp_urg", FieldType.BOOLEAN, (tcpflags & URG) == URG));

                    } catch (UnknownHostException e) {
                        logger.error("Bad host address");
                    }
                }

                list.add(evt);
            }


            // Parse Netflow event as JSON object
            // Map<String, Object> jsonNetflowEvent = JsonUtil.convertJsonToMap(recordValue);

        });
            if (debug) {
                logger.debug("Netflow Processor records output: " + records);
            }

        return list;
    }

    
    @Override
    public void onPropertyModified(PropertyDescriptor descriptor, String oldValue, String newValue) {

        logger.debug("property {} value changed from {} to {}", descriptor.getName(), oldValue, newValue);
        
        /**
         * Handle the debug property
         */
        if (descriptor.getName().equals(KEY_DEBUG))
        {
          if (newValue != null)
          {
              if (newValue.equalsIgnoreCase("true"))
              {
                  debug = true;
              }
          } else
          {
              debug = false;
          }
        }
    }

    private Map getHeader(byte[] header) {
        String[] header_label = {"version","count","sysuptime","unix_secs",
                "unix_nsecs","flow_sequence","engine_type","engine_id","sampling_interval"};
        String[] header_type = {"short","short","int","int","int","int","byte","byte","short"};
        Map<String, Object> header_value = new HashMap<>();

        byte shift = 0;
        byte[] ID = null;

        for (byte i=0; i < header_label.length; i++) {
            switch(header_type[i]) {
                case "byte":
                    ID = Arrays.copyOfRange(header, shift, shift+1);
                    header_value.put(header_label[i], getByte(ID));
                    shift += 1;
                    break;

                case "short":
                    ID = Arrays.copyOfRange(header, shift, shift+2);
                    header_value.put(header_label[i], getShort(ID));
                    shift += 2;
                    break;

                case "int":
                    ID = Arrays.copyOfRange(header, shift, shift+4);
                    header_value.put(header_label[i], getInt(ID));
                    shift += 4;
                    break;
            }
        }

        if (debug) {
            for (String label : header_label) {
                logger.debug(label + "=" + header_value.get(label));
            }
        }

        return (header_value);
    }

    private Map getRecord(byte[] record) {

        String[] record_label = {"src_ip4", "dst_ip4", "nexthop", "input", "output", "dPkts", "dOctets", "first",
                "last", "src_port", "dst_port", "pad1", "flags", "nprot", "tos", "src_as", "dst_as", "src_mask",
                "dst_mask", "pad2"};
        String[] record_type = {"ip","ip","ip","short","short","int","int","int","int","short","short","byte","byte",
                "byte","byte","short","short","byte","byte","short"};

        Map<String, Object> record_value = new HashMap<>();

        byte shift = 0;
        byte[] ID = null;

        for (byte i=0; i < record_label.length; i++) {
            switch(record_type[i]) {
                case "byte":
                    ID = Arrays.copyOfRange(record, shift, shift+1);
                    record_value.put(record_label[i], getByte(ID));
                    shift += 1;
                    break;

                case "short":
                    ID = Arrays.copyOfRange(record, shift, shift+2);
                    record_value.put(record_label[i], getShort(ID));
                    shift += 2;
                    break;

                case "int":
                    ID = Arrays.copyOfRange(record, shift, shift+4);
                    record_value.put(record_label[i], getInt(ID));
                    shift += 4;
                    break;

                case "ip":
                    ID = Arrays.copyOfRange(record, shift, shift+4);
                    record_value.put(record_label[i], getIP(ID));
                    shift += 4;
                    break;
            }
        }

        if (debug) {
            for (String label : record_label) {
                logger.debug(label + "=" + record_value.get(label));
            }
        }

        return (record_value);

    }

    private short getByte(byte[] buffer) {
        return((short) (buffer[0] & 0xFF));
    }

    private  int getShort(byte[] buffer) {
        return((int) (((buffer[0] & 0xFF) << 8) | (buffer[1] & 0xFF)));
    }

    private  long getInt(byte[] buffer) {
        return(  (long) (((buffer[0] & 0xFF) << 24) | ((buffer[1] & 0xFF) << 16)
                | ((buffer[2] & 0xFF) << 8) | (buffer[3] & 0xFF)));
    }

    private String getIP(byte[] buffer) {

        try {
            for (byte i = 1; i < 4; i++) {
                buffer[i] = (byte) (buffer[i] & 0xFF);
            }
            InetAddress address = InetAddress.getByAddress(buffer);
            return (address.toString().replace("/",""));
        } catch ( UnknownHostException e) {
            logger.error("Bad host address");
        }
        return("");
    }
}
