package org.bdware.analysis.taint;

import java.util.List;

import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.InvokeDynamicInsnNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.TypeInsnNode;
import org.objectweb.asm.tree.analysis.AnalyzerException;
import org.objectweb.asm.tree.analysis.Interpreter;

public class TaintInterpreter extends Interpreter<TaintValue> {
	TaintResult currentResult;
	TaintBits taintBits;

	public TaintInterpreter(int api) {
		super(api);

	}

	@Override
	public TaintValue newValue(Type type) {
		if (type != null)
			return new TaintValue(type.getSize());
		return new TaintValue(1);
	}

	@Override
	public TaintValue newOperation(AbstractInsnNode insn) throws AnalyzerException {
		TaintValue ret;
		switch (insn.getOpcode()) {
		case Opcodes.LCONST_0:
		case Opcodes.LCONST_1:
		case Opcodes.DCONST_0:
		case Opcodes.DCONST_1:
			ret = new TaintValue(2);
			break;
		case Opcodes.NEW:
			TypeInsnNode typeNode = (TypeInsnNode) insn;
			if (typeNode.desc.startsWith("wrp/jdk/nashorn/internal/scripts/JO")) {
				ret = new HeapObject();
			} else {
				ret = new TaintValue(1);
			}
			break;
		default:
			ret = new TaintValue(1);
		}
		return ret;
	}

	@Override
	public TaintValue copyOperation(AbstractInsnNode insn, TaintValue value) throws AnalyzerException {
		return value;
	}

	@Override
	public TaintValue unaryOperation(AbstractInsnNode insn, TaintValue value) throws AnalyzerException {
		return value;
	}

	@Override
	public TaintValue binaryOperation(AbstractInsnNode insn, TaintValue value1, TaintValue value2)
			throws AnalyzerException {
		/*
		 * IALOAD, LALOAD, FALOAD, DALOAD, AALOAD, BALOAD, CALOAD, SALOAD, IADD, LADD,
		 * FADD, DADD, ISUB, LSUB, FSUB, DSUB, IMUL, LMUL, FMUL, DMUL, IDIV, LDIV, FDIV,
		 * DDIV, IREM, LREM, FREM, DREM, ISHL, LSHL, ISHR, LSHR, IUSHR, LUSHR, IAND,
		 * LAND, IOR, LOR, IXOR, LXOR, LCMP, FCMPL, FCMPG, DCMPL, DCMPG, IF_ICMPEQ,
		 * IF_ICMPNE, IF_ICMPLT, IF_ICMPGE, IF_ICMPGT, IF_ICMPLE, IF_ACMPEQ, IF_ACMPNE,
		 * PUTFIELD
		 */
		TaintValue ret = new TaintValue(1);
		ret.isTainted = value1.isTainted | value2.isTainted;
		switch (insn.getOpcode()) {
		case Opcodes.DALOAD:
		case Opcodes.DADD:
		case Opcodes.DSUB:
		case Opcodes.DMUL:
		case Opcodes.DDIV:
		case Opcodes.DREM:
		case Opcodes.LALOAD:
		case Opcodes.LADD:
		case Opcodes.LSUB:
		case Opcodes.LMUL:
		case Opcodes.LDIV:
		case Opcodes.LREM:
		case Opcodes.LSHL:
		case Opcodes.LSHR:
		case Opcodes.LUSHR:
		case Opcodes.LAND:
		case Opcodes.LXOR:
		case Opcodes.LOR:
			ret.size = 2;
		default:
		}
		return ret;
	}

	@Override
	public TaintValue ternaryOperation(AbstractInsnNode insn, TaintValue value1, TaintValue value2, TaintValue value3)
			throws AnalyzerException {
		// TODO
		value1.isTainted |= value3.isTainted;
		return value1;
	}

	@Override
	public TaintValue naryOperation(AbstractInsnNode insn, List<? extends TaintValue> values) throws AnalyzerException {
		int size = 1;
		/*
		 * INVOKEVIRTUAL, INVOKESPECIAL, INVOKESTATIC, INVOKEINTERFACE, MULTIANEWARRAY
		 * and INVOKEDYNAMIC
		 */
		String desc;
		switch (insn.getOpcode()) {
		case Opcodes.INVOKEVIRTUAL:
		case Opcodes.INVOKESPECIAL:
		case Opcodes.INVOKESTATIC:
		case Opcodes.INVOKEINTERFACE:
			desc = ((MethodInsnNode) insn).desc;
			if (desc != null && (desc.charAt(desc.length() - 1) == 'D' || desc.charAt(desc.length() - 1) == 'J'))
				size = 2;
			break;
		case Opcodes.INVOKEDYNAMIC:
			desc = ((InvokeDynamicInsnNode) insn).desc;
			if (desc != null && (desc.charAt(desc.length() - 1) == 'D' || desc.charAt(desc.length() - 1) == 'J'))
				size = 2;

			// Extra..Judgeffff x
			if (((InvokeDynamicInsnNode) insn).name.startsWith("dyn:setElem|setProp:")) {
				if (values.size() == 2) {
					values.get(0).isTainted |= values.get(1).isTainted;
				}
			}
			break;
		default:
		}
		TaintValue ret = new TaintValue(size);
		for (TaintValue v : values)
			if (v != null)
				ret.isTainted |= v.isTainted;
		if (insn instanceof InvokeDynamicInsnNode) {
			long taint = (taintBits.getTaintValue(((InvokeDynamicInsnNode) insn).name + insn.hashCode()));

			if (taint > 0L)
				ret.isTainted |= taint;
		}
		return HeapObject.operate(insn, values, ret);
	}

	@Override
	public void returnOperation(AbstractInsnNode insn, TaintValue value, TaintValue expected) throws AnalyzerException {
		if (value instanceof HeapObject) {
			currentResult.ret.isTainted |= ((HeapObject) value).wholeTaint();
		} else if (value!=null)
			currentResult.ret.isTainted |= value.isTainted;
	}

	@Override
	public TaintValue merge(TaintValue v, TaintValue w) {
		TaintValue ret = new TaintValue(v.getSize());
		ret.isTainted |= v.isTainted;
		ret.isTainted |= w.isTainted;
		return ret;
	}

	public void setCurrentResult(TaintResult naiveTaintResult) {
		currentResult = naiveTaintResult;
	}

	public void setTaintBits(TaintBits tb) {
		taintBits = tb;
	}
}