Taint.qll

import python
import semmle.python.dataflow.TaintTracking
import semmle.python.security.strings.Untrusted

class SimpleSource extends TaintSource {
    SimpleSource() { this.(NameNode).getId() = "TAINTED_STRING" }

    override predicate isSourceOf(TaintKind kind) { kind instanceof ExternalStringKind }

    override string toString() { result = "taint source" }
}

class MySimpleSanitizer extends Sanitizer {
    MySimpleSanitizer() { this = "MySimpleSanitizer" }

    /**
     * The test `if is_safe(arg):` sanitizes `arg` on its `true` edge.
     *
     * Can't handle `if not is_safe(arg):` :\ that's why it's called MySimpleSanitizer
     */
    override predicate sanitizingEdge(TaintKind taint, PyEdgeRefinement test) {
        taint instanceof ExternalStringKind and
        exists(CallNode call | test.getTest() = call and test.getSense() = true |
            call = Value::named("test.is_safe").getACall() and
            test.getInput().getAUse() = call.getAnArg()
        )
    }
}

class MySanitizerHandlingNot extends Sanitizer {
    MySanitizerHandlingNot() { this = "MySanitizerHandlingNot" }

    /** The test `if is_safe(arg):` sanitizes `arg` on its `true` edge. */
    override predicate sanitizingEdge(TaintKind taint, PyEdgeRefinement test) {
        taint instanceof ExternalStringKind and
        clears_taint_on_true(test.getTest(), test.getSense(), test)
    }
}

/**
 * Helper predicate that recurses into any nesting of `not`
 *
 * To reduce the number of tuples this predicate holds for, we include the `PyEdgeRefinement` and
 * ensure that `test` is a part of this `PyEdgeRefinement` (instead of just taking the
 * `edge_refinement.getInput().getAUse()` part as a part of the predicate). Without including
 * `PyEdgeRefinement` as an argument *any* `CallNode c` to `test.is_safe` would be a result of
 * this predicate, since the tuple where `test = c` and `sense = true` would hold.
 */
private predicate clears_taint_on_true(
    ControlFlowNode test, boolean sense, PyEdgeRefinement edge_refinement
) {
    edge_refinement.getTest().getNode().(Expr).getASubExpression*() = test.getNode() and
    (
        test = Value::named("test.is_safe").getACall() and
        edge_refinement.getInput().getAUse() = test.(CallNode).getAnArg() and
        sense = true
        or
        test.(UnaryExprNode).getNode().getOp() instanceof Not and
        exists(ControlFlowNode nested_test |
            nested_test = test.(UnaryExprNode).getOperand() and
            clears_taint_on_true(nested_test, sense.booleanNot(), edge_refinement)
        )
    )
}

class TestConfig extends TaintTracking::Configuration {
    TestConfig() { this = "TestConfig" }

    override predicate isSanitizer(Sanitizer sanitizer) {
        sanitizer instanceof MySanitizerHandlingNot
    }

    override predicate isSource(TaintTracking::Source source) { source instanceof SimpleSource }

    override predicate isSink(TaintTracking::Sink sink) { none() }
}