Python: Better IntegerLiteral tracking for weak crypto key

Author: RasmusWL

python/ql/src/semmle/python/Concepts.qll

@@ -562,6 +562,21 @@ module Cryptography {
 
     /** Provides classes for modeling new key-pair generation APIs. */
     module KeyGeneration {
+      /**
+       * A data-flow configuration for tracking integer literals.
+       */
+      private class IntegerLiteralTrackerConfiguration extends DataFlow::Configuration {
+        IntegerLiteralTrackerConfiguration() { this = "IntegerLiteralTrackerConfiguration" }
+
+        override predicate isSource(DataFlow::Node source) {
+          source = DataFlow::exprNode(any(IntegerLiteral size))
+        }
+
+        override predicate isSink(DataFlow::Node sink) {
+          sink = any(KeyGeneration::Range kg).getKeySizeArg()
+        }
+      }
+
       /**
        * A data-flow node that generates a new key-pair for use with public-key cryptography.
        *
@@ -580,8 +595,9 @@ module Cryptography {
          * explains how we obtained this specific key size.
          */
         int getKeySizeWithOrigin(DataFlow::Node origin) {
-          exists(IntegerLiteral size | origin = DataFlow::exprNode(size) |
-            origin.(DataFlow::LocalSourceNode).flowsTo(this.getKeySizeArg()) and
+          exists(IntegerLiteral size, IntegerLiteralTrackerConfiguration config |
+            origin.asExpr() = size and
+            config.hasFlow(origin, this.getKeySizeArg()) and
             result = size.getValue()
           )
         }

python/ql/test/query-tests/Security/CWE-326/WeakCryptoKey.expected

@@ -6,3 +6,4 @@
 | weak_crypto.py:74:1:74:37 | ControlFlowNode for rsa_gen_key() | Creation of an RSA key uses $@ bits, which is below 2048 and considered breakable. | weak_crypto.py:12:12:12:15 | ControlFlowNode for IntegerLiteral | 1024 |
 | weak_crypto.py:76:1:76:22 | ControlFlowNode for Attribute() | Creation of an DSA key uses $@ bits, which is below 2048 and considered breakable. | weak_crypto.py:16:12:16:15 | ControlFlowNode for IntegerLiteral | 1024 |
 | weak_crypto.py:77:1:77:22 | ControlFlowNode for Attribute() | Creation of an RSA key uses $@ bits, which is below 2048 and considered breakable. | weak_crypto.py:12:12:12:15 | ControlFlowNode for IntegerLiteral | 1024 |
+| weak_crypto.py:84:12:84:29 | ControlFlowNode for Attribute() | Creation of an RSA key uses $@ bits, which is below 2048 and considered breakable. | weak_crypto.py:12:12:12:15 | ControlFlowNode for IntegerLiteral | 1024 |