gistfile1.java

import java.util.Iterator;
import java.util.Map;

import soot.Body;
import soot.BodyTransformer;
import soot.Local;
import soot.PackManager;
import soot.PatchingChain;
import soot.RefType;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.Transform;
import soot.Unit;
import soot.jimple.AbstractStmtSwitch;
import soot.jimple.InvokeExpr;
import soot.jimple.InvokeStmt;
import soot.jimple.Jimple;
import soot.jimple.StringConstant;
import soot.options.Options;

import soot.*;
import soot.jimple.*;
import soot.util.*;

import java.util.*;

public class AndroidInstrument {

    public static void main(String[] args) {

	// prefer Android APK files// -src-prec apk
	Options.v().set_src_prec(Options.src_prec_apk);

	// output as APK, too//-f J
	Options.v().set_output_format(Options.output_format_dex);

	// resolve the PrintStream and System soot-classes
	Scene.v().addBasicClass("java.io.PrintStream", SootClass.SIGNATURES);
	Scene.v().addBasicClass("java.lang.System", SootClass.SIGNATURES);

	PackManager.v().getPack("jtp").add(new Transform("jtp.myInstrumenter", new BodyTransformer() {

	    @Override
	    protected void internalTransform(final Body body, String phaseName, @SuppressWarnings("rawtypes") Map options) {
		final PatchingChain<Unit> units = body.getUnits();

		// important to use snapshotIterator here
		for (Iterator<Unit> iter = units.snapshotIterator(); iter.hasNext();) {
		    final Unit u = iter.next();
		    u.apply(new AbstractStmtSwitch() {

			public void caseAssignStmt(AssignStmt stmt) {
			    // InvokeExpr invokeExpr = stmt.getInvokeExpr();
			    // if(invokeExpr.getMethod().getName().equals("onDraw"))
			    // {

			    Local tmpRef = addTmpRef(body);
			    Local tmpRef2 = addTmpRef2(body);
			    Local tmpRef3 = addTmpRef3(body);

			    Local tmpString = addTmpString(body);
			    Local tmpString2 = addTmpString2(body);
			    Local tmpString3 = addTmpString3(body);

			    // insert "tmpRef = java.lang.System.out;"
			    units.insertBefore(Jimple.v().newAssignStmt(tmpRef,
				    Jimple.v().newStaticFieldRef(Scene.v().getField("<java.lang.System: java.io.PrintStream out>").makeRef())), stmt);
			    units.insertBefore(Jimple.v().newAssignStmt(
				    tmpRef2,
				    Jimple.v().newStaticInvokeExpr(
					    Scene.v().getMethod("<java.lang.Thread: java.lang.Thread currentThread()>").makeRef())), stmt);

			    SootMethod toCall2 = Scene.v().getSootClass("java.lang.Thread")
				    .getMethod("java.lang.StackTraceElement[] getStackTrace()");

			    units.insertBefore(Jimple.v().newAssignStmt(tmpString, Jimple.v().newVirtualInvokeExpr(tmpRef2, toCall2.makeRef())), stmt);

			    units.insertBefore(Jimple.v().newAssignStmt(
				    tmpString2,
				    Jimple.v().newStaticInvokeExpr(
					    Scene.v().getMethod("<java.lang.reflect.Array: java.lang.Object get(java.lang.Object,int)>").makeRef(),
					    tmpString, IntConstant.v(2))), stmt);

			    SootMethod toCall3 = Scene.v().getSootClass("java.lang.StackTraceElement").getMethod("java.lang.String getMethodName()");
			  
			    SootMethod toCall = Scene.v().getSootClass("java.io.PrintStream").getMethod("void println(java.lang.Object)");
			    units.insertBefore(Jimple.v().newInvokeStmt(Jimple.v().newVirtualInvokeExpr(tmpRef, toCall.makeRef(), tmpString2)), stmt);
			    // check that we did not mess up the Jimple
			    body.validate();
			    // }
			}

			public void caseInvokeStmt(InvokeStmt stmt) {
			    InvokeExpr invokeExpr = stmt.getInvokeExpr();
			    // if(invokeExpr.getMethod().getName().equals("onDraw"))
			    // {

			    Local tmpRef = addTmpRef(body);
			    Local tmpRef2 = addTmpRef2(body);
			    Local tmpRef3 = addTmpRef3(body);

			    Local tmpString = addTmpString(body);
			    Local tmpString2 = addTmpString2(body);
			    Local tmpString3 = addTmpString3(body);

			    // insert "tmpRef = java.lang.System.out;"
			    units.insertBefore(Jimple.v().newAssignStmt(tmpRef,
				    Jimple.v().newStaticFieldRef(Scene.v().getField("<java.lang.System: java.io.PrintStream out>").makeRef())), stmt);
			    units.insertBefore(Jimple.v().newAssignStmt(
				    tmpRef2,
				    Jimple.v().newStaticInvokeExpr(
					    Scene.v().getMethod("<java.lang.Thread: java.lang.Thread currentThread()>").makeRef())), stmt);

			    SootMethod toCall2 = Scene.v().getSootClass("java.lang.Thread")
				    .getMethod("java.lang.StackTraceElement[] getStackTrace()");

			    units.insertBefore(Jimple.v().newAssignStmt(tmpString, Jimple.v().newVirtualInvokeExpr(tmpRef2, toCall2.makeRef())), stmt);

			    units.insertBefore(Jimple.v().newAssignStmt(
				    tmpString2,
				    Jimple.v().newStaticInvokeExpr(
					    Scene.v().getMethod("<java.lang.reflect.Array: java.lang.Object get(java.lang.Object,int)>").makeRef(),
					    tmpString, IntConstant.v(2))), stmt);

			    SootMethod toCall3 = Scene.v().getSootClass("java.lang.StackTraceElement").getMethod("java.lang.String getMethodName()");
			   
			    SootMethod toCall = Scene.v().getSootClass("java.io.PrintStream").getMethod("void println(java.lang.Object)");
			    units.insertBefore(Jimple.v().newInvokeStmt(Jimple.v().newVirtualInvokeExpr(tmpRef, toCall.makeRef(), tmpString2)), stmt);
			    // check that we did not mess up the Jimple
			    body.validate();
			    // }
			}

		    });
		}
	    }

	}));
	Scene.v().addBasicClass("java.io.PrintStream", SootClass.SIGNATURES);
	Scene.v().addBasicClass("java.lang.System", SootClass.SIGNATURES);
	Scene.v().addBasicClass("java.lang.Thread", SootClass.SIGNATURES);
	Scene.v().addBasicClass("java.lang.StackTraceElement", SootClass.SIGNATURES);
	Scene.v().addBasicClass("java.lang.reflect.Array", SootClass.SIGNATURES);
	soot.Main.main(args);
    }

