InvokeStaticInstrumenter.java

import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import soot.Body;
import soot.BodyTransformer;
import soot.PatchingChain;
import soot.SootClass;
import soot.SootField;
import soot.SootMethod;
import soot.Unit;
import soot.Value;
import soot.jimple.AssignStmt;
import soot.jimple.DefinitionStmt;
import soot.jimple.FieldRef;
import soot.jimple.IntConstant;
import soot.jimple.InvokeExpr;
import soot.jimple.InvokeStmt;
import soot.jimple.NumericConstant;


public class InvokeStaticInstrumenter extends BodyTransformer {
	//private static String holder="";		/*For debugging*/
	private static String log_msg="";
	private static final String noViolation = "NoViolation";
	public static int flag = 0;

	@Override
	protected void internalTransform(Body body, String arg1,Map<String, String> arg2) {
		//System.out.println(body.getMethod().getSignature().startsWith("<android."));
		if((body.getMethod().getSignature().startsWith("<android.support")
				//body beginning with android.support is likely to be Android library-needn't analyze
				|| body.getMethod().getSignature().contains("$")))
					//body beginning containing "$" has been temporarily generated by jimple, needn't analyze
			return;
		else{
			SootMethod aMethod = body.getMethod();
			//System.out.println(aMethod);
			/*
			 * Rule 1:
			 * Checking for modes-Not ECB, Provider not given etc
			 */
			if(!check_mode(aMethod).equalsIgnoreCase(noViolation)){
				log_msg = log_msg.concat(check_mode(aMethod));
				flag = 1;
			}
			
			/*
			 * Rule 2:
			 * Checking for static IVs, in case mode used is not ECB
			 */
			if(!check_staticIV(aMethod).equalsIgnoreCase(noViolation)){
				log_msg = log_msg.concat(check_staticIV(aMethod));
				flag = 1;
			}
			
			/*
			 * Rule 3:
			 * Check for constant encryption keys
			 */
			if(!check_constKey(aMethod).equalsIgnoreCase(noViolation)){
				log_msg = log_msg.concat(check_constKey(aMethod));
				flag =1;
			}
			
			/*
			 * Rule 4:
			 * Do not use constant salts for PBE.
			 */
			if(!check_saltgen(aMethod).equalsIgnoreCase(noViolation)){
				log_msg = log_msg.concat(check_saltgen(aMethod));
				flag = 1;
			}
			
			 /* 
			  * Rule 5:
			  * Do not use fewer than 1,000 iterations for PBE.
			  */
			if(!check_PBEiterations(aMethod).equalsIgnoreCase(noViolation)){
				log_msg = log_msg.concat(check_PBEiterations(aMethod));
				flag = 1;
			}
			
			/*
			 * Rule 6:
			 * Checking for SecureRandom seeding
			 */
			if(!check_secureRandom(aMethod).equalsIgnoreCase(noViolation)){
				log_msg = log_msg.concat(check_secureRandom(aMethod));
				flag = 1;
			}
				
			try {
				FileWriter writer = new FileWriter(AnalysisAPK.logFile);
			    writer.write(log_msg+"\n");
			    writer.flush();
			    writer.close();
			    
			} catch (IOException e) {
				e.printStackTrace();
			}
			//System.out.println(body.getMethod().getSignature());
		}

	}
	

