import python

import semmle.python.dataflow.TaintTracking

private import LegacyPointsTo

class SimpleTest extends TaintKind {

SimpleTest() { this = "simple.test" }

}

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 SimpleTest }

}

class SimpleSource extends TaintSource {

SimpleSource() { this.(NameNode).getId() = "SOURCE" }

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

override string toString() { result = "simple.source" }

}

class SimpleSanitizer extends Sanitizer {

SimpleSanitizer() { this = "Simple sanitizer" }

override predicate sanitizingNode(TaintKind taint, ControlFlowNode node) {

node.(CallNode).getFunction().(NameNode).getId() = "SANITIZE" and

taint instanceof SimpleTest

}

override predicate sanitizingDefinition(TaintKind taint, EssaDefinition def) {

exists(CallNode call |

def.(ArgumentRefinement).getInput().getAUse() = call.getAnArg() and

call.getFunction().(NameNode).getId() = "SANITIZE"

) and

taint instanceof SimpleTest

}

}

class BasicCustomTaint extends TaintKind {

BasicCustomTaint() { this = "basic.custom" }

override TaintKind getTaintForFlowStep(ControlFlowNode fromnode, ControlFlowNode tonode) {

tonode.(CallNode).getAnArg() = fromnode and

tonode.(CallNode).getFunction().(NameNode).getId() = "TAINT_FROM_ARG" and

result = this

}

}

class BasicCustomSink extends TaintSink {

override string toString() { result = "Basic custom sink" }

BasicCustomSink() {

exists(CallNode call |

call.getFunction().(NameNode).getId() = "CUSTOM_SINK" and

this = call.getAnArg()

)

}

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

}

class BasicCustomSource extends TaintSource {

BasicCustomSource() { this.(NameNode).getId() = "CUSTOM_SOURCE" }

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

override string toString() { result = "basic.custom.source" }

}

class Rock extends TaintKind {

Rock() { this = "rock" }

override TaintKind getTaintOfMethodResult(string name) {

name = "prev" and result instanceof Scissors

}

predicate isSink(ControlFlowNode sink) {

exists(CallNode call |

call.getArg(0) = sink and

call.getFunction().(NameNode).getId() = "paper"

)

}

}

class Paper extends TaintKind {

Paper() { this = "paper" }

override TaintKind getTaintOfMethodResult(string name) {

name = "prev" and result instanceof Rock

}

predicate isSink(ControlFlowNode sink) {

exists(CallNode call |

call.getArg(0) = sink and

call.getFunction().(NameNode).getId() = "scissors"

)

}

}

class Scissors extends TaintKind {

Scissors() { this = "scissors" }

override TaintKind getTaintOfMethodResult(string name) {

name = "prev" and result instanceof Paper

}

predicate isSink(ControlFlowNode sink) {

exists(CallNode call |

call.getArg(0) = sink and

call.getFunction().(NameNode).getId() = "rock"

)

}

}

class RockPaperScissorSource extends TaintSource {

RockPaperScissorSource() {

exists(string name | this.(NameNode).getId() = name |

name = "ROCK" or name = "PAPER" or name = "SCISSORS"

)

}

override predicate isSourceOf(TaintKind kind) { kind = this.(NameNode).getId().toLowerCase() }

override string toString() { result = "rock.paper.scissors.source" }

}

private predicate function_param(string funcname, ControlFlowNode arg) {

exists(FunctionObject f |

f.getName() = funcname and

arg = f.getArgumentForCall(_, _)

)

}

class RockPaperScissorSink extends TaintSink {

RockPaperScissorSink() {

exists(string name | function_param(name, this) |

name = "rock" or name = "paper" or name = "scissors"

)

}

override predicate sinks(TaintKind taint) {

exists(string name | function_param(name, this) |

name = "paper" and taint = "rock"

or

name = "rock" and taint = "scissors"

or

name = "scissors" and taint = "paper"

)

}

override string toString() { result = "rock.paper.scissors.sink" }

}

class TaintCarrier extends TaintKind {

TaintCarrier() { this = "explicit.carrier" }

override TaintKind getTaintOfMethodResult(string name) {

name = "get_taint" and result instanceof SimpleTest

}

}

/* There is no sink for `TaintCarrier`. It is not "dangerous" in itself; it merely holds a `SimpleTest`. */

class TaintCarrierSource extends TaintSource {

TaintCarrierSource() { this.(NameNode).getId() = "TAINT_CARRIER_SOURCE" }

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

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

}

/* Some more realistic examples */

abstract class UserInput extends TaintKind {

bindingset[this]

UserInput() { any() }

}

class UserInputSource extends TaintSource {

UserInputSource() { this.(CallNode).getFunction().(NameNode).getId() = "user_input" }

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

override string toString() { result = "user.input.source" }

}

class SqlInjectionTaint extends UserInput {

SqlInjectionTaint() { this = "SQL injection" }

}

class CommandInjectionTaint extends UserInput {

CommandInjectionTaint() { this = "Command injection" }

}

class SqlSanitizer extends Sanitizer {

SqlSanitizer() { this = "SQL sanitizer" }

/** Holds if `test` shows value to be untainted with `taint` */

override predicate sanitizingEdge(TaintKind taint, PyEdgeRefinement test) {

exists(FunctionObject f, CallNode call |

f.getName() = "isEscapedSql" and

test.getTest() = call and

call.getAnArg() = test.getSourceVariable().getAUse() and

f.getACall() = call and

test.getSense() = true

) and

taint instanceof SqlInjectionTaint

}

}

class CommandSanitizer extends Sanitizer {

CommandSanitizer() { this = "Command sanitizer" }

/** Holds if `test` shows value to be untainted with `taint` */

override predicate sanitizingEdge(TaintKind taint, PyEdgeRefinement test) {

exists(FunctionObject f |

f.getName() = "isValidCommand" and

f.getACall().(CallNode).getAnArg() = test.getSourceVariable().getAUse() and

test.getSense() = true

) and

taint instanceof CommandInjectionTaint

}

}

class SqlQuery extends TaintSink {

SqlQuery() {

exists(CallNode call |

call.getFunction().(NameNode).getId() = "sql_query" and

call.getAnArg() = this

)

}

override string toString() { result = "SQL query" }

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

}

class OsCommand extends TaintSink {

OsCommand() {

exists(CallNode call |

call.getFunction().(NameNode).getId() = "os_command" and

call.getAnArg() = this

)

}

override string toString() { result = "OS command" }

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

}

class Falsey extends TaintKind {

Falsey() { this = "falsey" }

override boolean booleanValue() { result = false }

}

class FalseySource extends TaintSource {

FalseySource() { this.(NameNode).getId() = "FALSEY" }

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

override string toString() { result = "falsey.source" }

}

class TaintIterable extends TaintKind {

TaintIterable() { this = "iterable.simple" }

override TaintKind getTaintForIteration() { result instanceof SimpleTest }

}

class TaintIterableSource extends TaintSource {

TaintIterableSource() { this.(NameNode).getId() = "ITERABLE_SOURCE" }

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

}