refactor more python type-trackers to API-graphs

Author: erik-krogh

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

@@ -22,7 +22,7 @@ private module CryptographyModel {
      * Gets a predefined curve class from
      * `cryptography.hazmat.primitives.asymmetric.ec` with a specific key size (in bits).
      */
-    private API::Node predefinedCurveClass(int keySize) {
+    API::Node predefinedCurveClass(int keySize) {
       exists(string curveName |
         result =
           API::moduleImport("cryptography")
@@ -73,41 +73,6 @@ private module CryptographyModel {
         curveName = "BrainpoolP512R1" and keySize = 512
       )
     }
-
-    /** Gets a reference to a predefined curve class with a specific key size (in bits), as well as the origin of the class. */
-    private DataFlow::TypeTrackingNode curveClassWithKeySize(
-      DataFlow::TypeTracker t, int keySize, DataFlow::Node origin
-    ) {
-      t.start() and
-      result = predefinedCurveClass(keySize).getAnImmediateUse() and
-      origin = result
-      or
-      exists(DataFlow::TypeTracker t2 |
-        result = curveClassWithKeySize(t2, keySize, origin).track(t2, t)
-      )
-    }
-
-    /** Gets a reference to a predefined curve class with a specific key size (in bits), as well as the origin of the class. */
-    DataFlow::Node curveClassWithKeySize(int keySize, DataFlow::Node origin) {
-      curveClassWithKeySize(DataFlow::TypeTracker::end(), keySize, origin).flowsTo(result)
-    }
-
-    /** Gets a reference to a predefined curve class instance with a specific key size (in bits), as well as the origin of the class. */
-    private DataFlow::TypeTrackingNode curveClassInstanceWithKeySize(
-      DataFlow::TypeTracker t, int keySize, DataFlow::Node origin
-    ) {
-      t.start() and
-      result.(DataFlow::CallCfgNode).getFunction() = curveClassWithKeySize(keySize, origin)
-      or
-      exists(DataFlow::TypeTracker t2 |
-        result = curveClassInstanceWithKeySize(t2, keySize, origin).track(t2, t)
-      )
-    }
-
-    /** Gets a reference to a predefined curve class instance with a specific key size (in bits), as well as the origin of the class. */
-    DataFlow::Node curveClassInstanceWithKeySize(int keySize, DataFlow::Node origin) {
-      curveClassInstanceWithKeySize(DataFlow::TypeTracker::end(), keySize, origin).flowsTo(result)
-    }
   }
 
   // ---------------------------------------------------------------------------
@@ -179,9 +144,13 @@ private module CryptographyModel {
     DataFlow::Node getCurveArg() { result in [this.getArg(0), this.getArgByName("curve")] }
 
     override int getKeySizeWithOrigin(DataFlow::Node origin) {
-      this.getCurveArg() = Ecc::curveClassInstanceWithKeySize(result, origin)
-      or
-      this.getCurveArg() = Ecc::curveClassWithKeySize(result, origin)
+      exists(API::Node n |
+        n = Ecc::predefinedCurveClass(result) and origin = n.getAnImmediateUse()
+      |
+        this.getCurveArg() = n.getAUse()
+        or
+        this.getCurveArg() = n.getReturn().getAUse()
+      )
     }
 
     // Note: There is not really a key-size argument, since it's always specified by the curve.
@@ -202,9 +171,8 @@ private module CryptographyModel {
     }
 
     /** Gets a reference to a Cipher instance using algorithm with `algorithmName`. */
-    DataFlow::TypeTrackingNode cipherInstance(DataFlow::TypeTracker t, string algorithmName) {
-      t.start() and
-      exists(DataFlow::CallCfgNode call | result = call |
+    API::Node cipherInstance(string algorithmName) {
+      exists(API::CallNode call | result = call.getReturn() |
         call =
           API::moduleImport("cryptography")
               .getMember("hazmat")
@@ -216,47 +184,6 @@ private module CryptographyModel {
             call.getArg(0), call.getArgByName("algorithm")
           ]
       )
-      or
-      exists(DataFlow::TypeTracker t2 | result = cipherInstance(t2, algorithmName).track(t2, t))
-    }
-
-    /** Gets a reference to a Cipher instance using algorithm with `algorithmName`. */
-    DataFlow::Node cipherInstance(string algorithmName) {
-      cipherInstance(DataFlow::TypeTracker::end(), algorithmName).flowsTo(result)
-    }
-
-    /** Gets a reference to the encryptor of a Cipher instance using algorithm with `algorithmName`. */
-    DataFlow::TypeTrackingNode cipherEncryptor(DataFlow::TypeTracker t, string algorithmName) {
-      t.start() and
-      result.(DataFlow::MethodCallNode).calls(cipherInstance(algorithmName), "encryptor")
-      or
-      exists(DataFlow::TypeTracker t2 | result = cipherEncryptor(t2, algorithmName).track(t2, t))
-    }
-
-    /**
-     * Gets a reference to the encryptor of a Cipher instance using algorithm with `algorithmName`.
-     *
-     * You obtain an encryptor by using the `encryptor()` method on a Cipher instance.
-     */
-    DataFlow::Node cipherEncryptor(string algorithmName) {
-      cipherEncryptor(DataFlow::TypeTracker::end(), algorithmName).flowsTo(result)
-    }
-
-    /** Gets a reference to the dncryptor of a Cipher instance using algorithm with `algorithmName`. */
-    DataFlow::TypeTrackingNode cipherDecryptor(DataFlow::TypeTracker t, string algorithmName) {
-      t.start() and
-      result.(DataFlow::MethodCallNode).calls(cipherInstance(algorithmName), "decryptor")
-      or
-      exists(DataFlow::TypeTracker t2 | result = cipherDecryptor(t2, algorithmName).track(t2, t))
-    }
-
-    /**
-     * Gets a reference to the decryptor of a Cipher instance using algorithm with `algorithmName`.
-     *
-     * You obtain an decryptor by using the `decryptor()` method on a Cipher instance.
-     */
-    DataFlow::Node cipherDecryptor(string algorithmName) {
-      cipherDecryptor(DataFlow::TypeTracker::end(), algorithmName).flowsTo(result)
     }
 
     /**
@@ -267,11 +194,12 @@ private module CryptographyModel {
       string algorithmName;
 
       CryptographyGenericCipherOperation() {
-        exists(DataFlow::Node object, string method |
-          object in [cipherEncryptor(algorithmName), cipherDecryptor(algorithmName)] and
-          method in ["update", "update_into"] and
-          this.calls(object, method)
-        )
+        this =
+          cipherInstance(algorithmName)
+              .getMember(["decryptor", "encryptor"])
+              .getReturn()
+              .getMember(["update", "update_into"])
+              .getACall()
       }
 
       override Cryptography::CryptographicAlgorithm getAlgorithm() {
@@ -298,9 +226,8 @@ private module CryptographyModel {
     }
 
     /** Gets a reference to a Hash instance using algorithm with `algorithmName`. */
-    private DataFlow::TypeTrackingNode hashInstance(DataFlow::TypeTracker t, string algorithmName) {
-      t.start() and
-      exists(DataFlow::CallCfgNode call | result = call |
+    private API::Node hashInstance(string algorithmName) {
+      exists(API::CallNode call | result = call.getReturn() |
         call =
           API::moduleImport("cryptography")
               .getMember("hazmat")
@@ -312,13 +239,6 @@ private module CryptographyModel {
             call.getArg(0), call.getArgByName("algorithm")
           ]
       )
-      or
-      exists(DataFlow::TypeTracker t2 | result = hashInstance(t2, algorithmName).track(t2, t))
-    }
-
-    /** Gets a reference to a Hash instance using algorithm with `algorithmName`. */
-    DataFlow::Node hashInstance(string algorithmName) {
-      hashInstance(DataFlow::TypeTracker::end(), algorithmName).flowsTo(result)
     }
 
     /**
@@ -328,7 +248,9 @@ private module CryptographyModel {
       DataFlow::MethodCallNode {
       string algorithmName;
 
-      CryptographyGenericHashOperation() { this.calls(hashInstance(algorithmName), "update") }
+      CryptographyGenericHashOperation() {
+        this = hashInstance(algorithmName).getMember("update").getACall()
+      }
 
       override Cryptography::CryptographicAlgorithm getAlgorithm() {
         result.matchesName(algorithmName)

python/ql/src/experimental/semmle/python/frameworks/Xml.qll

@@ -119,7 +119,7 @@ private module SaxBasedParsing {
    *
    * See https://docs.python.org/3.10/library/xml.sax.reader.html#xml.sax.xmlreader.XMLReader.setFeature
    */
-  class SaxParserSetFeatureCall extends DataFlow::MethodCallNode {
+  class SaxParserSetFeatureCall extends API::CallNode, DataFlow::MethodCallNode {
     SaxParserSetFeatureCall() {
       this =
         API::moduleImport("xml")
@@ -132,27 +132,9 @@ private module SaxBasedParsing {
 
     // The keyword argument names does not match documentation. I checked (with Python
     // 3.9.5) that the names used here actually works.
-    DataFlow::Node getFeatureArg() { result in [this.getArg(0), this.getArgByName("name")] }
+    API::Node getFeatureArg() { result = this.getParameter(0, "name") }
 
-    DataFlow::Node getStateArg() { result in [this.getArg(1), this.getArgByName("state")] }
-  }
-
-  /** Gets a back-reference to the `setFeature` state argument `arg`. */
-  private DataFlow::TypeTrackingNode saxParserSetFeatureStateArgBacktracker(
-    DataFlow::TypeBackTracker t, DataFlow::Node arg
-  ) {
-    t.start() and
-    arg = any(SaxParserSetFeatureCall c).getStateArg() and
-    result = arg.getALocalSource()
-    or
-    exists(DataFlow::TypeBackTracker t2 |
-      result = saxParserSetFeatureStateArgBacktracker(t2, arg).backtrack(t2, t)
-    )
-  }
-
-  /** Gets a back-reference to the `setFeature` state argument `arg`. */
-  DataFlow::LocalSourceNode saxParserSetFeatureStateArgBacktracker(DataFlow::Node arg) {
-    result = saxParserSetFeatureStateArgBacktracker(DataFlow::TypeBackTracker::end(), arg)
+    API::Node getStateArg() { result = this.getParameter(1, "state") }
   }
 
   /**
@@ -163,16 +145,13 @@ private module SaxBasedParsing {
   private DataFlow::Node saxParserWithFeatureExternalGesTurnedOn(DataFlow::TypeTracker t) {
     t.start() and
     exists(SaxParserSetFeatureCall call |
-      call.getFeatureArg() =
+      call.getFeatureArg().getARhs() =
         API::moduleImport("xml")
             .getMember("sax")
             .getMember("handler")
             .getMember("feature_external_ges")
             .getAUse() and
-      saxParserSetFeatureStateArgBacktracker(call.getStateArg())
-          .asExpr()
-          .(BooleanLiteral)
-          .booleanValue() = true and
+      call.getStateArg().getAValueReachingRhs().asExpr().(BooleanLiteral).booleanValue() = true and
       result = call.getObject()
     )
     or
@@ -182,16 +161,13 @@ private module SaxBasedParsing {
     // take account of that we can set the feature to False, which makes the parser safe again
     not exists(SaxParserSetFeatureCall call |
       call.getObject() = result and
-      call.getFeatureArg() =
+      call.getFeatureArg().getARhs() =
         API::moduleImport("xml")
             .getMember("sax")
             .getMember("handler")
             .getMember("feature_external_ges")
             .getAUse() and
-      saxParserSetFeatureStateArgBacktracker(call.getStateArg())
-          .asExpr()
-          .(BooleanLiteral)
-          .booleanValue() = false
+      call.getStateArg().getAValueReachingRhs().asExpr().(BooleanLiteral).booleanValue() = false
     )
   }