package soot.jimple.infoflow.problems.rules;

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

import soot.Value;
import soot.jimple.DefinitionStmt;
import soot.jimple.Stmt;
import soot.jimple.infoflow.InfoflowManager;
import soot.jimple.infoflow.aliasing.Aliasing;
import soot.jimple.infoflow.data.Abstraction;
import soot.jimple.infoflow.data.AccessPath;
import soot.jimple.infoflow.problems.TaintPropagationResults;
import soot.jimple.infoflow.source.SourceInfo;
import soot.jimple.infoflow.util.ByReferenceBoolean;

/**
 * Rule to introduce unconditional taints at sources
 * 
 * @author Steven Arzt
 *
 */
public class SourcePropagationRule extends AbstractTaintPropagationRule {

	public SourcePropagationRule(InfoflowManager manager, Aliasing aliasing,
			Abstraction zeroValue, TaintPropagationResults results) {
		super(manager, aliasing, zeroValue, results);
	}

	private Collection<Abstraction> propagate(Abstraction d1,
			Abstraction source, Stmt stmt, ByReferenceBoolean killSource,
			ByReferenceBoolean killAll) {
		if (source == getZeroValue()) {
			// We never propagate zero facts onwards
			killSource.value = true;
			
			// Is this a source?
			final SourceInfo sourceInfo = getManager().getSourceSinkManager() != null
					? getManager().getSourceSinkManager().getSourceInfo(stmt, getManager().getICFG()) : null;
			if (sourceInfo != null && !sourceInfo.getAccessPaths().isEmpty()) {
				Set<Abstraction> res = new HashSet<>();
				Value leftOp = stmt instanceof DefinitionStmt ? ((DefinitionStmt) stmt).getLeftOp() : null;
				for (AccessPath ap : sourceInfo.getAccessPaths()) {
					Abstraction abs = new Abstraction(ap,
							stmt,
							sourceInfo.getUserData(),
							false,
							false);
					res.add(abs);
					
					// Compute the aliases
					if (leftOp != null)
						if (getAliasing().canHaveAliases(stmt, leftOp, abs))
							getAliasing().computeAliases(d1, stmt, leftOp,
									res, getManager().getICFG().getMethodOf(stmt), abs);
					
					// Set the corresponding call site
					if (stmt.containsInvokeExpr())
						abs.setCorrespondingCallSite(stmt);
				}
				return res;
			}
			if (killAll != null)
				killAll.value = true;
		}
		return null;
	}

	@Override
	public Collection<Abstraction> propagateNormalFlow(Abstraction d1,
			Abstraction source, Stmt stmt, ByReferenceBoolean killSource,
			ByReferenceBoolean killAll) {
		return propagate(d1, source, stmt, killSource, killAll);
	}

	@Override
	public Collection<Abstraction> propagateCallToReturnFlow(Abstraction d1,
			Abstraction source, Stmt stmt, ByReferenceBoolean killSource,
			ByReferenceBoolean killAll) {
		return propagate(d1, source, stmt, killSource, null);
	}

	@Override
	public Collection<Abstraction> propagateReturnFlow(
			Collection<Abstraction> callerD1s, Abstraction source, Stmt stmt,
			Stmt retSite, Stmt callSite, ByReferenceBoolean killAll) {
		return null;
	}

	@Override
	public Collection<Abstraction> propagateCallFlow(Abstraction d1,
			Abstraction source, Stmt stmt, ByReferenceBoolean killAll) {
		// Normally, we don't inspect source methods
		if (!getManager().getConfig().getInspectSources()
				&& getManager().getSourceSinkManager() != null) {
			final SourceInfo sourceInfo = getManager().getSourceSinkManager().getSourceInfo(
					stmt, getManager().getICFG());
			if (sourceInfo != null)
				killAll.value = true;
		}
		
		// By default, we don't inspect sinks either
		if (!getManager().getConfig().getInspectSinks()
				&& getManager().getSourceSinkManager() != null) {
			final boolean isSink = getManager().getSourceSinkManager().isSink(
					stmt, getManager().getICFG(), source.getAccessPath());
			if (isSink)
				killAll.value = true;
		}
		
		return null;
	}

}

@@ -136,16 +136,38 @@
 .Value;%0A
+import soot.ValueBox;%0A
 import s
@@ -1074,16 +1074,19 @@
 lAll) %7B%0A
+%09%09%0A
 %09%09if (so
@@ -1121,91 +1121,33 @@
 %09// 
-We never propagate zero facts onwards%0A%09%09%09killSource.value = true;%0A%09%09%09%0A%09%09%09// Is this
+Check whether this can be
  a s
@@ -1151,17 +1151,23 @@
 a source
-?
+ at all
 %0A%09%09%09fina
@@ -1331,16 +1331,122 @@
 : null;%0A
+%09%09%09%09%09%0A%09%09%09// We never propagate zero facts onwards%0A%09%09%09killSource.value = true;%0A%09%09%09%0A%09%09%09// Is this a source?%0A
 %09%09%09if (s
@@ -2280,16 +2280,599 @@
 ue;%0A%09%09%7D%0A
+%09%09else if (!getManager().getConfig().getInspectSources() && killAll != null) %7B%0A%09%09%09// Check whether this can be a source at all%0A%09%09%09final SourceInfo sourceInfo = getManager().getSourceSinkManager() != null%0A%09%09%09%09%09? getManager().getSourceSinkManager().getSourceInfo(stmt, getManager().getICFG()) : null;%0A%09%09%09%0A%09%09%09// If a non-zero values goes through a source, we kill it%0A%09%09%09if (sourceInfo != null) %7B%0A%09%09%09%09for (ValueBox vb : stmt.getUseBoxes()) %7B%0A%09%09%09%09%09if (getAliasing().mayAlias(source.getAccessPath(), vb.getValue()) != null) %7B%0A%09%09%09%09%09%09killAll.value = true;%0A%09%09%09%09%09%09break;%0A%09%09%09%09%09%7D%0A%09%09%09%09%7D%0A%09%09%09%7D%0A%09%09%7D%0A
 %09%09return

kill data flows that go through a source method