From 55648ac4dec436b2ca92d58db68a5b2e11df28c1 Mon Sep 17 00:00:00 2001
From: jorgectf <jorgectf@github.com>
Date: Thu, 20 Jul 2023 15:34:54 +0200
Subject: [PATCH] Add `shlex.quote` as sanitizer

---
 ...UnsafeShellCommandConstructionCustomizations.qll | 13 +++++++++++++
 .../UnsafeShellCommandConstructionQuery.qll         |  3 ++-
 .../src/unsafe_shell_test.py                        |  4 ++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/python/ql/lib/semmle/python/security/dataflow/UnsafeShellCommandConstructionCustomizations.qll b/python/ql/lib/semmle/python/security/dataflow/UnsafeShellCommandConstructionCustomizations.qll
index 390c9738cc3..a39ffc27e87 100644
--- a/python/ql/lib/semmle/python/security/dataflow/UnsafeShellCommandConstructionCustomizations.qll
+++ b/python/ql/lib/semmle/python/security/dataflow/UnsafeShellCommandConstructionCustomizations.qll
@@ -9,6 +9,7 @@ private import semmle.python.dataflow.new.DataFlow
 private import semmle.python.dataflow.new.TaintTracking
 private import CommandInjectionCustomizations::CommandInjection as CommandInjection
 private import semmle.python.Concepts as Concepts
+private import semmle.python.ApiGraphs
 
 /**
  * Module containing sources, sinks, and sanitizers for shell command constructed from library input.
@@ -17,6 +18,9 @@ module UnsafeShellCommandConstruction {
   /** A source for shell command constructed from library input vulnerabilities. */
   abstract class Source extends DataFlow::Node { }
 
+  /** A sanitizer for shell command constructed from library input vulnerabilities. */
+  abstract class Sanitizer extends DataFlow::Node { }
+
   private import semmle.python.frameworks.Setuptools
 
   /** An input parameter to a gem seen as a source. */
@@ -156,4 +160,13 @@ module UnsafeShellCommandConstruction {
 
     override DataFlow::Node getStringConstruction() { result = formatCall }
   }
+
+  /**
+   * A call to `shlex.quote`, considered as a sanitizer.
+   */
+  class ShlexQuoteAsSanitizer extends Sanitizer, DataFlow::Node {
+    ShlexQuoteAsSanitizer() {
+      this = API::moduleImport("shlex").getMember("quote").getACall().getArg(0)
+    }
+  }
 }
diff --git a/python/ql/lib/semmle/python/security/dataflow/UnsafeShellCommandConstructionQuery.qll b/python/ql/lib/semmle/python/security/dataflow/UnsafeShellCommandConstructionQuery.qll
index 0d9ebb8a472..80781a97ac6 100644
--- a/python/ql/lib/semmle/python/security/dataflow/UnsafeShellCommandConstructionQuery.qll
+++ b/python/ql/lib/semmle/python/security/dataflow/UnsafeShellCommandConstructionQuery.qll
@@ -24,7 +24,8 @@ class Configuration extends TaintTracking::Configuration {
   override predicate isSink(DataFlow::Node sink) { sink instanceof Sink }
 
   override predicate isSanitizer(DataFlow::Node node) {
-    node instanceof CommandInjection::Sanitizer // using all sanitizers from `rb/command-injection`
+    node instanceof Sanitizer or
+    node instanceof CommandInjection::Sanitizer // using all sanitizers from `py/command-injection`
   }
 
   // override to require the path doesn't have unmatched return steps
diff --git a/python/ql/test/query-tests/Security/CWE-078-UnsafeShellCommandConstruction/src/unsafe_shell_test.py b/python/ql/test/query-tests/Security/CWE-078-UnsafeShellCommandConstruction/src/unsafe_shell_test.py
index 1b4bc708c45..f83caf31eba 100644
--- a/python/ql/test/query-tests/Security/CWE-078-UnsafeShellCommandConstruction/src/unsafe_shell_test.py
+++ b/python/ql/test/query-tests/Security/CWE-078-UnsafeShellCommandConstruction/src/unsafe_shell_test.py
@@ -1,9 +1,13 @@
 import os
 import subprocess
+import shlex
 
 def unsafe_shell_one(name):
     os.system("ping " + name) # $result=BAD
 
+    # shlex.quote sanitizer
+    os.system("ping " + shlex.quote(name)) # $result=OK
+
     # f-strings
     os.system(f"ping {name}") # $result=BAD