import org.apache.bcel.classfile.*;
import org.apache.bcel.generic.*;
import org.apache.bcel.*;

import java.io.*;
import java.util.*;

public class BCELplay {

    JavaClass j;
    ClassGen cg;
    ConstantPoolGen cp;
    InstructionFactory _if;

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

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

    public void getMethods() {
        for(Method m : cg.getMethods()) {
            if(m.getName().equalsIgnoreCase("lol")) {
              cg.removeMethod(m);
               w("[NOTICE] REMOVED LOL()");
            } else
                System.out.println("Method in '"+cg.getClassName()+"': "+m);
        }
    }

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

    public void getFields() {
        for(Field m : cg.getFields()) {
            w("Field in ["+cg.getClassName()+"]: " + m.getName() + " with signature: " + m.getSignature() + " access flags: " + m.getAccessFlags());
        }
    }

    public void getSipushes() {
        for(Method m : j.getMethods()) {
            if(m.getName().equals("main")) {
                w("\tFound main method.");
                MethodGen mg = new MethodGen(m, m.getName(), new ConstantPoolGen(m.getConstantPool()));
                w("\tCreating InstructionList");
                InstructionList il = mg.getInstructionList();
                w("\tCreating InstructionHandlek");
                InstructionHandle ih = il.getEnd();
                InstructionHandle ih2 = ih.getPrev();
                while(ih2 != null) {
                    Instruction instruction = ih2.getInstruction();
                    if(instruction instanceof SIPUSH) {
                        SIPUSH push = (SIPUSH)instruction;
                        w("SIPUSH: " + push.getValue().intValue());
                    }
                    ih2 = ih2.getPrev();
                }
            }
        }
    }

    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() {
        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", "lol", Type.VOID, Type.NO_ARGS, Constants.INVOKESPECIAL));
                il.setPositions();
                method.setInstructionList(il);
                method.setMaxStack();
                method.setMaxLocals();
                cg.replaceMethod(m, method.getMethod());
                il.dispose();
            }
        }
    }

    public void inject() {
        try {
            InstructionList il = null;
            InstructionFactory factory = new InstructionFactory(cg, cp);
            MethodGen method = null;
            JavaClass _j = new ClassParser("Methods.class").parse();
            ClassGen _cg = new ClassGen(_j);
            Method m = null;
            for(Method d : _cg.getMethods()) {
                if(d.getName().equalsIgnoreCase("lol")) {
                    method = new MethodGen(d, d.getClass().getName(), cp);
                }
            }
            if(method != null) il = method.getInstructionList();
            InstructionHandle ih_0 = il.append(factory.createFieldAccess("java.lang.System", "out", new ObjectType("java.io.PrintStream"), Constants.GETSTATIC));
            il.append(new PUSH(cp, "This is a method i'm adding, oman did it work???"));
            il.append(factory.createInvoke("java.io.PrintStream", "println", Type.VOID, new Type[] { Type.STRING }, Constants.INVOKEVIRTUAL));
            InstructionHandle ih_8 = il.append(factory.createInvoke("HelloWorld", "sayHai", Type.VOID, Type.NO_ARGS, Constants.INVOKESTATIC));
            InstructionHandle ih_11 = il.append(factory.createReturn(Type.VOID));
            il.setPositions();
            method.setInstructionList(il);
            method.setMaxStack();
            method.setMaxLocals();
            cg.replaceMethod(m, method.getMethod());
        } catch(Exception e) {
            e.printStackTrace();
        }
    }
}