From 88feea34c691ccb2c547ba926c17ee68ffa63295 Mon Sep 17 00:00:00 2001
From: Christian Banse <oxisto@aybaze.com>
Date: Mon, 9 Dec 2024 23:00:52 +0100
Subject: [PATCH] Hotfix for python imports

Should make it better, but still not perfect probably.
---
 .../cpg/frontends/python/ExpressionHandler.kt | 16 +++++++++++++-
 .../frontends/python/PythonFrontendTest.kt    | 21 +++++++++++++++++++
 .../src/test/resources/python/import_test.py  |  5 +++++
 3 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 cpg-language-python/src/test/resources/python/import_test.py

diff --git a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt
index 897dfd617d..33d09cf81a 100644
--- a/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt
+++ b/cpg-language-python/src/main/kotlin/de/fraunhofer/aisec/cpg/frontends/python/ExpressionHandler.kt
@@ -480,7 +480,7 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
         var ref =
             if (isImport(base.name)) {
                 // Yes, it's an import, so we need to construct a reference with an FQN
-                newReference(base.name.fqn(node.attr), rawNode = node)
+                newReference(base.reconstructedImportName.fqn(node.attr), rawNode = node)
             } else {
                 newMemberExpression(name = node.attr, base = base, rawNode = node)
             }
@@ -596,3 +596,17 @@ class ExpressionHandler(frontend: PythonLanguageFrontend) :
         return lambda
     }
 }
+
+/**
+ * This utility function tries to reconstruct the name as if the expression was part of an imported
+ * symbol. This is needed because the [MemberExpression.name] includes the [MemberExpression.base]'s
+ * type instead of the name, and thus it might be "UNKNOWN".
+ */
+val Expression.reconstructedImportName: Name
+    get() {
+        return if (this is MemberExpression) {
+            this.base.reconstructedImportName.fqn(this.name.localName)
+        } else {
+            this.name
+        }
+    }
diff --git a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt
index e95489f9b4..6e767525a8 100644
--- a/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt
+++ b/cpg-language-python/src/test/kotlin/de/fraunhofer/aisec/cpg/frontends/python/PythonFrontendTest.kt
@@ -1590,6 +1590,27 @@ class PythonFrontendTest : BaseTest() {
         }
     }
 
+    @Test
+    fun testImportTest() {
+        val topLevel = Path.of("src", "test", "resources", "python")
+        val tu =
+            analyzeAndGetFirstTU(
+                listOf(topLevel.resolve("import_test.py").toFile()),
+                topLevel,
+                true
+            ) {
+                it.registerLanguage<PythonLanguage>()
+            }
+        assertNotNull(tu)
+
+        val refs = tu.refs
+        refs.forEach { assertIsNot<MemberExpression>(it) }
+        assertEquals(
+            setOf("a", "b", "pkg.module.foo", "another_module.foo"),
+            refs.map { it.name.toString() }.toSet()
+        )
+    }
+
     class PythonValueEvaluator : ValueEvaluator() {
         override fun computeBinaryOpEffect(
             lhsValue: Any?,
diff --git a/cpg-language-python/src/test/resources/python/import_test.py b/cpg-language-python/src/test/resources/python/import_test.py
new file mode 100644
index 0000000000..b8ba5b8148
--- /dev/null
+++ b/cpg-language-python/src/test/resources/python/import_test.py
@@ -0,0 +1,5 @@
+import pkg.module
+from pkg import another_module
+
+a = pkg.module.foo
+b = another_module.foo