Python: Clean up query a bit

Author: tausbn

python/ql/src/Security/CWE-295/MissingHostKeyValidation.ql

@@ -10,6 +10,7 @@
  */
 
 import python
+import semmle.python.dataflow.new.DataFlow
 import semmle.python.ApiGraphs
 
 private API::Node unsafe_paramiko_policy(string name) {
@@ -21,12 +22,12 @@ private API::Node paramikoSSHClientInstance() {
   result = API::moduleImport("paramiko").getMember("client").getMember("SSHClient").getReturn()
 }
 
-from CallNode call, ControlFlowNode arg, string name
+from DataFlow::CallCfgNode call, DataFlow::Node arg, string name
 where
-  call = paramikoSSHClientInstance().getMember("set_missing_host_key_policy").getACall().asCfgNode() and
+  call = paramikoSSHClientInstance().getMember("set_missing_host_key_policy").getACall() and
   arg = call.getAnArg() and
   (
-    arg = unsafe_paramiko_policy(name).getAUse().asCfgNode() or
-    arg = unsafe_paramiko_policy(name).getReturn().getAUse().asCfgNode()
+    arg = unsafe_paramiko_policy(name).getAUse() or
+    arg = unsafe_paramiko_policy(name).getReturn().getAUse()
   )
 select call, "Setting missing host key policy to " + name + " may be unsafe."

python/ql/src/semmle/python/dataflow/new/internal/DataFlowPublic.qll

@@ -193,6 +193,13 @@ class CallCfgNode extends CfgNode {
 
   /** Gets the data-flow node corresponding to the named argument of the call corresponding to this data-flow node */
   Node getArgByName(string name) { result.asCfgNode() = node.getArgByName(name) }
+
+  /** Gets the data-flow node corresponding to an argument of the call corresponding to this data-flow node */
+  Node getAnArg() {
+    exists(int n | result = this.getArg(n))
+    or
+    exists(string name | result = this.getArgByName(name))
+  }
 }
 
 /**