package fr.lifl.stc.stan.analyser;

import java.util.Hashtable;
import java.util.Enumeration;

import org.apache.bcel.SignatureConstants;
import org.apache.bcel.classfile.Attribute;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.classfile.Field;
import org.apache.bcel.classfile.ExternalMethodsAttribute;
import org.apache.bcel.classfile.ExternalFieldsAttribute;
import org.apache.bcel.generic.ConstantPoolGen;
import org.apache.bcel.generic.InvokeInstruction;
import org.apache.bcel.generic.FieldInstruction;
import org.apache.bcel.generic.ReferenceType;
import org.apache.bcel.generic.BasicType;
import org.apache.bcel.Constants;

import fr.lifl.stc.stan.analyser.GlobalAnalyser;
import fr.lifl.stc.stan.execution.stack.JvmType;
import fr.lifl.stc.stan.link.Link;
import fr.lifl.stc.stan.link.TableLink;
import fr.lifl.stc.stan.policy.Policy;
import fr.lifl.stc.stan.samples.flow.FlowAnnotation;
import fr.lifl.stc.stan.samples.flow.FlowPartialSignature;
import fr.lifl.stc.stan.samples.flow.FlowPartialSignatureData;
import fr.lifl.stc.stan.signature.Signature;
import fr.lifl.stc.stan.signature.SignatureDictionary;
import fr.lifl.stc.stan.annotation.Annotation;
import fr.lifl.stc.stan.tools.HTMLOutput;
import fr.lifl.stc.stan.util.Printer;
import fr.lifl.stc.stan.dsl.FlowChecker;

/**
 * @author <a href="mailto:graziano@lifl.fr">Julien Graziano </a>
 *
 */
public class ClassAnalyser {

	public static long memory = 0L;

	/**
	 * dictionary containing signatures
	 */
	private SignatureDictionary dictionary;

	private Policy policy;

    private FlowChecker flowchecker; /** Checker for dsl policies */

	private JavaClass _class;

	private Method methods[];

	private ConstantPoolGen cpg;

	public int proof_index;

	private int externalMethods_index;
	private int externalFields_index;

	/**
	 * list the external methods invoked in the class
	 *
	 * key => method type
	 *
	 * where
	 * 		key = the key of the method in the dictionary
	 * 		method type = a combination between <link href=fr.lifl.rd2p.jits.tools.ea.signature.SignatureDictionary>
	 * 		SignatureDictionary.GLOBAL_SIGN_IDX</link> and <link href=fr.lifl.rd2p.jits.tools.ea.signature.SignatureDictionary>
	 * 		SignatureDictionary.EXACT_SIGN_IDX</link>, specifing which of global/exact signatures are nedeed
	 */
	private Hashtable<String,Integer> externalMethods;

	/**
	 * key => security level
	 * Where key is = className.fieldName
	 * the entry in the constant pool corresponding to the field
	 */
	private Hashtable<String,Integer> externalFields;


	/**
	 * Creates a new <code>ClassAnalyser</code> instance.
	 * @param dict
	 * @param pol
	 * @param c
	 */
	public ClassAnalyser(SignatureDictionary dict, Policy pol, FlowChecker fc, JavaClass c) {
		dictionary = dict;
		policy = pol;
        flowchecker = fc;
		_class = c;
		methods = c.getMethods();
		externalMethods = new Hashtable<String,Integer>();
		externalFields = new Hashtable<String,Integer>();
		cpg = new ConstantPoolGen(c.getConstantPool());
		prepareConstantPool();
	}

	/**
	 *
	 *
	 */
	private void prepareConstantPool() {
		proof_index = cpg.addUtf8("ProofMap");
		externalMethods_index = cpg.addUtf8("ExternalMethods");
		externalFields_index = cpg.addUtf8("ExternalFields");
		_class.setConstantPool(cpg.getFinalConstantPool());
	}

