package org.example.jmssampler;

import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.camel.component.crypto.DigitalSignatureConstants;
import org.apache.jmeter.config.Arguments;
import org.apache.jmeter.protocol.java.sampler.AbstractJavaSamplerClient;
import org.apache.jmeter.protocol.java.sampler.JavaSamplerContext;
import org.apache.jmeter.samplers.SampleResult;

import javax.jms.*;
import java.io.FileInputStream;
import java.io.Serializable;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.Signature;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;


public class JMSSampler extends AbstractJavaSamplerClient implements Serializable {
    private static final long serialVersionUID = 1L;

    // set up default arguments for the JMeter GUI
    @Override
    public Arguments getDefaultParameters() {
        Arguments defaultParameters = new Arguments();
        defaultParameters.addArgument("username", "");
        defaultParameters.addArgument("password", "");
        defaultParameters.addArgument("queueName", "");
        defaultParameters.addArgument("protocol", "tcp");
        defaultParameters.addArgument("hostname", "localhost");
        defaultParameters.addArgument("port", "61616");
        defaultParameters.addArgument("keystorePath", "keystore.jks");
        defaultParameters.addArgument("keystorePassword", "password");
        defaultParameters.addArgument("truststorePath", "truststore.jks");
        defaultParameters.addArgument("truststorePassword", "password");
        defaultParameters.addArgument("signingAlias", "localclient");
        defaultParameters.addArgument("messagePath", "jmstest.xml");
        return defaultParameters;
    }

    @Override
    public SampleResult runTest(JavaSamplerContext context) {
        // pull parameters
        String username = context.getParameter("username");
        String password = context.getParameter("password");
        String queueName = context.getParameter("queueName");
        String protocol = context.getParameter("protocol");
        String hostname = context.getParameter("hostname");
        String port = context.getParameter("port");
        String endpoint = protocol + "://" + hostname + ":" + port;
        String keystorePath = context.getParameter("keystorePath");
        String keystorePassword = context.getParameter("keystorePassword");
        String trustStorePath = context.getParameter("trustStorePath");
        String trustStorePassword = context.getParameter("trustStorePassword");
        String signingAlias = context.getParameter("signingAlias");
        String message;
        String messagePath = context.getParameter("messagePath");
        SampleResult result = new SampleResult();
        result.sampleStart(); // start stopwatch

        try {
            // Get message from file
            // System.out.println("step 0: Load message");
            // System.out.println("Working Directory = " + System.getProperty("user.dir"));
            FileInputStream xml = null;
            StringBuilder builder = new StringBuilder();
            xml = new FileInputStream(messagePath);
            int content;
            while ((content = xml.read()) != -1) {
                // convert to char and display it
                builder.append((char) content);
            }
            message = builder.toString();

            // Get private key from keystore
            // System.out.println("step 1: Get keystore");
            KeyStore ks = KeyStore.getInstance("jks");
            char [] keyPwd = keystorePassword.toCharArray();
            ks.load(new FileInputStream(keystorePath), keyPwd);
            PrivateKey key = (PrivateKey) ks.getKey(signingAlias, keyPwd);

            // If private key available proceed
            if (key != null) {
                // System.out.println("step 2: Get signature instance");
                Signature signature = Signature.getInstance("SHA256withRSA");

                // System.out.println("step 3: Initiate signature");
                signature.initSign(key);

                // System.out.println("step 4: Update signature with message content");
                signature.update(message.getBytes());

                // System.out.println("step 5: Sign message with private key");
                byte[] messageSignatureByteArray = signature.sign();

                // System.out.println("step 6: Encode message to base64");
                String messageSignature =  Base64.getEncoder().encodeToString(messageSignatureByteArray);

                // System.out.println("step 7: Get jms header content");
//                // System.out.println(messageSignature);

                // System.out.println("step 8: Start jms connection");
                ConnectionFactory factory = new ActiveMQConnectionFactory(endpoint);
                Connection jmsConnection = factory.createConnection(username, password);
                jmsConnection.start();

                // System.out.println("step 9: Start session on jms queue");
                Session session = jmsConnection.createSession(false, Session.AUTO_ACKNOWLEDGE);
                Destination destinationQueue = session.createQueue(queueName);
                MessageProducer producer = session.createProducer(destinationQueue);
                producer.setDeliveryMode(DeliveryMode.NON_PERSISTENT);

                // Create the message
                // System.out.println("step 10: Create the message");
                Message outgoingMessage = session.createTextMessage(message);

                // Set JMS headers
                // System.out.println("step 11: Add header to jms request");
                Map<String,String> jmsHeaders = new HashMap<String, String>();
                jmsHeaders.put(DigitalSignatureConstants.SIGNATURE, messageSignature);
                outgoingMessage.setJMSCorrelationID(jmsHeaders.get("id"));
                for (Map.Entry<String,String> header: jmsHeaders.entrySet()) {
//                    // System.out.println("key: " + header.getKey() + "\tvalue: " + header.getValue());
                    outgoingMessage.setStringProperty(header.getKey(),header.getValue());
                }

                // System.out.println("step 12: Send message to queue");
                producer.send(outgoingMessage);

                // System.out.println("step 13: Close queue session");
                session.close();

                // System.out.println("step 14: Close jms connection");
                jmsConnection.close();

                result.sampleEnd(); // stop stopwatch
                result.setSuccessful(true);
                result.setResponseMessage("Successfully posted message into the queue");
                result.setResponseCodeOK(); // 200 code

            } else {
                result.sampleEnd(); // stop stopwatch
                result.setSuccessful(false);
                result.setResponseMessage("\tError: No private key found on keystore.");
                result.setResponseCode("500");
            }

        } catch (Exception e) {
            result.sampleEnd(); // stop stopwatch
            result.setSuccessful(false);
            result.setResponseMessage("Exception: " + e );

            // get stack trace as a String to return as document data
            java.io.StringWriter stringWriter = new java.io.StringWriter();
            e.printStackTrace(new java.io.PrintWriter(stringWriter));
            result.setDataType(org.apache.jmeter.samplers.SampleResult.TEXT);
            result.setResponseCode("500");
        }

        return result;
    }
}