import org.apache.bcel.classfile.ClassParser;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.*;
import org.apache.bcel.Constants;

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

public class HookCheckerBorker implements ClassFileTransformer {

    public byte[] transform(ClassLoader loader, String className, Class<?> redefinedClass, ProtectionDomain domain,
                            byte[] classBuffer)
            throws IllegalClassFormatException {
        /* logic here */
        if (className.equals("java.lang.Class")) {
            try {
                ClassGen cgen = new ClassGen(new ClassParser(new ByteArrayInputStream(classBuffer), className).parse());
                Method gdm = cgen.containsMethod("getDeclaredMethod",
                        "(Ljava/lang/String;[Ljava/lang/Class;)Ljava/lang/reflect/Method;");
                MethodGen gdmGen = new MethodGen(gdm, cgen.getClassName(),
                        cgen.getConstantPool());
                InstructionList injected = new InstructionList();
                InstructionFactory factory = new InstructionFactory(cgen, cgen.getConstantPool());
                injected.append(factory.createInvoke("java.lang.Thread", "currentThread",
                        new ObjectType("java.lang.Thread"), new Type[0], Constants.INVOKESTATIC));
                injected.append(factory.createInvoke("java.lang.Thread", "getStackTrace",
                        new ArrayType(new ObjectType("java.lang.StackTraceElement"), 1), new Type[0],
                        Constants.INVOKEVIRTUAL));
                injected.append(new ICONST(3));
                injected.append(new AALOAD());
                injected.append(factory.createInvoke("java.lang.StackTraceElement", "getClassName", Type.STRING,
                        new Type[0], Constants.INVOKEVIRTUAL));
                injected.append(factory.createConstant("run.checkHooked"));
                injected.append(factory.createInvoke("java.lang.String", "equals", Type.BOOLEAN,
                        new Type[] { Type.OBJECT }, Constants.INVOKEVIRTUAL));
                BranchHandle ifEqBranch = injected.append(new IFEQ(null));
                injected.append(factory.createNew(new ObjectType("java.lang.NoSuchMethodException")));
                injected.append(new DUP());
                injected.append(factory.createInvoke("java.lang.NoSuchMethodException", "<init>", Type.VOID,
                        new Type[0], Constants.INVOKESPECIAL));
                injected.append(new ATHROW());
                InstructionHandle ldThisHandle = null;
                InstructionList gdmCodeList = gdmGen.getInstructionList();
                for (InstructionHandle handle : gdmCodeList.getInstructionHandles()) {
                    if (handle.getInstruction() instanceof ALOAD
                            && handle.getPrev().getInstruction() instanceof ALOAD) {
                        ALOAD aload = (ALOAD) handle.getInstruction();
                        ALOAD aloadThis = (ALOAD) handle.getPrev().getInstruction();
                        if (aload.getIndex() == 2 && aloadThis.getIndex() == 0) {
                            ldThisHandle = handle.getPrev();
                            break;
                        }
                    }
                }
                if (ldThisHandle != null) {
                    ifEqBranch.setTarget(ldThisHandle);
                    gdmCodeList.insert(ldThisHandle, injected);
                    gdmCodeList.setPositions();
                    gdmGen.setInstructionList(gdmCodeList); //redundant -- gdmCodeList == gdmGen.getInstructionList()
                    gdmGen.setMaxLocals();
                    gdmGen.setMaxStack();
                    cgen.replaceMethod(gdm, gdmGen.getMethod());
                    return cgen.getJavaClass().getBytes();
                }
            } catch (IOException ioex) {
                ioex.printStackTrace();
            }
        }
        return null;
    }


    public static void premain(String args, Instrumentation instr) {
        instr.addTransformer(new HookCheckerBorker());
    }
}