Python: Proper threat-model handling for argparse

Author: RasmusWL

python/ql/lib/semmle/python/frameworks/Stdlib.model.yml

@@ -12,5 +12,13 @@ extensions:
       - ['sys', 'Member[argv]', 'commandargs']
       - ['sys', 'Member[orig_argv]', 'commandargs']
 
-      # TODO: argparse
+      # if no argument is given, the default is to use sys.argv[1:]
+      - ['argparse.ArgumentParser', 'Member[parse_args,parse_known_args].WithArity[0].ReturnValue', 'commandargs']
+  - addsTo:
+      pack: codeql/python-all
+      extensible: summaryModel
+    data:
+      - ['argparse.ArgumentParser', 'Member[parse_args,parse_known_args]', 'Argument[0,args:]', 'ReturnValue', 'taint']
+      # note: taint of attribute lookups is handled in QL
+
       # TODO: input / read from stdin

python/ql/lib/semmle/python/frameworks/Stdlib.qll

@@ -4989,6 +4989,26 @@ module StdlibPrivate {
 
     override string getKind() { result = Escaping::getHtmlKind() }
   }
+
+  // ---------------------------------------------------------------------------
+  // argparse
+  // ---------------------------------------------------------------------------
+  /**
+   * if result of `parse_args` is tainted (because it uses command-line arguments),
+   *    then the parsed values accesssed on any attribute lookup is also tainted.
+   */
+  private class ArgumentParserAnyAttributeStep extends TaintTracking::AdditionalTaintStep {
+    override predicate step(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
+      nodeFrom =
+        API::moduleImport("argparse")
+            .getMember("ArgumentParser")
+            .getReturn()
+            .getMember("parse_args")
+            .getReturn()
+            .getAValueReachableFromSource() and
+      nodeTo.(DataFlow::AttrRead).getObject() = nodeFrom
+    }
+  }
 }
 
 // ---------------------------------------------------------------------------

python/ql/test/library-tests/frameworks/stdlib/threat_models.py

@@ -30,14 +30,14 @@
 parser = argparse.ArgumentParser()
 parser.add_argument("foo")
 
-args = parser.parse_args() # $ MISSING: threatModelSource[commandargs]=parser.parse_args()
-ensure_tainted(args.foo) # $ MISSING: tainted
+args = parser.parse_args() # $ threatModelSource[commandargs]=parser.parse_args()
+ensure_tainted(args.foo) # $ tainted
 
 explicit_argv_parsing = parser.parse_args(sys.argv) # $ threatModelSource[commandargs]=sys.argv
-ensure_tainted(explicit_argv_parsing.foo) # $ MISSING: tainted
+ensure_tainted(explicit_argv_parsing.foo) # $ tainted
 
 fake_args = parser.parse_args(["<foo>"])
-ensure_not_tainted(fake_args.foo)
+ensure_not_tainted(fake_args.foo) # $ SPURIOUS: tainted
 
 ########################################
 # reading input from stdin