package eu.gity.jmeter.ddos.trafgencustom;

import eu.gity.jmeter.ddos.common.DDoSUtils;
import eu.gity.jmeter.ddos.common.Trafgen;
import eu.gity.jmeter.ddos.common.TrafgenControl;
import eu.gity.jmeter.ddos.common.TrafgenControl.TrafgenThread;
import org.apache.jmeter.samplers.AbstractSampler;
import org.apache.jmeter.samplers.Entry;
import org.apache.jmeter.samplers.SampleResult;
import org.apache.jmeter.testelement.TestElement;
import org.apache.jmeter.testelement.TestStateListener;
import org.apache.jmeter.testelement.property.PropertyIterator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.concurrent.TimeUnit;

/**
 * This Test Element class stores properties of Trafgen Custom Config DoS attack starts trafgen generator using instance of TrafgenControl class
 *
 * @author Gity a.s.
 */
public class TrafgenCustomSampler extends AbstractSampler implements TestStateListener, Trafgen {
    public static final String NUMBER = "TrafgenCustomSampler.number";
    public static final String RATE = "TrafgenCustomSampler.rate";
    public static final String INTERF = "TrafgenCustomSampler.interface";
    public static final String CONFIG = "TrafgenCustomSampler.config";
    //This property is set from DDoSThreadGroup and is just forwarded to TrafgenControl
    public static final String DYNAMICRATE = "TrafgenCustomSampler.dynamicRate";
    public static final String configFileSuffix = "TrafgenCustom.cfg";
    private static final long serialVersionUID = 240L;
    private static final Logger log = LoggerFactory.getLogger(TrafgenCustomSampler.class);
    private transient TrafgenControl trafgenControl;

    /**
     * Constructs trafgenControl object if not done yet
     */
    void checkInit() {
        if (trafgenControl == null) {
            trafgenControl = new TrafgenControl();
        }
    }

    /**
     * Syncing of the information to GUI do disable/enable specific options
     *
     * @return true if dynamicRate will be used
     */
    boolean isDynamicRate() {
        return trafgenControl.isDynamicRate();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void setDynamicRate(boolean dynamic) {
        this.setProperty(DYNAMICRATE, dynamic);
        trafgenControl.setDynamicRate(dynamic);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public SampleResult sample(Entry e) {
        if (this.getPropertyAsString(INTERF).isEmpty()) {
            return null; //This error is already reported to log in DDoSUtils class
        }

        SampleResult res = new SampleResult();
        res.setSampleLabel(getName());
        res.sampleStart();
        res.setDataType(SampleResult.TEXT);
        int result = trafgenControl.runTrafgen();
        switch (result) {
            case -3:    //Trafgen has finished (static rate)
                log.debug("Trafgen has finished generating predefined amount of packets. Stopping JMeter thread.");
                res.setResponseOK();
                this.getThreadContext().getThread().stop();
                break;
            case -2:    //Rate is constant;
                res.setResponseOK();
                break;
            case -1:
                log.error("There was problem with Trafgen configuration file - see details above.");
                res.setResponseMessage("There was problem with Trafgen configuration file - see log for details.");
                break;
            case 0:
                log.debug("Rate set to 0 - terminating generator");
                res.setResponseOK();
                break;
            default:
                log.debug("Packet generation rate changed to " + result);
                res.setResponseOK();
        }
        res.sampleEnd();
        try {
            //Prevents generation of excessive amounts of samples
            TimeUnit.MILLISECONDS.sleep(50);
        } catch (InterruptedException e1) {
            //Never mind
        }
        return res;
    }

    /**
     * Generates configuration file based on text from CONFIG Test Element property
     *
     * @return name of the newly created configuration file or null on failure
     */
    private String generateConfig() {
        String jmeterHomeDir = DDoSUtils.getInstance().getJmeterHomeDir();

        DateFormat dateFormat = new SimpleDateFormat("yyyy_MM_dd_HH_mm_ss");
        Date date = new Date();

        String nanos = Long.toString(System.nanoTime());
        String timestamp = dateFormat.format(date) + "_" + nanos.substring(nanos.length() - 6);

        String configFile = jmeterHomeDir + TrafgenControl.configRelativePath + timestamp + configFileSuffix;
        trafgenControl.setTimestamp(timestamp);

        try (Writer writer = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(configFile), StandardCharsets.UTF_8))) {
            writer.write(this.getPropertyAsString(CONFIG));
        } catch (IOException e) {
            log.error("Error while writting configuration to file " + configFile + ":\n" + e.getStackTrace().toString());
            return null;
        }
        return configFile;
    }

