Python: Support EC keygen without class-instance for cryptography

I also added a new test to show off how what the origin ends up looking
like... I think it looks ok

Author: RasmusWL

python/ql/src/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 DataFlow::Node curveClassWithKeySize(int keySize) {
+    private API::Node predefinedCurveClass(int keySize) {
       exists(string curveName |
         result =
           API::moduleImport("cryptography")
@@ -31,7 +31,6 @@ private module CryptographyModel {
               .getMember("asymmetric")
               .getMember("ec")
               .getMember(curveName)
-              .getAUse()
       |
         // obtained by manually looking at source code in
         // https://github.com/pyca/cryptography/blob/cba69f1922803f4f29a3fde01741890d88b8e217/src/cryptography/hazmat/primitives/asymmetric/ec.py#L208-L300
@@ -75,13 +74,30 @@ private module CryptographyModel {
       )
     }
 
+    /** 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::LocalSourceNode 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::LocalSourceNode curveClassInstanceWithKeySize(
       DataFlow::TypeTracker t, int keySize, DataFlow::Node origin
     ) {
       t.start() and
-      result.(DataFlow::CallCfgNode).getFunction() = curveClassWithKeySize(keySize) and
-      origin = result
+      result.(DataFlow::CallCfgNode).getFunction() = curveClassWithKeySize(keySize, origin)
       or
       exists(DataFlow::TypeTracker t2 |
         result = curveClassInstanceWithKeySize(t2, keySize, origin).track(t2, t)
@@ -164,6 +180,8 @@ private module CryptographyModel {
 
     override int getKeySizeWithOrigin(DataFlow::Node origin) {
       this.getCurveArg() = Ecc::curveClassInstanceWithKeySize(result, origin)
+      or
+      this.getCurveArg() = Ecc::curveClassWithKeySize(result, origin)
     }
 
     // Note: There is not really a key-size argument, since it's always specified by the curve.

python/ql/test/library-tests/frameworks/cryptography/EcKeygenOrigin.expected

@@ -0,0 +1,5 @@
+| ec_keygen_origin.py:8:1:8:45 | ControlFlowNode for Attribute() | 384 | ec_keygen_origin.py:8:31:8:42 | ControlFlowNode for Attribute |
+| ec_keygen_origin.py:9:1:9:43 | ControlFlowNode for Attribute() | 384 | ec_keygen_origin.py:9:31:9:42 | ControlFlowNode for Attribute |
+| ec_keygen_origin.py:12:1:12:36 | ControlFlowNode for Attribute() | 384 | ec_keygen_origin.py:11:9:11:20 | ControlFlowNode for Attribute |
+| ec_keygen_origin.py:15:1:15:39 | ControlFlowNode for Attribute() | 384 | ec_keygen_origin.py:11:9:11:20 | ControlFlowNode for Attribute |
+| ec_keygen_origin.py:20:1:20:32 | ControlFlowNode for Attribute() | 384 | ec_keygen_origin.py:6:58:6:66 | ControlFlowNode for ImportMember |

python/ql/test/library-tests/frameworks/cryptography/EcKeygenOrigin.ql

@@ -0,0 +1,8 @@
+import semmle.python.dataflow.new.DataFlow
+import semmle.python.Concepts
+
+from Cryptography::PublicKey::KeyGeneration keyGen, int keySize, DataFlow::Node origin
+where
+  keyGen.getLocation().getFile().getShortName() = "ec_keygen_origin.py" and
+  keySize = keyGen.getKeySizeWithOrigin(origin)
+select keyGen, keySize, origin

python/ql/test/library-tests/frameworks/cryptography/ec_keygen_origin.py

@@ -0,0 +1,20 @@
+# Since key-size is not specified explicitly as an integer for the predefined
+# classes in the `cryptography.hazmat.primitives.asymmetric.ec` module, we need
+# special handling of the origin... this test is simply to show off how we handle this.
+
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives.asymmetric.ec import SECP384R1
+
+ec.generate_private_key(curve=ec.SECP384R1()) # $ PublicKeyGeneration keySize=384
+ec.generate_private_key(curve=ec.SECP384R1) # $ PublicKeyGeneration keySize=384
+
+alias = ec.SECP384R1
+ec.generate_private_key(curve=alias) # $ PublicKeyGeneration keySize=384
+
+instance = alias()
+ec.generate_private_key(curve=instance) # $ PublicKeyGeneration keySize=384
+
+
+x = SECP384R1
+y = x
+ec.generate_private_key(curve=y) # $ PublicKeyGeneration keySize=384

python/ql/test/library-tests/frameworks/cryptography/test_ec.py

@@ -6,7 +6,7 @@
 
 
 private_key = ec.generate_private_key(curve=ec.SECP384R1()) # $ PublicKeyGeneration keySize=384
-private_key = ec.generate_private_key(curve=ec.SECP384R1) # $ MISSING: PublicKeyGeneration keySize=384
+private_key = ec.generate_private_key(curve=ec.SECP384R1) # $ PublicKeyGeneration keySize=384
 public_key = private_key.public_key()
 
 HASH_ALGORITHM = hashes.SHA256()