public final class AgentFactory {
    private final static Class<?> vmClass;
    private final static Method attach;
    private final static Method loadAgent;
    private final static Method detach;

    private final static List<Path> generatedAgents = new ArrayList<>();

    /**
     * Starts building a new Java agent jar
     *
     * @param agentClass Agent main class name
     * @param jarBuilder Jar builder
     * @return Agent jar path
     */
    @NotNull
    public static Path newJavaAgent(@NotNull String agentClass, @NotNull BiConsumer<Manifest, AgentJarOutputStream> jarBuilder) throws Exception {
        // Create new temporary file for agent jar
        // Note: might have an issue with systems cleaning temporary directory automatically...
        // Also it's up to user to clean up agent jar later.
        // https://stackoverflow.com/questions/28752006/alternative-to-file-deleteonexit-in-java-nio
        // https://stackoverflow.com/questions/18146637/delete-on-close-deletes-files-before-close-on-linux
        Path jarFile = Files.createTempFile("agent", ".jar");

        // Generate manifest for agent
        Manifest manifest = new Manifest();
        Attributes mainAttributes = manifest.getMainAttributes();
        mainAttributes.put(Attributes.Name.MANIFEST_VERSION, "1.0");
        mainAttributes.put(new Attributes.Name("Agent-Class"), agentClass);
        mainAttributes.put(new Attributes.Name("Can-Retransform-Classes"), "true");
        mainAttributes.put(new Attributes.Name("Can-Redefine-Classes"), "true");

        // Build jar contents
        try(AgentJarOutputStream os = new AgentJarOutputStream(Files.newOutputStream(jarFile), manifest, agentClass)) {
            jarBuilder.accept(manifest, os);
        }

        generatedAgents.add(jarFile);

        return jarFile;
    }

    /**
     * Starts building a new Java agent jar
     *
     * @param agentClass Agent main class name
     * @param jarBuilder Jar builder
     * @return Agent jar path
     */
    @NotNull
    public static Path newJavaAgent(@NotNull String agentClass, @NotNull Consumer<AgentJarOutputStream> jarBuilder) throws Exception {
        return newJavaAgent(agentClass, (manifest, jarOutputStream) -> jarBuilder.accept(jarOutputStream));
    }

    /**
     * Builds a new Java agent jar from {@code agentClass}
     *
     * @param agentClass Agent class
     * @return Agent jar path
     */
    @NotNull
    public static Path newJavaAgent(@NotNull Class<?> agentClass) throws Exception {
        return newJavaAgent(agentClass.getName(), (UncheckedConsumer<AgentJarOutputStream, IOException>) jarBuilder -> jarBuilder.addClassToJar(agentClass));
    }

    /**
     * Gets JVM PID
     *
     * @return JVM PID
     */
    public static long getCurrentPID() {
        String jvm = ManagementFactory.getRuntimeMXBean().getName();
        return Long.parseLong(jvm.substring(0, jvm.indexOf('@')));
    }

    /**
     * Attaches an agent to JVM
     *
     * @param jarPath Jar path
     */
    public static void attachAgent(@NotNull Path jarPath) throws Exception {
        Object vm = attach.invoke(null, "" + getCurrentPID());
        loadAgent.invoke(vm, jarPath.toString());
        detach.invoke(vm);
    }