	/**
	 *
	 */
	public void analyse() {
		int i;
		boolean printed = false;
		MethodAnalyzer manalyser;

		for (i = 0; i < methods.length; i++) {
			HTMLOutput.newMethod(methods[i].getName(), methods[i].getSignature());
			printed = false;


			if (methods[i].isNative()) {
				//Printer.println("[run] Native method " + SignatureDictionary.key(_class,methods[i]));
				String msg = "native ";
				if (methods[i].isStatic())
					msg += "static\t";
				else
					msg += "\t";

				String key = SignatureDictionary.key(_class, methods[i]);
				if (dictionary.contains(key))
					continue;

				System.err.println("[run] Native method "+msg + this._class.getClassName() + "."
						+ methods[i].getName() + methods[i].getSignature()+" - signature not found: worst signature considered ");

				TableLink b = computeWorstSignature(methods[i]);

				FlowPartialSignature lps = new FlowPartialSignature(FlowAnnotation.class.getName(), b);
				Signature s = new Signature();
				s.add(lps);

				// HTMLOutput.newMethod(methods[i].getName(), methods[i].getSignature());
				//HTMLOutput.appendln(b);
				dictionary.putSignature(SignatureDictionary.key(_class, methods[i]), s, true);
			} else if (methods[i].isAbstract()) {
				if (dictionary.contains(SignatureDictionary.key(_class,
						methods[i])))
					continue;

				TableLink b = null;
				if (Config.worstSignatureEnabled() && methods[i].isPublic()) {
					b = computeWorstSignature(methods[i]);
				} else {
					int n = methods[i].getArgumentTypes().length;
					int dec;
					if (methods[i].isStatic())
						dec = 4;
					else
						dec = 5;
					b = new TableLink(n + dec, n + dec);
				}
				FlowPartialSignature lps = new FlowPartialSignature(FlowAnnotation.class.getName(), b);
				Signature s = new Signature();
				s.add(lps);

				dictionary.putSignature(SignatureDictionary.key(_class, methods[i]), s , false);
			} else {
			//	String key = SignatureDictionary.key(_class, methods[i]);
			//	Signature sgn = dictionary.findExactSignature(key);
				{

					manalyser = null;
					Annotation[] annotations = null;

					Printer.println("[run] Analysing method "
							+ methods[i].getName() + methods[i].getSignature());
					manalyser = new MethodAnalyzer(dictionary, policy, flowchecker, cpg, _class,
                                                   methods[i], this, proof_index);

					//if (Config.implicitFlowEnabled())
					//	manalyser.analyseImplicitFlow();

					/**
					 * analyses the method
					 */
					manalyser.analyse();

//					System.gc();
//					System.gc();
//					long memoryEnd = Runtime.getRuntime().freeMemory();
//					Stats.addMemory(memoryStart - memoryEnd);

					/**
					 * prints friendly-html
					 */
					annotations = manalyser.getAnnotations();
					for (int j = 0; j < annotations.length; j++) {
						HTMLOutput.appendln(annotations[j].toHTML());
						annotations[1].addResults();
					}
					printed = true;

					/**
					 * annotates LinkMap and ProofMap
					 */
					if (Config.annotateEnabled())
						manalyser.annotateProof();

					Stats.incAnalyzedMethod();

				}
				memory = Math.max(memoryUsed(), memory);
			}
			if(!printed) {
				Signature aux =	dictionary.findGlobalSignature(SignatureDictionary.key(_class, methods[i]));
				HTMLOutput.appendln("<h4>Information Flow</h4>");
				if(aux!=null) {
					HTMLOutput.appendln("<p>global</p>");
					HTMLOutput.appendln(aux.toHTMLString());
				}
				aux = dictionary.findExactSignature(SignatureDictionary.key(_class, methods[i]));
				if(aux!=null) {
					HTMLOutput.appendln("<p>exact</p>");
					HTMLOutput.appendln(aux.toHTMLString());
				}
			}



		}
		if (Config.annotateEnabled()){
			this.annotateFields();
			this.annotateExternalMethods();
			this.annotateExternalFields();
		}
	}

