package eu.gity.jmeter.ddos.slowloris;

import eu.gity.jmeter.ddos.common.Trafgen;
import org.apache.jmeter.gui.action.ActionNames;
import org.apache.jmeter.gui.action.ActionRouter;
import org.apache.jmeter.gui.action.Start;
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.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.reflect.Field;
import java.util.ArrayList;

/**
 * This Test Element class stores properties of Slowloris DoS attack starts python script executing the attack
 *
 * @author Gity a.s.
 */
public class SlowLorisSampler extends AbstractSampler implements TestStateListener, Trafgen {
    public static final String TARGETIP = "SlowLorisSampler.targetIP";
    public static final String NUMBER = "SlowLorisSampler.number";
    public static final String PORT = "SlowLorisSampler.port";
    public static final String USEHTTPS = "SlowLorisSampler.usehttps";
    public static final String RANDOMIZE = "SlowLorisSampler.randomize";
    public static final String scriptRelativePath = "/lib/ext/ddos";
    private static final long serialVersionUID = 240L;
    private static final Logger log = LoggerFactory.getLogger(SlowLorisSampler.class);
    SlowlorisThread slowlorisThread;
    private StringBuilder pid; //StringBuilder must be used to be able to clone a reference to test runtime instance of this TE (Strings are immutable in Java)

