import org.apache.bcel.classfile.JavaClass;
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.FileOutputStream;

public class BCELplay {
    public static void main(String[] args) {
        new BCELplay();
    }

    public BCELplay() {
        try {
            w("Parsing HelloWorld.class");
            JavaClass j = new ClassParser("HelloWorld.class").parse();
            w("Making ClassGen for HelloWorld.class");
            ClassGen cg = new ClassGen(j);
            w("Getting the ConstantPoolGen.");
            ConstantPoolGen cp = cg.getConstantPool();
            w("Creating Instruction Factory");
            InstructionFactory _if = new InstructionFactory(cg, cp);
            w("Listing the methods!");
            getMethods(cg);
            w("Injecting Method");
            inject(cp, cg);
            w("Dumping the file!");
            cg.getJavaClass().dump(new FileOutputStream("HelloWorld.class"));
            w("Relisting Methods, seeing if it REALLY injected..");
            getMethods(cg);
        } catch(Exception e) {
            e.printStackTrace();
        }
    }

    public void getMethods(ClassGen _cg) {
        for(Method m : _cg.getMethods()) {
            System.out.println("Method in '"+_cg.getClassName()+"': "+m);
        }
    }

    public void w(String l) {
        System.out.println(l);
    }

    //injectInMain() -- Was gonna try to get a reference into the main method, so it calls the method.
    public void injectInMain(ClassGen cg, ConstantPoolGen cp, InstructionFactory _if) {
        for(Method m : cg.getMethods()) {
            if(m.getName().equalsIgnoreCase("main")) {
                Method mainMethod = m;
                MethodGen method = new MethodGen(mainMethod, mainMethod.getClass().getName(), cp);
                InstructionList il = method.getInstructionList();
                il.append(_if.createInvoke("HelloWorld", "hai", Type.NULL, Type.NO_ARGS, Constants.INVOKESTATIC));
                method.setMaxStack();
                method.setMaxLocals();
                il.dispose();
            }
        }
    }

    public void inject(ConstantPoolGen cp, ClassGen cg) {
        try {
            InstructionList il = new InstructionList();
            MethodGen method = null;
            JavaClass j = new ClassParser("Methods.class").parse();
            ClassGen _cg = new ClassGen(j);
            for(Method m : _cg.getMethods()) {
                if(m.getName().equalsIgnoreCase("lol")) {
                    method = new MethodGen(m, m.getClass().getName(), cp);
                }
            }
            
            method.setMaxStack();
            method.setMaxLocals();
            cg.addMethod(method.getMethod());
        } catch(Exception e) {
            e.printStackTrace();
        }
    }
}