	/**
	 * computes the worst signature possible for method <code>method</code>
	 * @param method the method for which the signature will be computed
	 * @return the worst possible signture for method <code>method</code>
	 */
	TableLink computeWorstSignature(Method method) {
		int n = method.getArgumentTypes().length;
		int dec;
		if (method.isStatic())
			dec = 4;
		else
			dec = 5;
		TableLink b = new TableLink(n + dec, n + dec);

		for (int ii = 0; ii < n + dec; ii++) {
			int fromType = 0;
			int fromPart = 0;
			if ((ii > 0 && ii < dec)
					|| (ii == 0 && method.getReturnType().getSignature().charAt(0) == '[')
					|| (ii >= dec && method.getArgumentTypes()[ii - dec] instanceof ReferenceType && method
									.getArgumentTypes()[ii - dec].getSignature().charAt(0) == '[')) {
				fromPart = JvmType.PUBLIC_PART;
				fromType = Link.REFERENCE_LINK;
				// REFERENCE PUBLIC PART
			} else if ((ii == 0 && method.getReturnType() instanceof BasicType && method.getReturnType().getSignature().charAt(0)!='V')
					|| (ii >= dec && method.getArgumentTypes()[ii - dec] instanceof BasicType)) {
				fromPart = JvmType.PUBLIC_PART;
				fromType = Link.VALUE_LINK;
				// VALUE PUBLIC PART
			} else {
				fromPart = JvmType.WHOLE_PART;
				fromType = Link.REFERENCE_LINK;
				// WHOLE REFERENCE
			}
			for (int jj = 0; jj < n + dec; jj++) {

				int toType = 0;
				int toPart = 0;
				if ((jj > 0 && jj < dec)
						|| (jj == 0 && method.getReturnType().getSignature()
								.charAt(0) == '[')
								|| (jj >= dec
										&& method.getArgumentTypes()[jj - dec] instanceof ReferenceType && method
										.getArgumentTypes()[jj - dec].getSignature()
										.charAt(0) == '[')) {
					toPart = JvmType.PUBLIC_PART;
					toType = Link.REFERENCE_LINK;
					// REFERENCE PUBLIC PART
				} else if ((jj == 0 && method.getReturnType() instanceof BasicType  && method.getReturnType().getSignature().charAt(0)!='V')
						|| (jj >= dec && method.getArgumentTypes()[jj - dec] instanceof BasicType)) {
					toPart = JvmType.PUBLIC_PART;
					toType = Link.VALUE_LINK;
					// VALUE PUBLIC PART
				} else {
					toPart = JvmType.WHOLE_PART;
					toType = Link.REFERENCE_LINK;
					// WHOLE REFERENCE
				}

				b.set(ii, jj, GlobalAnalyser.linkFactory.createLink(fromPart,
						fromType, toPart, toType));

			}
		}
		if (method.getReturnType().getSignature().charAt(0) == 'V')
			for (int ii = 0; ii < n + dec; ii++) {
				b.put(ii, 0, null);
				b.put(0, ii, null);

			}
		return b;
	}

	/**
	 * annotates the <code>.class</code> file with security policies associated to class attributes.
	 *
	 */
	private void annotateFields() {
		Field[] field = this._class.getFields();
		for (int i = 0; i < field.length; i++) {
			if (policy.findAttribute(this._class.getClassName(), field[i]
			                                                           .getName()) == Policy.SECRET_ATTRIBUTE) // secret field
			{
				field[i].setAccessFlags(field[i].getAccessFlags()
						| Constants.ACC_SECRET);
			}
		}
	}