    /**
     * Finds out where JMeter resides
     *
     * @return path to JMeter main directory
     */
    public String getJmeterHomeDir() {
        if (System.getProperty("jmeter.home") == null) {
            File userDir = new File(System.getProperty("user.dir"));
            return userDir.getAbsoluteFile().getParent();
        } else {
            return System.getProperty("jmeter.home");
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public SampleResult sample(Entry e) {
        SampleResult res = new SampleResult();
        res.setSampleLabel(getName());
        res.sampleStart();

        String port = this.getPropertyAsString(PORT);
        if (port.equals("80")) {
            port = ""; //Port 80 is set automatically by the script
        } else {
            port = " -p " + port;
        }

        String useHttps = "";
        if (this.getPropertyAsBoolean(USEHTTPS)) {
            useHttps = " --https";
        }

        String randomize = "";
        if (this.getPropertyAsBoolean(RANDOMIZE)) {
            randomize = " -ua";
        }

        String command = "python " + getJmeterHomeDir() + scriptRelativePath + "/slowloris.py -v" + randomize + useHttps + port
                + " -s " + this.getPropertyAsString(NUMBER) + " " + this.getPropertyAsString(TARGETIP);

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

        slowlorisThread.setArgs(args);
        slowlorisThread.start();
        if (slowlorisThread.getResult()) {
            res.setDataType(SampleResult.TEXT);
            res.setResponseOK();
        } else {
            res.setResponseCode("500");// $NON-NLS-1$
            res.setResponseMessage("FAILED");
        }
        res.sampleEnd();

        return res;
    }

    /**
     * Terminates slowloris processes for a specific sampler instance (i.e. when DDoS Thread Group is terminated by scheduler)
     * or all processes (i.e. when "Stop Test" button is pressed)
     *
     * @param killall determines if all processes regardless of Thread Group should be terminated
     *                but true means terminating all running python scripts so it is used only in emergency situations
     */
    public void terminateSlowloris(boolean killall) {
        String termCommand;
        String killCommand;
        if (killall) {
            //TODO search for slowloris PIDs instead of killing all python scripts
            termCommand = "killall -SIGTERM python";
            killCommand = "killall -SIGKILL python";
        } else {
            termCommand = "pkexec kill -SIGTERM " + pid;
            killCommand = "pkexec kill -SIGKILL " + pid;
        }
        try {
            //Tries to terminate specific/all slowloris processes
            String[] termArgs = new String[]{"/bin/bash", "-c", termCommand};
            Process termPr = new ProcessBuilder(termArgs).start();
            termPr.waitFor();
            if (termPr.exitValue() != 0) {
                //Tries to kill specific/all slowloris processes if termination did not succeeded
                String[] killArgs = new String[]{"/bin/bash", "-c", killCommand};
                Process killPr = new ProcessBuilder(killArgs).start();
                killPr.waitFor();
                if (killPr.exitValue() != 0) {
                    //Tries to kill all slowloris processes if killing specific processes did not succeeded
                    String[] killAllArgs = new String[]{"/bin/bash", "-c", "killall -SIGKILL python"};
                    Process killAllPr = new ProcessBuilder(killAllArgs).start();
                    killAllPr.waitFor();
                }
            }
        } catch (IOException e) {
            log.error("Error occured while terminating slowloris processes:\n" + e.getStackTrace().toString());
        } catch (InterruptedException e) {
            log.error("Termination of slowloris processes was interrupted:\n" + e.getStackTrace().toString());
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void stopTrafgen() {
        terminateSlowloris(false);
    }

    @Override
    public void checkPacketCount(String tgName, long duration) {
        // Not used as there is no actual packet count parameter for the Slowloris script
    }

    /**
     * {@inheritDoc}
     * <p>
     * A {@code pid} StringBuilder object is instantiated here. A reference to it is then passed into test runtime
     * version of this TE when clonning. Finally a value is assingned when launching the process in sample()
     */
    @Override
    public void testStarted() {
        pid = new StringBuilder(); //A new value to the pid object will be assigned in sample()
        slowlorisThread = new SlowlorisThread();
    }

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

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

    /**
     * {@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());
            ((SlowLorisSampler) clonedElement).pid = this.pid;
            ((SlowLorisSampler) clonedElement).slowlorisThread = this.slowlorisThread;
            return clonedElement;
        } catch (InstantiationException | IllegalAccessException e) {
            throw new AssertionError(e); // clone should never return null
        }
    }

    @Override
    public void setDynamicRate(boolean dynamic) {
    }

    @Override
    public void changeRate(int rateChange) {
    }

    /**
     * A class which implements thread for running slowloris process
     *
     * @author Gity a.s.
     */
    public class SlowlorisThread extends Thread {
        private String[] args;
        private boolean result;

        /**
         * Allows to set the command including arguments
         *
         * @param args including arguments
         */
        public void setArgs(String[] args) {
            this.args = args;
        }

        /**
         * Allows access to result of the script run
         *
         * @return true if there were no errors
         */
        public boolean getResult() {
            return result;
        }

        /**
         * {@inheritDoc}
         */
        public void run() {
            ActionRouter.getInstance().addPostActionListener(Start.class, new ActionListener() {
                @Override
                public void actionPerformed(ActionEvent e) {
                    if (e.getActionCommand().equals(ActionNames.ACTION_STOP) || e.getActionCommand().equals(ActionNames.ACTION_SHUTDOWN)) {
                        terminateSlowloris(false);
                    }
                }
            });
            try {
                Runtime rt = Runtime.getRuntime();
                Process process;
                process = rt.exec(args);

                try {
                    Field f = process.getClass().getDeclaredField("pid");
                    f.setAccessible(true);
                    pid.append(f.get(process).toString());
                } catch (Exception e) {
                    log.error("Error getting PID of newly started slowloris script. All running python scripts will be terminated now! \n" + e.getStackTrace().toString());
                    terminateSlowloris(true);
                }

                //Redirects errors from stdout to JMeter log
                BufferedReader errorStream = new BufferedReader(new InputStreamReader(process.getErrorStream()));
                ArrayList<String> errorList = new ArrayList<String>();
                String line = null;
                while ((line = errorStream.readLine()) != null) {
                    errorList.add(line);
                }
                //Unfortunately python script writes info/debug messages into error output and then process finishes with error code !=0
                //thus here all error outputs are taken only as info messages
                log.info(errorList.toString());

    			/*
    			if (errorList.size() > 0) {
    				log.error("Slowloris error: \n" + errorList.toString());
    			}
    			*/

                process.waitFor();
                /*
                if (process.exitValue() == 0) {
                	result = true;
                } else {
                	result = false;
                }
                */
                result = true;
            } catch (IOException e) {
                log.error("Error occured while generating traffic:\n" + e.getStackTrace().toString());
                result = false;
            } catch (InterruptedException e) {
                log.warn("Slowloris was interrupted");
                result = true;
            }
        }
    }
}
