codeql/python/ql/test/library-tests/taint/general/ParamSource.ql

import python
import semmle.python.security.TaintTracking

/* Standard library sink */
import semmle.python.security.injection.Command

class TestKind extends TaintKind {
    TestKind() {
        this = "test"
    }

}

class CustomSource extends TaintSource {

    CustomSource() {
      exists(Parameter p |
            p.asName().getId() = "arg" and
            this.(ControlFlowNode).getNode() = p
        )
    }

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

    override string toString() {
        result = "Source of untrusted input"
    }

}

class SimpleSink extends TaintSink {

    override string toString() { result = "Simple sink" }

    SimpleSink() {
        exists(CallNode call |
            call.getFunction().(NameNode).getId() = "SINK" and
            this = call.getAnArg()
        )
    }

    override predicate sinks(TaintKind taint) {
        taint instanceof TestKind
    }

}

from TaintSource src, TaintSink sink, TaintKind srckind, TaintKind sinkkind

where src.flowsToSink(srckind, sink) and sink.sinks(sinkkind)
select srckind, src.getLocation().toString(), sink.getLocation().getStartLine(), sink.(ControlFlowNode).getNode().toString(), sinkkind