Agent.java

import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;
import com.sun.tools.attach.AttachNotSupportedException;
import com.sun.tools.attach.VirtualMachine;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.lang.management.RuntimeMXBean;
import java.util.Properties;
import java.util.jar.Attributes;
import java.util.jar.JarOutputStream;
import java.util.jar.Manifest;

/**
 * User: Tyler Sedlar, TSedlar
 * License: Creative Commons Attribution 3.0 Unported License
 * Date: 11/10/13
 * Time: 9:32 AM
 */
public class Agent {

    public final String pid;
    public final Manifest manifest;
    public final File jar;

    public final VirtualMachine vm;
    public final Properties properties;

    private boolean detached = false;

    private String getPid() {
        RuntimeMXBean bean = ManagementFactory.getRuntimeMXBean();
        String pid = bean.getName();
        if (pid.contains("@")) {
            pid = pid.substring(0, pid.indexOf("@"));
        }
        return pid;
    }

    public Agent(Class<?> agent) throws IOException, AttachNotSupportedException, AgentLoadException, AgentInitializationException {
        this.pid = getPid();
        this.manifest = new Manifest();
        this.jar = File.createTempFile("agent-" + pid, ".jar");
        jar.deleteOnExit();
        Attributes attributes = manifest.getMainAttributes();
        attributes.put(Attributes.Name.MANIFEST_VERSION, "1.0");
        attributes.put(new Attributes.Name("Agent-Class"), agent.getName());
        attributes.put(new Attributes.Name("Can-Retransform-Classes"), "true");
        attributes.put(new Attributes.Name("Can-Redefine-Classes"), "true");
        try (JarOutputStream output = new JarOutputStream(new FileOutputStream(jar), manifest)) {
            output.flush();
        }
        this.vm = VirtualMachine.attach(pid);
        this.properties = vm.getSystemProperties();
        vm.loadAgent(jar.getAbsolutePath());
        Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {
            public void run() {
                try {
                    detached = true;
                    vm.detach();
                    System.out.println("Agent#" + pid + " detached");
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }));
    }

    public void detach() {
        detached = true;
    }

    public boolean attached() {
        return !detached;
    }
}