package samples;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.IllegalClassFormatException;
import java.io.ByteArrayInputStream;
import java.security.ProtectionDomain;

import javassist.*;

public class InstrumentationSample3 {
    private static ClassPool classPool;

    public static void premain(String agentArgs, Instrumentation instrumentation) {
        classPool = ClassPool.getDefault();

        instrumentation.addTransformer(new ClassFileTransformer() {
                public byte[] transform(ClassLoader loader,
                                 String className,
                                 Class<?> classBeingRedefined,
                                 ProtectionDomain protectionDomain,
                                 byte[] classfileBuffer) throws IllegalClassFormatException {

                    if (className.equals("HelloWorld")) {
                        try {
                            ByteArrayInputStream stream = new ByteArrayInputStream(classfileBuffer);
                            CtClass ctClass = classPool.makeClass(stream);
                            
                            CtMethod ctMethod = ctClass.getDeclaredMethod("main");
                            ctMethod.insertBefore("System.out.println(\"Hello, Tiger!\");");
                            
                            return ctClass.toBytecode();
                        } catch (Exception ex) {
                            IllegalClassFormatException e =  new IllegalClassFormatException();
                            e.initCause(ex);
                            throw e;
                        }
                    } else {
                        return null;
                    }
                }
            });
    }
}