    static {
        try {
            // TODO: very very very likely not compatible with Java 9 right now
            Path toolsJar = Paths.get(System.getProperty("java.home"), "..", "lib/tools.jar")
                    .toRealPath()
                    .toAbsolutePath();

            if(Files.notExists(toolsJar))
                throw new IllegalStateException("JDK is required for runtime transformations!");

            Class<?> _vmClass;
            try {
                _vmClass = Class.forName("com.sun.tools.attach.VirtualMachine");
            } catch (ClassNotFoundException e) {
                ClassLoaderTools.URLClassLoaderTools uclTool = new ClassLoaderTools.URLClassLoaderTools(
                        ClassLoader.getSystemClassLoader());
                uclTool.addURL(toolsJar.toUri().toURL());
                uclTool.resetCache();
                _vmClass = Class.forName("com.sun.tools.attach.VirtualMachine");
            }

            vmClass = _vmClass;
            attach = vmClass.getMethod("attach", String.class);
            loadAgent = vmClass.getMethod("loadAgent", String.class);
            detach = vmClass.getMethod("detach");
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

        // Add runtime hook to clean up generated agents
        if(!Boolean.getBoolean("shuriken.dontCleanUpAgents")) {
            Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                generatedAgents.forEach((SilentConsumer<Path>) Files::deleteIfExists);
            }, "Shuriken Generated Agent cleanup"));
        }
    }
} public class TestAgent {
    public static synchronized void agentmain(String args, Instrumentation instrumentation) throws Exception {
        System.setProperty("shuriken.testagent1", "present");
    }
} public class TestAgent2 implements ClassFileTransformer {
    private final static String TARGET_CL = "eu/mikroskeem/test/shuriken/instrumentation/testclasses/TestTransformable2";
    private final static int TARGET_A = Opcodes.ACC_PUBLIC + Opcodes.ACC_SUPER;
    private final static String TARGET_M = "a";
    private final static String TARGET_S = new Descriptor().returns(String.class).build();

    public static synchronized void agentmain(String args, Instrumentation instrumentation) throws Exception {
        new TestAgent2(instrumentation);
    }

    private TestAgent2(Instrumentation instrumentation) throws Exception {
        instrumentation.addTransformer(this, false);
        //Class<?> target = Class.forName(TARGET_CL.replace('/', '.'));
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if(!className.equals(TARGET_CL))
            return classfileBuffer;

        ClassReader cr = new ClassReader(classfileBuffer);
        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
        cr.accept(cw, 0);

        // Implement method `a`
        MethodVisitor mv = cw.visitMethod(TARGET_A, TARGET_M, TARGET_S, null, null);
        GeneratorAdapter ga = new GeneratorAdapter(mv, TARGET_A, TARGET_M, TARGET_S);
        ga.visitCode();
        ga.visitLdcInsn("foobarbaz");
        ga.returnValue();
        ga.visitMaxs(0, 0);
        ga.visitEnd();

        return cw.toByteArray();
    }
} public class TestAgent3 implements ClassFileTransformer {
    private final static String TARGET_CL = "eu/mikroskeem/test/shuriken/instrumentation/testclasses/TestTransformable3";
    private final static int TARGET_A = Opcodes.ACC_PUBLIC + Opcodes.ACC_STATIC;
    private final static String TARGET_M = "incrementFirst";
    private final static String TARGET_S = new Descriptor().accepts(int[].class).build();

    public static synchronized void agentmain(String args, Instrumentation instrumentation) throws Exception {
        new TestAgent3(instrumentation);
    }

    private TestAgent3(Instrumentation instrumentation) throws Exception {
        instrumentation.addTransformer(this, true);
        Class<?> target = Class.forName(TARGET_CL.replace('/', '.'));
        instrumentation.retransformClasses(target);
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if(!className.equals(TARGET_CL))
            return classfileBuffer;

        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
        ClassNode node = readClass(classfileBuffer);

        // Try to find target method
        MethodNode method = findMethod(node, TARGET_A, TARGET_M, TARGET_S);
        Ensure.notNull(method, "Could not find target method!");

        // Find ICONST_1 instruction (note: very shitty way, but then again since it's only usage so it's fine)
        // Note: x++ --> x = x + 1
        InsnNode insnNode = findInstruction(method.instructions, InsnNode.class, insn -> insn.getOpcode() == Opcodes.ICONST_1);
        Ensure.notNull(insnNode, "Could not find target instruction!");

        // Replace instruction with new one (x + 1 --> x + 3)
        method.instructions.set(insnNode, new InsnNode(Opcodes.ICONST_3));

        node.accept(cw);
        return cw.toByteArray();
    }
} public class TestAgent4 implements ClassFileTransformer {
    private final static String TARGET_CL = "eu/mikroskeem/test/shuriken/instrumentation/testclasses/TestTransformable4";
    private final static int TARGET_A = Opcodes.ACC_PUBLIC;
    private final static String TARGET_M = "testProperty";
    private final static String TARGET_M2 = "testPrint";
    private final static String TARGET_S = Descriptor.DEFAULT;

    private final static String ROUTE_CL = "java/lang/System";
    private final static String ROUTE_M = "getProperties";
    private final static String ROUTE_S = new Descriptor().returns(Properties.class).build();

    private final static String ROUTE_F = "out";
    private final static String ROUTE_F_T = Type.getDescriptor(PrintStream.class);
    private final static String ROUTE_F_T_M = new Descriptor().returns(PrintStream.class).build();

    public static synchronized void agentmain(String args, Instrumentation instrumentation) throws Exception {
        new TestAgent4(instrumentation);
    }

    private TestAgent4(Instrumentation instrumentation) throws Exception {
        instrumentation.addTransformer(this, false);
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if(!className.equals(TARGET_CL))
            return classfileBuffer;
        try {
            ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
            ClassNode node = readClass(classfileBuffer);

            // Try to find target method
            MethodNode method = findMethodWithAccessAtleast(node, TARGET_A, TARGET_M, TARGET_S);
            Ensure.notNull(method, "Could not find target method!");

            MethodNode method2 = findMethodWithAccessAtleast(node, TARGET_A, TARGET_M2, TARGET_S);
            Ensure.notNull(method2, "Could not find target method 2!");

            // Find method invoke insn
            MethodInsnNode insnNode = findMethodInvocation(method.instructions, STATIC, ROUTE_CL, ROUTE_M, ROUTE_S);
            Ensure.notNull(insnNode, "Could not find System.getProperties() call!");

            // Find field get insn
            FieldInsnNode fInsnNode = findFieldInstruction(method2.instructions, GETSTATIC, ROUTE_CL, ROUTE_F, ROUTE_F_T);
            Ensure.notNull(fInsnNode, "Could not find GETSTATIC System.out");

            // Reroute method
            rerouteMethodInvocation(method.instructions, insnNode, TARGET_CL, "getDummyProperties", ROUTE_S);

            // Reroute field
            rerouteFieldGetter(method2.instructions, fInsnNode, TARGET_CL, "getDummyStream", ROUTE_F_T_M);

            node.accept(cw);
            return cw.toByteArray();
        } catch (Throwable t) {
            t.printStackTrace();
            return classfileBuffer;
        }
    }
} public class TestAgent5 implements ClassFileTransformer {
    private final static String TARGET_CL = "eu/mikroskeem/test/shuriken/instrumentation/testclasses/TestTransformable5";
    private final static int TARGET_A = Opcodes.ACC_PUBLIC;
    private final static String TARGET_M = "testProperty";
    private final static String TARGET_S = Descriptor.DEFAULT;

    private final static String ROUTE_CL = "java/util/Properties";
    private final static String ROUTE_M = "getProperty";
    private final static String ROUTE_S = new Descriptor().accepts(String.class).returns(String.class).build();

    private final static String ROUTE_T_CL = "eu/mikroskeem/test/shuriken/instrumentation/testclasses/TestRerouteTarget5";
    private final static String ROUTE_T_M_M = "rerouteGetProperty";
    private final static String ROUTE_T_M_S = new Descriptor().accepts(Properties.class, String.class).returns(String.class).build();
    private final static String ROUTE_T_F_M = "rerouteStringSet";
    private final static String ROUTE_T_F_S = new Descriptor().accepts(Type.getType("L" + TARGET_CL + ";"), Type.getType(String.class)).build();

    public static synchronized void agentmain(String args, Instrumentation instrumentation) throws Exception {
        new TestAgent5(instrumentation);
    }

    private TestAgent5(Instrumentation instrumentation) throws Exception {
        instrumentation.addTransformer(this, false);
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if(!className.equals(TARGET_CL))
            return classfileBuffer;
        try {
            ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
            ClassNode node = readClass(classfileBuffer);

            // Try to find target method
            MethodNode method = findMethodWithAccessAtleast(node, TARGET_A, TARGET_M, TARGET_S);
            Ensure.notNull(method, "Could not find target method!");

            // Find method invoke insn
            MethodInsnNode insnNode = findMethodInvocation(method.instructions, VIRTUAL, ROUTE_CL, ROUTE_M, ROUTE_S);
            Ensure.notNull(insnNode, "Could not find Properties.getProperty(String) call!");

            // Find field insn node
            FieldInsnNode fInsnNode = findFieldInstruction(method.instructions, PUT, TARGET_CL, "foundProperty", Type.getDescriptor(String.class));
            Ensure.notNull(fInsnNode, "Could not find TestTransformable.foundProperty setter");

            // Reroute method
            rerouteMethodInvocation(method.instructions, insnNode, ROUTE_T_CL, ROUTE_T_M_M, ROUTE_T_M_S);

            // Reroute field
            rerouteFieldSetter(method.instructions, fInsnNode, ROUTE_T_CL, ROUTE_T_F_M, ROUTE_T_F_S);

            node.accept(cw);
            return cw.toByteArray();
        } catch (Throwable t) {
            t.printStackTrace();
            return classfileBuffer;
        }
    }
} public class TestTransformable2 implements TestTransformable2Interface {

} public class TestTransformable3 {
    public static void incrementFirst(int[] array) {
        array[0]++;
    }
} public class TestTransformable4 {
    private static PrintStream dummyStream = new PrintStream(System.out) {
        @Override
        public void println(String text) {
            super.println("[Retransformed] " + text);
        }
    };

    public void testProperty() {
        Properties properties = System.getProperties();
        if(properties.getProperty("shuriken.testagent4").equals("true")) {
            System.out.println("Method reroute failed");
        }
    }

    public void testPrint() {
        System.out.println("Test");
    }

    // System.getProperties() gets replaced with this
    public static Properties getDummyProperties() {
        Properties props = new Properties();
        props.setProperty("shuriken.testagent4", "false");
        return props;
    }

    // GETSTATIC System.out gets replaced with this
    public static PrintStream getDummyStream() {
        return dummyStream;
    }
} public class TestTransformable5 {
    private String foundProperty;

    public void testProperty() {
        Properties properties = System.getProperties();
        foundProperty = properties.getProperty("java.home");
    }
}
import eu.mikroskeem.shuriken.instrumentation.runtime.AgentFactory;
import eu.mikroskeem.test.shuriken.instrumentation.testagent.TestAgent;
import eu.mikroskeem.test.shuriken.instrumentation.testagent.TestAgent2;
import eu.mikroskeem.test.shuriken.instrumentation.testagent.TestAgent3;
import eu.mikroskeem.test.shuriken.instrumentation.testagent.TestAgent4;
import eu.mikroskeem.test.shuriken.instrumentation.testagent.TestAgent5;
import eu.mikroskeem.test.shuriken.instrumentation.testclasses.TestTransformable2;
import eu.mikroskeem.test.shuriken.instrumentation.testclasses.TestTransformable3;
import eu.mikroskeem.test.shuriken.instrumentation.testclasses.TestTransformable4;
import eu.mikroskeem.test.shuriken.instrumentation.testclasses.TestTransformable5;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.nio.file.Path;
package eu.mikroskeem.test.shuriken.instrumentation;



/**
 * @author Mark Vainomaa
 */
public class RuntimeTransformationTester {
    private static boolean agentFactoryWorks = false;

    @BeforeAll
    public static void testIfAgentFactoryWorks() throws Exception {
        try {
            AgentFactory.class.getName();
            agentFactoryWorks = true;
        } catch (Exception e) {
            System.out.println("Skipping runtime transformation unit tests, because they are not working with this JVM");
            e.printStackTrace();
        }
    }

    @Test
    public void testAgentLoading() throws Exception {
        Assumptions.assumeTrue(agentFactoryWorks, "Agent factory is not working");

        Path agentFile = AgentFactory.newJavaAgent(TestAgent.class);
        AgentFactory.attachAgent(agentFile);

        Assertions.assertEquals("present", System.getProperty("shuriken.testagent1"));
    }

    @Test
    public void testTransforming() throws Exception {
        Assumptions.assumeTrue(agentFactoryWorks, "Agent factory is not working");

        Path agentFile = AgentFactory.newJavaAgent(TestAgent2.class);
        AgentFactory.attachAgent(agentFile);

        // Check if method `a` is implemented
        TestTransformable2 test = new TestTransformable2();