	/**
	 *
	 *
	 */
	private void annotateExternalFields(){

		byte[] byteStream = this.getExternalFieldsAsBytes();
		if (byteStream == null) {
			return;
		}

		Attribute[] orig = this._class.getAttributes();
		if (byteStream.length > 0) {
			ExternalFieldsAttribute linkAttr = new ExternalFieldsAttribute(externalFields_index,
					byteStream.length, byteStream, _class.getConstantPool());
			int j = 0;
			while (j < orig.length) {
				if (orig[j] instanceof ExternalFieldsAttribute)
					break;
				j++;
			}
			if (j < orig.length) {
				orig[j] = linkAttr;
				_class.setAttributes(orig);
			} else {
				Attribute[] new_attr = new Attribute[orig.length + 1];
				new_attr[0] = linkAttr;
				System.arraycopy(orig, 0, new_attr, 1, orig.length);
				_class.setAttributes(new_attr);
			}
		} else {
			//suppress the attribute
			int j = 0;
			while (j < orig.length) {
				if (orig[j] instanceof ExternalFieldsAttribute) {
					break;
				}
				j++;
			}
			if (j < orig.length) {
				Attribute[] new_attr = new Attribute[orig.length - 1];
				if (j > 0)
					System.arraycopy(orig, 0, new_attr, 0, j);
				if (j < orig.length - 1)
					System.arraycopy(orig, j + 1, new_attr, j, orig.length - j
							- 1);
				_class.setAttributes(new_attr);
			}
		}
	}

	/**
	 * returns the EXTERNAL_FIELDS attibute
	 *
	 * length array_of_external_fields
	 *
	 * where the array has the form:
	 *
	 * index  security_level
	 *
	 * where
	 * 		index [2bytes] is the constant pool index
	 * 		scurity_level [1 byte] the policy of the field : one of the values #Policy.SECRET_ATTRIBUTE or #Policy.PUBLIC_ATTRIBUTE
	 *
	 *
	 *
	 * @return
	 */
	byte[] getExternalFieldsAsBytes() {

		byte[] output = null;
		if(this.externalFields.size() < 1)
			return null;
		int length = SignatureConstants.TOTAL_LENGTH_SIZE;

		//compute length = noOfFields * ( siezOfIndex (2) + sizeOfPolicy (1) )
		length += this.externalFields.size() * ( SignatureConstants.CSTPOOL_SIZE + 1);

		output = new byte[length];
		int index = 0;
		byte[] arrSize = SignatureConstants.toByteArray(this.externalFields.size(),
				SignatureConstants.TOTAL_LENGTH_SIZE);
		System.arraycopy(arrSize, 0, output, index, arrSize.length);
		index += arrSize.length;


		//total number of distinct signatures to be embedded
		Enumeration<String> it = this.externalFields.keys();
		for(; it.hasMoreElements();) {
			String description = it.nextElement();

			String className = SignatureDictionary.keyClassName(description);
			String fieldName = SignatureDictionary.keyFieldName(description);
			int cstPoolIdx = this.externalFields.get(description).intValue();
			int securityLevel = policy.findAttribute( className, fieldName);

			//add constant pool index
			arrSize = SignatureConstants.toByteArray(cstPoolIdx, SignatureConstants.CSTPOOL_SIZE);
			System.arraycopy(arrSize, 0, output, index, arrSize.length);
			index += arrSize.length;

			//add security level
			output[index++] = (byte)securityLevel;
		}

		return output;
	}


	/**
	 *
	 *
	 */
	private void annotateExternalMethods(){

		byte[] byteStream = this.getExternalMethodsAsBytes();
		if (byteStream == null) {
			return;
		}

		Attribute[] orig = this._class.getAttributes();
		if (byteStream.length > 0) {
			ExternalMethodsAttribute linkAttr = new ExternalMethodsAttribute(externalMethods_index,
					byteStream.length, byteStream, _class.getConstantPool());
			int j = 0;
			while (j < orig.length) {
				if (orig[j] instanceof ExternalMethodsAttribute)
					break;
				j++;
			}
			if (j < orig.length) {
				orig[j] = linkAttr;
				_class.setAttributes(orig);
			} else {
				Attribute[] new_attr = new Attribute[orig.length + 1];
				new_attr[0] = linkAttr;
				System.arraycopy(orig, 0, new_attr, 1, orig.length);
				_class.setAttributes(new_attr);
			}
		} else {
			//suppress the attribute
			int j = 0;
			while (j < orig.length) {
				if (orig[j] instanceof ExternalMethodsAttribute) {
					break;
				}
				j++;
			}
			if (j < orig.length) {
				Attribute[] new_attr = new Attribute[orig.length - 1];
				if (j > 0)
					System.arraycopy(orig, 0, new_attr, 0, j);
				if (j < orig.length - 1)
					System.arraycopy(orig, j + 1, new_attr, j, orig.length - j
							- 1);
				_class.setAttributes(new_attr);
			}
		}
	}