	/*
	 * Checks the mode used for encryption/decryption processes
	 * @args: SootMethod
	 * @return: String Message corresponding to results found
	 */
	private static String check_mode(SootMethod aMethod) {
		String msg="NoViolation";		
		Body methodBody = aMethod.retrieveActiveBody();
		PatchingChain<Unit> units = methodBody.getUnits();
		InvokeExpr called;
		Value param;
		int flag = 0;
		for(Unit eachUnit:units){
			if (eachUnit instanceof AssignStmt){
				if(((AssignStmt) eachUnit).getRightOp() instanceof InvokeExpr){
					called = ((AssignStmt) eachUnit).getInvokeExpr();
					/*
					 * Looking for Cipher.getInstance(___) calls
					 */
					if(called.toString().contains("Cipher")&& called.toString().contains("getInstance")){
						/*
						 * Checks if no provider was given, simply issues a warning
						 * In this case the JCA framework looks for the first provider 
						 * instance that supports the given algorithm
						 */
						if(called.getArgCount()==1){
							flag = -1;
							msg = "\n"+aMethod.getSignature()+"\n\tCode 1: Warning: Provider not given in call to Cipher.getInstance()";
						}
					
						param = called.getArg(0);
						/*
						 * Checks if algorithm got defaulted to ECB mode
						 */
						if(param.toString().equalsIgnoreCase("\"AES\"") || param.toString().equalsIgnoreCase("\"DES\"") || param.toString().equalsIgnoreCase("\"DESEDE\"") || param.toString().equalsIgnoreCase("\"RSA\"")){
							if(flag==-1){
								msg = msg.concat("\n\tCode 1: Error: No mode specified, defaults to ECB!");
							}
							else{
								flag = 1;
								msg = "\n"+aMethod.getSignature()+"\n\tCode 1: Error: No mode specified defaults to ECB!";
							}							
						}
						/*
						 * Checks if the code intentionally uses ECB mode
						 */
						if(param.toString().contains("/") && param.toString().contains("/ECB/")){
							if(flag==-1){
								msg = msg.concat("\n\tCode 1: Error: ECB mode intentionally chosen!");
							}
							else{
								flag = 1;
								msg = "\n"+aMethod.getSignature()+"\n\tCode 1: Error: ECB mode intentionally chosen!";
							}
						}

						//Just a newline
						if(flag!=0){
							msg = msg.concat("\n");
						}
					}					
				}
			}
		}
		return msg;
	}

	
	/*
	 * Checks if the IVs used in the process are static
	 * @args: SootMethod
	 * @returns: String message corresponding to the analysis
	 */
	private static String check_staticIV(SootMethod aMethod) {
		String msg = "NoViolation";		
		Body methodBody = aMethod.retrieveActiveBody();
		int IVflag = -1;
		InvokeExpr expr;
		PatchingChain<Unit> units = methodBody.getUnits();
		for(Unit eachUnit:units){
			if(eachUnit instanceof InvokeStmt)
			{
				InvokeExpr ie = ((InvokeStmt) eachUnit).getInvokeExpr();
				SootMethod sm = ie.getMethod();
				String name= sm.getName();
				if((name.contains("init(int,java.security.Key,java.security.spec.AlgorithmParameterSpec)")))
						{
					IVflag = 0;
					break;
				}
			}
		}
		
		if(IVflag == 0){
			for(Unit eachUnit:units){
				if(eachUnit instanceof AssignStmt){
					if(((AssignStmt) eachUnit).getRightOp() instanceof InvokeExpr){
						expr = ((AssignStmt) eachUnit).getInvokeExpr();
						if(expr.toString().contains("IvParameterSpec") && expr.toString().contains("getIV()"))
								IVflag = 1;
					}
				}
			}
		}
		
		/*
		 * IVflag == 0 implies there was no call to getIV() method present in IVParameterSpec found
		 */
		else if(IVflag==0){
			msg = "\n"+aMethod.getSignature()+"\n\tCode 2: Error: IV generation is not done using proper library methods; Might as well be static";
			msg = msg.concat("\n");
		}

		return msg;
		
	}
	
	/*
	 * Checks if the keys being used for encryption are constant or randomly generated
	 * @args: SootMethod
	 * @returns: String message
	 */
	private static String check_constKey(SootMethod aMethod) {
		String msg="NoViolation";
		Body methodBody = aMethod.retrieveActiveBody();
		PatchingChain<Unit> units = methodBody.getUnits();
		InvokeExpr called;
		int flag = -1;
		for(Unit eachUnit:units){
			if(eachUnit instanceof AssignStmt){
				if(((AssignStmt) eachUnit).getRightOp() instanceof InvokeExpr){
					called = ((AssignStmt) eachUnit).getInvokeExpr();
					//If key was initialized using SecretKeyFactory
					if(called.toString().contains("keyFactory")){
						flag = 0;
						if(called.getMethod().getName().equalsIgnoreCase("generateSecret")){
							flag = 1;
						}
					}
					//If key generated using KeyGenerator
					else if(called.toString().contains("KeyGenerator")){
						flag = 0;
						if(called.getMethod().getName().equalsIgnoreCase("generateKey")){
							flag = 1;
						}
					}
				}
			}
		}
		if(flag==0){
			msg = "\n"+aMethod.getSignature()+"\n\tCode 3: Error: Key generation done without call to key generation methods, constant key usage suspected";
			msg = msg.concat("\n");
		}		
		return msg;
	}
	
	/*
	 * Checks if a constant salt is used or not for PBE
	 * @args: SootMethod
	 * @returns: String
	 */
	private static String check_saltgen(SootMethod aMethod){
		String msg = "NoViolation";
		ArrayList<String> approvedSources = new ArrayList<String>();
		approvedSources.add("<java.security.SecureRandom: void nextBytes(byte[])>");
		approvedSources.add("<java.security.SecureRandom: void setSeed(byte[])>");
		approvedSources.add("<java.lang.String: byte[] getBytes(java.lang.String)>");

		String pbeParam = "<javax.crypto.spec.PBEParameterSpec: void <init>(byte[],int)>";
		if(TaintAnalysis.myMap.containsKey(pbeParam)){
			for(String each:TaintAnalysis.myMap.get(pbeParam)){
				if(!approvedSources.contains(each))
					msg = "";
					break;
			}
		}
		if(msg.isEmpty()){
			msg =  "\n"+aMethod.getSignature()+"\n\tCode 4: Error: Salt generation not done properly";
		}
		return msg;
	}