    private static Local addTmpRef(Body body) {
	Local tmpRef = Jimple.v().newLocal("tmpRef", RefType.v("java.io.PrintStream"));
	body.getLocals().add(tmpRef);
	return tmpRef;
    }

    private static Local addTmpRef2(Body body) {
	Local tmpRef2 = Jimple.v().newLocal("tmpRef2", RefType.v("java.lang.Thread"));
	body.getLocals().add(tmpRef2);
	return tmpRef2;
    }

    private static Local addTmpRef3(Body body) {
	Local tmpRef3 = Jimple.v().newLocal("tmpRef3", RefType.v("java.lang.String"));
	body.getLocals().add(tmpRef3);
	return tmpRef3;
    }

    private static Local addTmpString(Body body) {
	Local tmpString = Jimple.v().newLocal("tmpString", ArrayType.v(RefType.v("java.lang.StackTraceElement"), 1));
	body.getLocals().add(tmpString);
	return tmpString;
    }

    private static Local addTmpString2(Body body) {
	Local tmpString2 = Jimple.v().newLocal("tmpString2", RefType.v("java.lang.StackTraceElement"));
	body.getLocals().add(tmpString2);
	return tmpString2;
    }

    private static Local addTmpString3(Body body) {
	Local tmpString3 = Jimple.v().newLocal("tmpString3", RefType.v("java.lang.String"));
	body.getLocals().add(tmpString3);
	return tmpString3;
    }
}