feat(bytecode): init import calls

archguard · Mar 29, 2022 · 79720f2 · 79720f2
1 parent 0e5a244
commit 79720f2
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 9 deletions.
diff --git a/scan_bytecode/src/main/kotlin/org/archguard/scanner/bytecode/ByteCodeParser.kt b/scan_bytecode/src/main/kotlin/org/archguard/scanner/bytecode/ByteCodeParser.kt
@@ -53,7 +53,13 @@ class ByteCodeParser {
         logger.debug("ByteCodeParser parser: {}", file)
         val classNode = getAsmClasNode(file)
         val module = getModule(file.toPath())
-        return createDataStruct(classNode, module)
+        val ds = createDataStruct(classNode, module)
+
+        ds.Imports = importCollector.packImports().map {
+            CodeImport(Source = it.key, AsName = it.value)
+        }.toTypedArray()
+
+        return ds
     }
 
     @Throws(IOException::class)
@@ -92,7 +98,7 @@ class ByteCodeParser {
 
     private fun createDataStruct(node: ClassNode, module: CodeModule): CodeDataStruct {
         val ds = CodeDataStruct()
-        val names = importCollector.splitClassAndPackageName(Type.getObjectType(node.name).className)
+        val names = importCollector.splitPackageAndClassName(Type.getObjectType(node.name).className)
         ds.Package = names.first
         ds.NodeName = names.second
 
@@ -131,16 +137,22 @@ class ByteCodeParser {
 //            createAnnotation(it)
 //        }?.toTypedArray() ?: arrayOf()
 
+        val className = Type.getType(field.desc).className
+
+        importCollector.processClassName(className)
+
         return CodeField(
-            TypeType = Type.getType(field.desc).className,
+            TypeType = className,
             TypeValue = field.name,
             Modifiers = createModifiers(field.access, FIELD_ALLOWED, isInterface, FIELD_EXCLUDED)
         )
     }
 
     private fun createAnnotation(annotation: AnnotationNode): CodeAnnotation {
-        val name: String = Type.getType(annotation.desc).className
-        val codeAnnotation = CodeAnnotation(Name = name)
+        val className: String = Type.getType(annotation.desc).className
+        val codeAnnotation = CodeAnnotation(Name = className)
+
+        importCollector.processClassName(className)
 
         if (annotation.values != null) {
             val values: List<Any> = annotation.values
@@ -180,8 +192,9 @@ class ByteCodeParser {
     }
 
     private fun getParamsFromDesc(desc: String, parameters: MutableList<ParameterNode>): Array<CodeProperty> {
-        Type.getType(desc).argumentsAndReturnSizes
         return Type.getType(desc).argumentTypes.mapIndexed { index, it ->
+            importCollector.processClassName(it.className)
+
             CodeProperty(
                 TypeType = it.className,
                 TypeValue = parameters[index].name

diff --git a/scan_bytecode/src/main/kotlin/org/archguard/scanner/bytecode/ImportCollector.kt b/scan_bytecode/src/main/kotlin/org/archguard/scanner/bytecode/ImportCollector.kt
@@ -16,7 +16,7 @@ class ImportCollector {
 
     }
 
-    fun splitClassAndPackageName(fullName: String): Pair<String, String> {
+    fun splitPackageAndClassName(fullName: String): Pair<String, String> {
         val lastDot = fullName.lastIndexOf('.')
         var packageName = fullName
         var className = fullName
@@ -33,7 +33,17 @@ class ImportCollector {
         return ""
     }
 
-    private fun packImports(): List<String> {
-        return arrayListOf()
+    fun packImports(): MutableMap<String, String> {
+        return mapSimpleNames
+    }
+
+    fun processClassName(className: String) {
+        if (className.startsWith(JAVA_LANG_PACKAGE)) {
+            return
+        }
+
+        val names = splitPackageAndClassName(className)
+
+        this.mapSimpleNames[className] = names.second
     }
 }
diff --git a/scan_bytecode/src/test/kotlin/org/archguard/scanner/bytecode/ByteCodeParserTest.kt b/scan_bytecode/src/test/kotlin/org/archguard/scanner/bytecode/ByteCodeParserTest.kt
@@ -72,6 +72,8 @@ internal class ByteCodeParserTest {
         assertEquals(2, ds.Functions.size)
         assertEquals("args", ds.Functions[1].Parameters[0].TypeValue)
         assertEquals("java.lang.String[]", ds.Functions[1].Parameters[0].TypeType)
+
+//        assertEquals(2, ds.Imports.size)
     }
 
     @Test