	/**
	 * returns the EXTERNAL_METHODS attibute
	 *
	 * length array_of_external_methods
	 *
	 * where an entry in the array has the form:
	 *
	 * entry_size(in bytes) index  type  signature
	 *
	 * where
	 * 		entry_size [TOTAL_LENGHT_SIZE] is the number of bytes to follow for the entry (in bytes)
	 * 		index [2bytes] is the constant pool entry for the method
	 * 		type [1 byte] the type of signature (global or local)
	 * 		signature the method signature
	 *
	 * @return
	 */
	byte[] getExternalMethodsAsBytes() {
		byte[] output = null;
		if(this.externalMethods.size() < 1)
			return null;
		int length = SignatureConstants.TOTAL_LENGTH_SIZE;

		//compute length

		//total number of distinct signatures to be embedded
		int nSign = 0;
		Enumeration<String> it = this.externalMethods.keys();
		for(; it.hasMoreElements();) {
			String key = it.nextElement();
			int value = this.externalMethods.get(key).intValue();


			if((value & SignatureDictionary.EXACT_SIGN_IDX) != 0 ) {
				//we have the exact signature
				Signature sign = this.dictionary.findExactSignature(key);
				if(sign!=null) {
					nSign ++;
					FlowPartialSignature link = (FlowPartialSignature) sign
					.getPartialSignature(FlowAnnotation.class.getName());
					//System.out.println("FLW for :"+key+" = "+link);
					byte[] byteStream = ((FlowPartialSignatureData) link.getComponent())
					.getDataAsBytes();
					length += SignatureConstants.TOTAL_LENGTH_SIZE;
					length += SignatureConstants.CSTPOOL_SIZE;
					length += byteStream.length;
					//1 byte for type of signature (exact or global)
					length += 1;
				}
				else
					Printer.println("Exact signature for "+key+" not found in the dictionary!");

			}
			if((value & SignatureDictionary.GLOBAL_SIGN_IDX) != 0 ) {
				//we have the global signature
				Signature sign = this.dictionary.findGlobalSignature(key);
				if(sign!=null){
					nSign++;
					FlowPartialSignature link = (FlowPartialSignature) sign
					.getPartialSignature(FlowAnnotation.class.getName());
					byte[] byteStream = ((FlowPartialSignatureData) link.getComponent())
					.getDataAsBytes();
					length += SignatureConstants.TOTAL_LENGTH_SIZE;
					length += SignatureConstants.CSTPOOL_SIZE;
					length += byteStream.length;
					length += 1;
				}
				else
					Printer.println("Global signature for "+key+" not found in the dictionary!");

			}
		}

		output = new byte[length];
		int index = 0;
		byte[] arrSize = SignatureConstants.toByteArray(nSign,
				SignatureConstants.TOTAL_LENGTH_SIZE);
		System.arraycopy(arrSize, 0, output, index, arrSize.length);
		index += arrSize.length;

		//add data to
		it = this.externalMethods.keys();

		for(; it.hasMoreElements();) {
			String key = it.nextElement();
			int value = this.externalMethods.get(key).intValue();
			String className = SignatureDictionary.keyClassName(key);
			int cstPoolIdx = cpg.lookupMethodref( className,
					SignatureDictionary.keyOnlyMethodName(key),
					SignatureDictionary.keyDescription(key));
			if(cstPoolIdx == -1) {
				cstPoolIdx = cpg.lookupInterfaceMethodref( className,
						SignatureDictionary.keyOnlyMethodName(key),
						SignatureDictionary.keyDescription(key));

				if(cstPoolIdx == -1) {
					className = SignatureDictionary.keyJits(className);
					cstPoolIdx = cpg.lookupMethodref( className,
							SignatureDictionary.keyOnlyMethodName(key),
							SignatureDictionary.keyDescription(key));
					if(cstPoolIdx == -1)
						cstPoolIdx = cpg.lookupInterfaceMethodref( className,
								SignatureDictionary.keyOnlyMethodName(key),
								SignatureDictionary.keyDescription(key));
				}

			}

			if((value & SignatureDictionary.EXACT_SIGN_IDX) != 0 ) {
				//we have the exact signature
				Signature sign = this.dictionary.findExactSignature(key);
				if(sign!=null) {
					FlowPartialSignature link = (FlowPartialSignature) sign
					.getPartialSignature(FlowAnnotation.class.getName());
					byte[] byteStream = ((FlowPartialSignatureData) link.getComponent())
					.getDataAsBytes();
					arrSize = SignatureConstants.toByteArray(cstPoolIdx,
							SignatureConstants.CSTPOOL_SIZE);
					System.arraycopy(arrSize, 0, output, index, arrSize.length);
					index += arrSize.length;

					arrSize = SignatureConstants.toByteArray(byteStream.length,
							SignatureConstants.TOTAL_LENGTH_SIZE);
					System.arraycopy(arrSize, 0, output, index, arrSize.length);
					index += arrSize.length;

					output[index++] = SignatureDictionary.EXACT_SIGN_IDX;

					System.arraycopy(byteStream, 0, output, index, byteStream.length);
					index += byteStream.length;
				}
			}
			if((value & SignatureDictionary.GLOBAL_SIGN_IDX) != 0 ) {
				//we have the global signature
				Signature sign = this.dictionary.findGlobalSignature(key);
				if(sign!=null){
					FlowPartialSignature link = (FlowPartialSignature) sign
					.getPartialSignature(FlowAnnotation.class.getName());
					byte[] byteStream = ((FlowPartialSignatureData) link.getComponent())
					.getDataAsBytes();
					arrSize = SignatureConstants.toByteArray(cstPoolIdx,
							SignatureConstants.CSTPOOL_SIZE);
					System.arraycopy(arrSize, 0, output, index, arrSize.length);
					index += arrSize.length;

					arrSize = SignatureConstants.toByteArray(byteStream.length,
							SignatureConstants.TOTAL_LENGTH_SIZE);
					System.arraycopy(arrSize, 0, output, index, arrSize.length);
					index += arrSize.length;

					output[index++] = SignatureDictionary.GLOBAL_SIGN_IDX;

					System.arraycopy(byteStream, 0, output, index, byteStream.length);
					index += byteStream.length;
				}
			}
		}
		return output;
	}

	/**
	 * adds an external method invoked by one of the class' methods
	 * @param i
	 * @param typeSignature the type of signature of invoked method (global/exact)
	 */
	void addExternalMethod(InvokeInstruction i, int typeSignature) {
		int newVal;
		if( i.getClassName(cpg).equals(this._class.getClassName()))
			return;
		String key = SignatureDictionary.key(i, this.cpg);
		if(this.externalMethods.containsKey(key)){
			newVal = this.externalMethods.get(key).intValue() | typeSignature;
		}
		else
			newVal = typeSignature;
		this.externalMethods.put(key, new Integer(newVal));
	}

	/**
	 *
	 * @param i
	 */
	void addExternalField(FieldInstruction i) {
		if( i.getClassName(cpg).equals(this._class.getClassName()))
			return;
		String key = SignatureDictionary.key(i.getClassName(this.cpg), i.getFieldName(this.cpg), "");
		this.externalFields.put(key, new Integer(i.getIndex()));
	}

	/**
	 *
	 * @return
	 */
	public long memoryUsed() {
		return Runtime.getRuntime().totalMemory()
		- Runtime.getRuntime().freeMemory();
	}

}