	/*
	 * Checks if <1000 iterations are used for PBE
	 * @args: SootMethod
	 * @returns: String 
	 */
	private static String check_PBEiterations(SootMethod aMethod){
		String msg = "NoViolation";
		Body methodBody = aMethod.retrieveActiveBody();
		PatchingChain<Unit> units = methodBody.getUnits();

		int flag =-1;
		InvokeExpr expression;
		Value iterCount=null;
		for(Unit eachUnit:units){
			if(eachUnit instanceof InvokeStmt){
				expression = ((InvokeStmt) eachUnit).getInvokeExpr();
				if(expression.toString().contains("PBEParameterSpec") 
						&& expression.getMethod().getName().contains("<init>")){
					flag = 0;
					iterCount = expression.getArg(1);
					break;
				}
			}
		}
		
		if(iterCount instanceof NumericConstant){
			NumericConstant const1 = IntConstant.v(1000);
			NumericConstant result = ((NumericConstant) iterCount).greaterThan(const1);
			NumericConstant const2 = IntConstant.v(1);
			if(!result.equals(const2)){
				flag = 1;
			}
		}
		else{
			
			for(Unit eachUnit:units){
				if(eachUnit instanceof AssignStmt &&((AssignStmt) eachUnit).getLeftOp().equals(iterCount)){
					if(((AssignStmt) eachUnit).getRightOp() instanceof FieldRef){
						FieldRef fr = (FieldRef)((AssignStmt) eachUnit).getRightOp();
						
						SootClass aClass = fr.getFieldRef().declaringClass();
						
						if(!check_PBEIterCount(aClass,fr.getField()))
							flag = 1;
					}
				}
			}
		}		
		
		if(flag==1){
			msg = "\n"+aMethod.getSignature()+"\n\tCode 5: Error: Call to PBEParameterSpec runs for less than 1000 iterations.";
			msg = msg.concat("\n");
		}
		
		return msg;
		}
		
	
	/*
	 * Checks if SecureRandom calls have been used with the default seed, or overridden seed
	 * @args: SootMethod
	 * @returns: String
	 */
	private static String check_secureRandom(SootMethod aMethod) {
		String msg = "NoViolation";
		Body methodBody = aMethod.retrieveActiveBody();
		InvokeExpr expr;
		PatchingChain<Unit> units = methodBody.getUnits();
		for(Unit eachUnit:units){
			if(eachUnit instanceof InvokeStmt){
				expr = ((InvokeStmt) eachUnit).getInvokeExpr();
				if(expr.getMethod().toString().contains("SecureRandom")){
					if(expr.getMethod().getName().equalsIgnoreCase("<init>"))
						if(expr.getMethod().getParameterCount()==1){
							msg = "\n"+aMethod.getSignature()+"\n\tCode 6: Error: Seed given to SecureRandom object, overrides default algorithm seed";
						}
					if(((InvokeStmt) eachUnit).getInvokeExpr().getMethod().getName().equalsIgnoreCase("setSeed"))
						msg = "\n"+aMethod.getSignature()+"\n\tCode 6: Warning: Seeded SecureRandom object using setSeed(bytes[]) method, overrides default algorithm seed."+
					"Was this intended?";
				}
			}
		}
		if(!msg.equalsIgnoreCase("NoViolation")){
			msg.concat("\n");
		}
		return msg;
	}
	
	
	/*
	 * Returns true if the PBE Iteration parameter is greater than 1000, false otherwise
	 * Used by check_PBEiterations(SootMethod)
	 * @args: SootClass
	 * @returns: boolean
	 */
	private static boolean check_PBEIterCount(SootClass aClass, SootField theField){
		List<SootMethod> methodList = aClass.getMethods();
		NumericConstant value = IntConstant.v(0);
		for(SootMethod aMethod:methodList){
			Body mBody = aMethod.retrieveActiveBody();
			PatchingChain<Unit> unitChain = mBody.getUnits();
			for(Unit unit:unitChain){
				if(unit instanceof DefinitionStmt && ((DefinitionStmt) unit).getLeftOp() instanceof FieldRef){
					if(((DefinitionStmt) unit).getLeftOp().toString().contains(theField.getName()) 
							&& ((DefinitionStmt) unit).getRightOp() instanceof NumericConstant)
						value = ((NumericConstant)((DefinitionStmt)unit).getRightOp());
				}
			}
			NumericConstant const1 = IntConstant.v(1000);
			NumericConstant result = ((NumericConstant) value.greaterThanOrEqual(const1));
			NumericConstant const2 = IntConstant.v(1);
			if(result.equals(const2)){
				return true;
			}
			
		}
		return false;
	}
}