fix(python): support with/for statement variables (#1608)

internal/languages/python/.snapshots/flow/TestFlow--different-line.yml → internal/languages/python/.snapshots/TestDatatypes--different-line.yml

@@ -1,7 +1,7 @@
 low:
     - rule:
         cwe_ids: []
-        id: rule_logger_test
+        id: datatypes_test
         title: ""
         description: ""
         documentation_url: ""
@@ -30,6 +30,6 @@ low:
                 end: 20
         content: ""
       parent_line_number: 3
-      fingerprint: c94602447d6771c00b72425485a6cf6c_0
-      old_fingerprint: c94602447d6771c00b72425485a6cf6c_0
+      fingerprint: d3079d939d16cfca99001c1eeb9eda3c_0
+      old_fingerprint: d3079d939d16cfca99001c1eeb9eda3c_0
 

internal/languages/python/.snapshots/flow/TestFlow--same-line.yml → internal/languages/python/.snapshots/TestDatatypes--same-line.yml

@@ -1,7 +1,7 @@
 low:
     - rule:
         cwe_ids: []
-        id: rule_logger_test
+        id: datatypes_test
         title: ""
         description: ""
         documentation_url: ""
@@ -30,11 +30,11 @@ low:
                 end: 25
         content: ""
       parent_line_number: 1
-      fingerprint: 4b26059938bf9c55dcda8d08bcf6a4bd_0
-      old_fingerprint: 4b26059938bf9c55dcda8d08bcf6a4bd_0
+      fingerprint: 19205ed1e11f9a2acd2a0f2ed6b1cd6c_0
+      old_fingerprint: 19205ed1e11f9a2acd2a0f2ed6b1cd6c_0
     - rule:
         cwe_ids: []
-        id: rule_logger_test
+        id: datatypes_test
         title: ""
         description: ""
         documentation_url: ""
@@ -63,6 +63,6 @@ low:
                 end: 27
         content: ""
       parent_line_number: 2
-      fingerprint: 4b26059938bf9c55dcda8d08bcf6a4bd_1
-      old_fingerprint: 4b26059938bf9c55dcda8d08bcf6a4bd_1
+      fingerprint: 19205ed1e11f9a2acd2a0f2ed6b1cd6c_1
+      old_fingerprint: 19205ed1e11f9a2acd2a0f2ed6b1cd6c_1
 

internal/languages/python/.snapshots/TestFlow--flow.yml

@@ -0,0 +1,58 @@
+high:
+    - rule:
+        cwe_ids:
+            - "42"
+        id: flow_test
+        title: Test dataflow and variables
+        description: Test dataflow and variables
+        documentation_url: ""
+      line_number: 3
+      full_filename: flow.py
+      filename: flow.py
+      source:
+        location:
+            start: 3
+            end: 3
+            column:
+                start: 9
+                end: 27
+      sink:
+        location:
+            start: 3
+            end: 3
+            column:
+                start: 9
+                end: 27
+        content: ""
+      parent_line_number: 3
+      fingerprint: 22039dd750c8bd604904ee9f5bc626f0_0
+      old_fingerprint: 22039dd750c8bd604904ee9f5bc626f0_0
+    - rule:
+        cwe_ids:
+            - "42"
+        id: flow_test
+        title: Test dataflow and variables
+        description: Test dataflow and variables
+        documentation_url: ""
+      line_number: 7
+      full_filename: flow.py
+      filename: flow.py
+      source:
+        location:
+            start: 7
+            end: 7
+            column:
+                start: 9
+                end: 27
+      sink:
+        location:
+            start: 7
+            end: 7
+            column:
+                start: 9
+                end: 27
+        content: ""
+      parent_line_number: 7
+      fingerprint: 22039dd750c8bd604904ee9f5bc626f0_1
+      old_fingerprint: 22039dd750c8bd604904ee9f5bc626f0_1
+

internal/languages/python/analyzer/analyzer.go

@@ -45,6 +45,10 @@ func (analyzer *analyzer) Analyze(node *sitter.Node, visitChildren func() error)
 		return analyzer.analyzeParameters(node, visitChildren)
 	case "keyword_argument":
 		return analyzer.analyzeKeywordArgument(node, visitChildren)
+	case "for_statement":
+		return analyzer.analyzeForStatement(node, visitChildren)
+	case "with_item":
+		return analyzer.analyzeWithItem(node, visitChildren)
 	case "while_statement", "try_statement", "if_statement": // statements don't have results
 		return visitChildren()
 	case "conditional_expression":
@@ -195,6 +199,35 @@ func (analyzer *analyzer) analyzeKeywordArgument(node *sitter.Node, visitChildre
 	return visitChildren()
 }
 
+func (analyzer *analyzer) analyzeForStatement(node *sitter.Node, visitChildren func() error) error {
+	left := node.ChildByFieldName("left")
+	right := node.ChildByFieldName("right")
+
+	analyzer.builder.Dataflow(left, right)
+	analyzer.scope.Declare(analyzer.builder.ContentFor(left), left)
+	analyzer.lookupVariable(right)
+
+	return visitChildren()
+}
+
+// with x, y as foo:
+func (analyzer *analyzer) analyzeWithItem(node *sitter.Node, visitChildren func() error) error {
+	value := node.ChildByFieldName("value")
+	analyzer.lookupVariable(value)
+
+	err := visitChildren()
+
+	if value.Type() == "as_pattern" {
+		aliasValue := value.NamedChild(0)
+		alias := value.ChildByFieldName("alias")
+		analyzer.lookupVariable(aliasValue)
+		analyzer.builder.Alias(alias, aliasValue)
+		analyzer.scope.Declare(analyzer.builder.ContentFor(alias), alias)
+	}
+
+	return err
+}
+
 // import x
 // import a.b
 // from z import x

internal/languages/python/python_test.go

@@ -7,26 +7,33 @@ import (
 	"github.com/bearer/bearer/internal/languages/testhelper"
 )
 
-//go:embed testdata/logger.yml
-var loggerRule []byte
+//go:embed testdata/datatypes_rule.yml
+var datatypesRule []byte
 
 //go:embed testdata/scope_rule.yml
 var scopeRule []byte
 
+//go:embed testdata/flow_rule.yml
+var flowRule []byte
+
 //go:embed testdata/import_rule.yml
 var importRule []byte
 
 //go:embed testdata/subscript_rule.yml
 var subscriptRule []byte
 
-func TestFlow(t *testing.T) {
-	testhelper.GetRunner(t, loggerRule, "python").RunTest(t, "./testdata/testcases/flow", ".snapshots/flow/")
+func TestDatatypes(t *testing.T) {
+	testhelper.GetRunner(t, datatypesRule, "python").RunTest(t, "./testdata/datatypes", ".snapshots/")
 }
 
 func TestScope(t *testing.T) {
 	testhelper.GetRunner(t, scopeRule, "python").RunTest(t, "./testdata/scope", ".snapshots/")
 }
 
+func TestFlow(t *testing.T) {
+	testhelper.GetRunner(t, flowRule, "python").RunTest(t, "./testdata/flow", ".snapshots/")
+}
+
 func TestImport(t *testing.T) {
 	testhelper.GetRunner(t, importRule, "python").RunTest(t, "./testdata/import", ".snapshots/")
 }

internal/languages/python/testdata/logger.yml → internal/languages/python/testdata/datatypes_rule.yml

@@ -10,4 +10,4 @@ patterns:
       - variable: DATA_TYPE
         detection: datatype
 metadata:
-  id: rule_logger_test
+  id: datatypes_test

internal/languages/python/testdata/flow/flow.py

@@ -0,0 +1,10 @@
+def with_statement():
+    with source() as value, other:
+        cursor_sink(value)
+
+def for_statement():
+    for value in source():
+        result_sink(value)
+        cursor_sink(value) # no match
+
+cursor_sink(value) # no match

internal/languages/python/testdata/flow_rule.yml

@@ -0,0 +1,24 @@
+languages:
+  - python
+patterns:
+  - pattern: cursor_sink($<SOURCE>)
+    filters:
+      - variable: SOURCE
+        detection: flow_test_source
+        scope: cursor
+  - pattern: result_sink($<SOURCE>)
+    filters:
+      - variable: SOURCE
+        detection: flow_test_source
+        scope: result
+auxiliary:
+  - id: flow_test_source
+    patterns:
+      - source()
+severity: high
+metadata:
+  description: Test dataflow and variables
+  remediation_message: Test dataflow and variables
+  cwe_id:
+    - 42
+  id: flow_test