fix: handle identifiers in subscript

internal/languages/python/.snapshots/TestSubscript--subscript.yml

@@ -0,0 +1,142 @@
+high:
+    - rule:
+        cwe_ids:
+            - "42"
+        id: subscript_test
+        title: Test detection filter subscript statements
+        description: Test detection filter subscript statements
+        documentation_url: ""
+      line_number: 1
+      full_filename: subscript.py
+      filename: subscript.py
+      source:
+        location:
+            start: 1
+            end: 1
+            column:
+                start: 1
+                end: 13
+      sink:
+        location:
+            start: 1
+            end: 1
+            column:
+                start: 1
+                end: 13
+        content: ""
+      parent_line_number: 1
+      fingerprint: dbeafee632db3831065048aabd949c8b_0
+      old_fingerprint: dbeafee632db3831065048aabd949c8b_0
+    - rule:
+        cwe_ids:
+            - "42"
+        id: subscript_test
+        title: Test detection filter subscript statements
+        description: Test detection filter subscript statements
+        documentation_url: ""
+      line_number: 3
+      full_filename: subscript.py
+      filename: subscript.py
+      source:
+        location:
+            start: 3
+            end: 3
+            column:
+                start: 1
+                end: 19
+      sink:
+        location:
+            start: 3
+            end: 3
+            column:
+                start: 1
+                end: 19
+        content: ""
+      parent_line_number: 3
+      fingerprint: dbeafee632db3831065048aabd949c8b_1
+      old_fingerprint: dbeafee632db3831065048aabd949c8b_1
+    - rule:
+        cwe_ids:
+            - "42"
+        id: subscript_test
+        title: Test detection filter subscript statements
+        description: Test detection filter subscript statements
+        documentation_url: ""
+      line_number: 7
+      full_filename: subscript.py
+      filename: subscript.py
+      source:
+        location:
+            start: 7
+            end: 7
+            column:
+                start: 1
+                end: 7
+      sink:
+        location:
+            start: 7
+            end: 7
+            column:
+                start: 1
+                end: 7
+        content: ""
+      parent_line_number: 7
+      fingerprint: dbeafee632db3831065048aabd949c8b_2
+      old_fingerprint: dbeafee632db3831065048aabd949c8b_2
+    - rule:
+        cwe_ids:
+            - "42"
+        id: subscript_test
+        title: Test detection filter subscript statements
+        description: Test detection filter subscript statements
+        documentation_url: ""
+      line_number: 9
+      full_filename: subscript.py
+      filename: subscript.py
+      source:
+        location:
+            start: 9
+            end: 9
+            column:
+                start: 1
+                end: 13
+      sink:
+        location:
+            start: 9
+            end: 9
+            column:
+                start: 1
+                end: 13
+        content: ""
+      parent_line_number: 9
+      fingerprint: dbeafee632db3831065048aabd949c8b_3
+      old_fingerprint: dbeafee632db3831065048aabd949c8b_3
+    - rule:
+        cwe_ids:
+            - "42"
+        id: subscript_test
+        title: Test detection filter subscript statements
+        description: Test detection filter subscript statements
+        documentation_url: ""
+      line_number: 10
+      full_filename: subscript.py
+      filename: subscript.py
+      source:
+        location:
+            start: 10
+            end: 10
+            column:
+                start: 1
+                end: 13
+      sink:
+        location:
+            start: 10
+            end: 10
+            column:
+                start: 1
+                end: 13
+        content: ""
+      parent_line_number: 10
+      fingerprint: dbeafee632db3831065048aabd949c8b_4
+      old_fingerprint: dbeafee632db3831065048aabd949c8b_4
+

internal/languages/python/analyzer/analyzer.go

@@ -120,10 +120,20 @@ func (analyzer *analyzer) analyzeAttribute(node *sitter.Node, visitChildren func
 }
 
 // foo["bar"]
+// foo[x]
+// globals()[y]
 func (analyzer *analyzer) analyzeSubscript(node *sitter.Node, visitChildren func() error) error {
 	objectNode := node.ChildByFieldName("value")
-	analyzer.builder.Dataflow(node, objectNode)
-	analyzer.lookupVariable(objectNode)
+	if objectNode.Type() == "identifier" {
+		analyzer.builder.Dataflow(node, objectNode)
+		analyzer.lookupVariable(objectNode)
+	}
+
+	subscriptNode := node.ChildByFieldName("subscript")
+	if subscriptNode.Type() == "identifier" {
+		analyzer.builder.Dataflow(node, subscriptNode)
+		analyzer.lookupVariable(subscriptNode)
+	}
 
 	return visitChildren()
 }

internal/languages/python/python_test.go

@@ -16,6 +16,9 @@ var scopeRule []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/")
 }
@@ -27,3 +30,7 @@ func TestScope(t *testing.T) {
 func TestImport(t *testing.T) {
 	testhelper.GetRunner(t, importRule, "python").RunTest(t, "./testdata/import", ".snapshots/")
 }
+
+func TestSubscript(t *testing.T) {
+	testhelper.GetRunner(t, subscriptRule, "python").RunTest(t, "./testdata/subscript", ".snapshots/")
+}

internal/languages/python/testdata/subscript/subscript.py

@@ -0,0 +1,17 @@
+foo[input()]
+
+globals()[input()]
+
+x = input()
+
+foo[x]
+
+globals()[x]
+globals()[x]("something")
+
+# don't match the below
+
+foo[y]
+foo["world"]
+globals()[y]
+globals()["world"]

internal/languages/python/testdata/subscript_rule.yml

@@ -0,0 +1,24 @@
+languages:
+  - python
+patterns:
+  - pattern: foo[$<USER_INPUT>]
+    filters:
+      - variable: USER_INPUT
+        detection: subscript_test_user_input
+        scope: result
+  - pattern: globals()[$<USER_INPUT>]$<...>
+    filters:
+      - variable: USER_INPUT
+        detection: subscript_test_user_input
+        scope: result
+auxiliary:
+  - id: subscript_test_user_input
+    patterns:
+      - input()
+severity: high
+metadata:
+  description: Test detection filter subscript statements
+  remediation_message: Test detection filter subscript statements
+  cwe_id:
+    - 42
+  id: subscript_test