package alterrs.jbct.deob;

import static alterrs.asm.Opcodes.BIPUSH;
import static alterrs.asm.Opcodes.GOTO;
import static alterrs.asm.Opcodes.ICONST_0;
import static alterrs.asm.Opcodes.ICONST_1;
import static alterrs.asm.Opcodes.ICONST_2;
import static alterrs.asm.Opcodes.ICONST_3;
import static alterrs.asm.Opcodes.ICONST_4;
import static alterrs.asm.Opcodes.ICONST_5;
import static alterrs.asm.Opcodes.IFEQ;
import static alterrs.asm.Opcodes.IFGE;
import static alterrs.asm.Opcodes.IFGT;
import static alterrs.asm.Opcodes.IFLE;
import static alterrs.asm.Opcodes.IFLT;
import static alterrs.asm.Opcodes.IFNE;
import static alterrs.asm.Opcodes.IF_ACMPEQ;
import static alterrs.asm.Opcodes.IF_ACMPNE;
import static alterrs.asm.Opcodes.IF_ICMPEQ;
import static alterrs.asm.Opcodes.IF_ICMPGE;
import static alterrs.asm.Opcodes.IF_ICMPGT;
import static alterrs.asm.Opcodes.IF_ICMPLE;
import static alterrs.asm.Opcodes.IF_ICMPLT;
import static alterrs.asm.Opcodes.IF_ICMPNE;
import static alterrs.asm.Opcodes.ILOAD;
import static alterrs.asm.Opcodes.LDC;
import static alterrs.asm.Opcodes.NEW;
import static alterrs.asm.Opcodes.RETURN;
import static alterrs.asm.Opcodes.SIPUSH;

import java.util.HashSet;
import java.util.Set;

import org.slf4j.Logger;

import alterrs.asm.tree.AbstractInsnNode;
import alterrs.asm.tree.ClassNode;
import alterrs.asm.tree.JumpInsnNode;
import alterrs.asm.tree.MethodNode;
import alterrs.asm.tree.TypeInsnNode;
import alterrs.asm.tree.VarInsnNode;
import alterrs.commons.LoggingUtil;
import alterrs.jbct.utility.ASMUtil;
import alterrs.jbct.utility.AbstractTransformer;

/**
 * @author Lazaro Brito
 */
public class ControlFlowDeobfuscator extends AbstractTransformer {
	private static final Logger logger = LoggingUtil.log();

	private int dCount = 0;

	public ControlFlowDeobfuscator() {
		super(true);
	}

	@Override
	public void visitVarInsn(ClassNode c, MethodNode m, VarInsnNode n) {
		if (n.getOpcode() == ILOAD) {
			AbstractInsnNode n2 = ASMUtil.nextRealInsn(n.getNext());// constant
			if (n2 == null) {
				return;
			}

			AbstractInsnNode n3 = ASMUtil.nextRealInsn(n2.getNext());// comparator
			if (n3 == null) {
				return;
			}

			AbstractInsnNode n4 = ASMUtil.nextRealInsn(n3.getNext());// fail
			if (n4 == null) {
				return;
			}

			if (!isConst(n2) || !isComp(n3)) {
				return;
			}

			int n3JumpIndex = m.instructions.indexOf(((JumpInsnNode) n3).label);

			do {
				if (n4.getOpcode() == RETURN && n2.getOpcode() == LDC) {
					break;
				}

				if (n4.getOpcode() == GOTO) {
					JumpInsnNode n4j = (JumpInsnNode) n4;

					int myIndex = m.instructions.indexOf(n4);
					if (myIndex == m.instructions.indexOf(ASMUtil
							.nextRealInsn(n4j.label))) {
						break;
					} else if (m.instructions.indexOf(n4j.label) == n3JumpIndex) {
						break;
					}
				}

				if (n4.getOpcode() == NEW) {
					if (((TypeInsnNode) n4).desc
							.equals("java/lang/IllegalStateException")) {
						break;
					}
				}
				return;
			} while (false);

			Set<AbstractInsnNode> del = new HashSet<AbstractInsnNode>();

			for (AbstractInsnNode insn : m.instructions.toArray()) {
				if (insn.getOpcode() == ILOAD
						&& ((VarInsnNode) insn).var == n.var) {
					AbstractInsnNode insn2 = ASMUtil.nextRealInsn(insn
							.getNext());
					if (insn2 == null) {
						break;
					}

					AbstractInsnNode comp = ASMUtil.nextRealInsn(insn2
							.getNext());
					if (comp == null) {
						break;
					}

					if (!isComp(comp)) {
						insn = insn.getNext();
						continue;
					}

					int jumpIndex = m.instructions
							.indexOf(((JumpInsnNode) comp).label);

					AbstractInsnNode i = insn;
					while (i != null) {
						if (m.instructions.indexOf(i) >= jumpIndex) {
							break;
						}

						del.add(i);
						i = i.getNext();
					}

					dCount++;
				}
			}

			for (AbstractInsnNode i : del) {
				m.instructions.remove(i);
			}
		}
	}

	private boolean isConst(AbstractInsnNode n) {
		switch (n.getOpcode()) {
		case BIPUSH:
		case SIPUSH:
		case LDC:
		case ICONST_0:
		case ICONST_1:
		case ICONST_2:
		case ICONST_3:
		case ICONST_4:
		case ICONST_5:
			return true;
		}
		return false;
	}

	private boolean isComp(AbstractInsnNode n) {
		switch (n.getOpcode()) {
		case IFEQ:
		case IFNE:
		case IFLT:
		case IFGE:
		case IFGT:
		case IFLE:
		case IF_ICMPEQ:
		case IF_ICMPNE:
		case IF_ICMPLT:
		case IF_ICMPGE:
		case IF_ICMPGT:
		case IF_ICMPLE:
		case IF_ACMPEQ:
		case IF_ACMPNE:
			return true;
		}
		return false;
	}

	@Override
	public void onFinish() {
		logger.info("Removed " + dCount + " dummy conditionals!");
	}
}