    /**
     * Starts trafgen process in a separate thread or restarts it if dynamic rate setting is being used and the setting has changed since last run
     *
     * @return -2 when no rate change was required, -1 on error, 0 when rate is set to 0 (trafgen is terminated) or any other number which corresponds to a new rate
     */
    public int runTrafgen() {
        if (trafgenControl.getConfigFile() == null) {
            trafgenControl.setConfigFile(generateConfig());
            if (trafgenControl.getConfigFile() == null) {
                log.error("Error occured while creating Trafgen config file. Aborting!");
                return -1;
            }
        }

        String netInterface = this.getPropertyAsString(INTERF);
        String number = null;
        String rate = null;

        if (!trafgenControl.isDynamicRate()) {
            String correctedPacketCount = trafgenControl.getCorrectedPacketCount();
            String packetCount = correctedPacketCount == null ? this.getPropertyAsString(NUMBER) : correctedPacketCount;
            if (packetCount.isEmpty()) {
                number = "";
            } else {
                number = "--num " + packetCount;
            }

            String packetRate = this.getPropertyAsString(RATE);
            if (packetRate.isEmpty()) {
                rate = "";
            } else {
                rate = " --rate " + packetRate + "pps";
            }
        } else {
            if (trafgenControl.getNewRate() != trafgenControl.getOldRate()) {
                int newRate = trafgenControl.getNewRate();
                trafgenControl.setOldRate(newRate);
                number = "";
                rate = " --rate " + newRate + "pps";
                trafgenControl.terminateTrafgen(false);
                if (newRate == 0) {
                    try {
                        Files.delete(Paths.get(trafgenControl.getConfigFile()));
                    } catch (IOException e1) {
                        log.error("Error occured while deleting trafgen configuration file " + trafgenControl.getConfigFile() + ":\n" + e1.getStackTrace().toString());
                    }
                    trafgenControl.setConfigFile(null);
                    return 0;
                }
            } else {
                return -2; //Rate has not changed so trafgen doesn't have to be restarted
            }
        }

        String command = "pkexec trafgen --in " + trafgenControl.getConfigFile() + " --out "
                + netInterface + " --cpp " + number + rate;

        String[] args = new String[]{"/bin/bash", "-c", command};

        trafgenControl.setTrafgenThread(trafgenControl.new TrafgenThread(args));
        trafgenControl.getTrafgenThread().start();

        return trafgenControl.getNewRate();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void stopTrafgen() {
        trafgenControl.terminateTrafgen(false);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void checkPacketCount(String tgName, long duration) {
        trafgenControl.verifyPacketCount(tgName, duration, this.getPropertyAsLong(RATE), this.getPropertyAsLong(NUMBER));
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void changeRate(int rateChange) {
        trafgenControl.changeRate(rateChange);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void testStarted() {
        checkInit();
        trafgenControl.generateConfig();
        trafgenControl.setDynamicRate(this.getPropertyAsBoolean(DYNAMICRATE));
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void testStarted(String host) {
        testStarted();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void testEnded() {
        TrafgenThread tt = trafgenControl.getTrafgenThread();
        if (tt == null) {
            return;
        }
        try {
            tt.join();
        } catch (InterruptedException e) {
            log.error("The sampler was interrupted while waiting for trafgen thread to finish:\n" + e.getStackTrace().toString());
        }
        if (!tt.getResult()) {
            log.error("A problem occured while generating traffic");
        }
        trafgenControl.reset();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void testEnded(String host) {
        testEnded();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public Object clone() {
        try {
            TestElement clonedElement = this.getClass().newInstance();

            PropertyIterator iter = propertyIterator();
            while (iter.hasNext()) {
                clonedElement.setProperty(iter.next().clone());
            }
            clonedElement.setRunningVersion(isRunningVersion());
            ((TrafgenCustomSampler) clonedElement).trafgenControl = this.trafgenControl;
            return clonedElement;
        } catch (InstantiationException | IllegalAccessException e) {
            throw new AssertionError(e); // clone should never return null
        }
    }
}
