9 files changed, 322 insertions, 0 deletions
diff --git a/src/main/kotlin/moe/nea/shot/ClassRef.kt b/src/main/kotlin/moe/nea/shot/ClassRef.kt
new file mode 100644
index 0000000..6fdd9d1
--- /dev/null
+++ b/src/main/kotlin/moe/nea/shot/ClassRef.kt
@@ -0,0 +1,17 @@
+package moe.nea.shot
+
+import java.io.Serializable
+
+
+data class ClassRef(val className: String) : Serializable {
+    val path: String get() = "$nudeJvmRef.class"
+    val jvmRef: String get() = "L$nudeJvmRef;"
+    val nudeJvmRef: String get() = className.replace(".", "/")
+
+    companion object {
+        fun fromPath(path: String): ClassRef? {
+            if (!path.endsWith(".class")) return null
+            return ClassRef(path.removeSuffix(".class").replace("/", "."))
+        }
+    }
+}
+\ No newline at end of file
diff --git a/src/main/kotlin/moe/nea/shot/ClassShots.kt b/src/main/kotlin/moe/nea/shot/ClassShots.kt
new file mode 100644
index 0000000..066ba82
--- /dev/null
+++ b/src/main/kotlin/moe/nea/shot/ClassShots.kt
@@ -0,0 +1,9 @@
+package moe.nea.shot
+
+import java.io.Serializable
+
+data class ClassShots(
+    val annotations: List<ClassRef>,
+    val methodShots: Map<MethodRef, MethodShots>,
+    val fieldShots: Map<FieldRef, FieldShots>,
+) : Serializable
+\ No newline at end of file
diff --git a/src/main/kotlin/moe/nea/shot/FieldRef.kt b/src/main/kotlin/moe/nea/shot/FieldRef.kt
new file mode 100644
index 0000000..007c79b
--- /dev/null
+++ b/src/main/kotlin/moe/nea/shot/FieldRef.kt
@@ -0,0 +1,5 @@
+package moe.nea.shot
+
+import java.io.Serializable
+
+data class FieldRef(val name: String) : Serializable
+\ No newline at end of file
diff --git a/src/main/kotlin/moe/nea/shot/FieldShots.kt b/src/main/kotlin/moe/nea/shot/FieldShots.kt
new file mode 100644
index 0000000..36217c5
--- /dev/null
+++ b/src/main/kotlin/moe/nea/shot/FieldShots.kt
@@ -0,0 +1,5 @@
+package moe.nea.shot
+
+import java.io.Serializable
+
+data class FieldShots(val annotations: List<ClassRef>) : Serializable
+\ No newline at end of file
diff --git a/src/main/kotlin/moe/nea/shot/InjectionApplicator.kt b/src/main/kotlin/moe/nea/shot/InjectionApplicator.kt
new file mode 100644
index 0000000..82b05a6
--- /dev/null
+++ b/src/main/kotlin/moe/nea/shot/InjectionApplicator.kt
@@ -0,0 +1,69 @@
+package moe.nea.shot
+
+import org.objectweb.asm.*
+
+class InjectionApplicator(
+    val injections: ClassShots,
+    writer: ClassWriter
+) : ClassVisitor(Opcodes.ASM9, writer) {
+
+    override fun visit(
+        version: Int,
+        access: Int,
+        name: String?,
+        signature: String?,
+        superName: String?,
+        interfaces: Array<out String>?
+    ) {
+        super.visit(version, access, name, signature, superName, interfaces)
+        for (annotation in injections.annotations) {
+            super.visitAnnotation(annotation.jvmRef, false).visitEnd()
+        }
+    }
+
+    override fun visitField(
+        access: Int,
+        name: String,
+        descriptor: String?,
+        signature: String?,
+        value: Any?
+    ): FieldVisitor {
+        val fieldVisitor = super.visitField(access, name, descriptor, signature, value)
+        val fieldShots = injections.fieldShots[FieldRef(name)] ?: return fieldVisitor
+        for (annotation in fieldShots.annotations) {
+            fieldVisitor.visitAnnotation(annotation.jvmRef, false)
+                .visitEnd()
+        }
+        return fieldVisitor
+    }
+
+    override fun visitMethod(
+        access: Int,
+        name: String,
+        descriptor: String,
+        signature: String?,
+        exceptions: Array<out String>?
+    ): MethodVisitor {
+        val parentVisitor = super.visitMethod(access, name, descriptor, signature, exceptions)
+        val methodShots = injections.methodShots[MethodRef(name, descriptor.replaceAfter(")", ""))] ?: return parentVisitor
+        for (injection in methodShots.annotations) {
+            parentVisitor.visitAnnotation(injection.jvmRef, false)
+                .visitEnd()
+        }
+        var maxParameterCount = 0
+        for ((parameter, annotations) in methodShots.parameterAnnotations) {
+            maxParameterCount = maxOf(maxParameterCount, parameter + 1)
+            for (annotation in annotations) {
+                parentVisitor
+                    .visitParameterAnnotation(parameter, annotation.jvmRef, false)
+                    .visitEnd()
+            }
+        }
+        return object : MethodVisitor(Opcodes.ASM8, parentVisitor) {
+            override fun visitAnnotableParameterCount(parameterCount: Int, visible: Boolean) {
+                super.visitAnnotableParameterCount(
+                    if (visible) parameterCount else maxOf(parameterCount, maxParameterCount), visible)
+            }
+        }
+    }
+}
+\ No newline at end of file
diff --git a/src/main/kotlin/moe/nea/shot/MethodRef.kt b/src/main/kotlin/moe/nea/shot/MethodRef.kt
new file mode 100644
index 0000000..7b07efc
--- /dev/null
+++ b/src/main/kotlin/moe/nea/shot/MethodRef.kt
@@ -0,0 +1,5 @@
+package moe.nea.shot
+
+import java.io.Serializable
+
+data class MethodRef(val name: String, val argumentDescriptor: String) : Serializable
diff --git a/src/main/kotlin/moe/nea/shot/MethodShots.kt b/src/main/kotlin/moe/nea/shot/MethodShots.kt
new file mode 100644
index 0000000..59893e2
--- /dev/null
+++ b/src/main/kotlin/moe/nea/shot/MethodShots.kt
@@ -0,0 +1,8 @@
+package moe.nea.shot
+
+import java.io.Serializable
+
+data class MethodShots(
+    val annotations: List<ClassRef>,
+    val parameterAnnotations: Map<Int, List<ClassRef>>,
+) : Serializable
+\ No newline at end of file
diff --git a/src/main/kotlin/moe/nea/shot/ShotParser.kt b/src/main/kotlin/moe/nea/shot/ShotParser.kt
new file mode 100644
index 0000000..648baa4
--- /dev/null
+++ b/src/main/kotlin/moe/nea/shot/ShotParser.kt
@@ -0,0 +1,149 @@
+package moe.nea.shot
+
+import java.util.*
+
+class ShotParser {
+
+    private data class Builder(
+        val indentation: Int,
+        val handler: Handler
+    )
+
+    private interface Handler {
+        fun handleLine(line: String): Handler?
+    }
+
+    private inner class MethodInjectionBuilder : Handler {
+        val annotations = mutableListOf<ClassRef>()
+        val parameterAnnotations = mutableMapOf<Int, MutableList<ClassRef>>()
+        override fun handleLine(line: String): Handler? {
+            if (line.startsWith("annotate ")) {
+                annotations.add(ClassRef(line.substring("annotate ".length).trim()))
+                return null
+            }
+            if (line.startsWith("annotateParameter ")) {
+                val (idx, name) = line.substring("annotateParameter ".length).split(" ")
+                parameterAnnotations.getOrPut(idx.toInt()) { mutableListOf() }.add(ClassRef(name))
+                return null
+            }
+            error("Unknown directive")
+        }
+    }
+
+    private inner class FieldInjectionBuilder : Handler {
+        val annotations = mutableListOf<ClassRef>()
+        override fun handleLine(line: String): Handler? {
+            if (line.startsWith("annotate ")) {
+                annotations.add(ClassRef(line.substring("annotate ".length).trim()))
+                return null
+            }
+            error("Unknown directive")
+        }
+    }
+
+    private inner class ClassInjectionBuilder : Handler {
+        val methodHandlers = mutableMapOf<MethodRef, MethodInjectionBuilder>()
+        val fieldHandlers = mutableMapOf<FieldRef, FieldInjectionBuilder>()
+        val annotations = mutableListOf<ClassRef>()
+        override fun handleLine(line: String): Handler? {
+            if (line.endsWith(":")) {
+                val data = line.substring(0, line.length - 1).trim()
+                if (data.contains("(")) {
+                    return methodHandlers.getOrPut(parseMethod(data)) { MethodInjectionBuilder() }
+                }
+                if (" " !in data) {
+                    return fieldHandlers.getOrPut(FieldRef(data)) { FieldInjectionBuilder() }
+                }
+                error("Unknown condition")
+            }
+            if (line.startsWith("annotate ")) {
+                annotations.add(ClassRef(line.substring("annotate ".length).trim()))
+                return null
+            }
+            error("Unknown line $line")
+        }
+
+        private fun parseMethod(data: String): MethodRef {
+            require(data.endsWith(")"))
+            val name = data.substringBefore("(")
+            val parameterDescriptor = data.substringAfter("(").dropLast(1)
+            val parameters = parameterDescriptor.split(",").filter { it.isNotBlank() }.map {
+                mapTypeToDescriptor(false, it.trim())
+            }
+            return MethodRef(name, parameters.joinToString("", "(", ")"))
+        }
+    }
+
+    private fun mapTypeToDescriptor(allowVoid: Boolean, name: String): String {
+        require(" " !in name)
+        if (name.endsWith("[]")) {
+            return "[" + mapTypeToDescriptor(allowVoid, name.substring(0, name.length - 2))
+        }
+        return when (name) {
+            "void" -> if (allowVoid) "V" else error("Void not allowed")
+            "boolean" -> "Z"
+            "byte" -> "B"
+            "int" -> "I"
+            "double" -> "D"
+            "long" -> "J"
+            "short" -> "S"
+            "float" -> "F"
+            "char" -> "C"
+            else -> "L" + name.replace(".", "/") + ";"
+        }
+    }
+
+    private inner class Root : Handler {
+        val map = mutableMapOf<ClassRef, ClassInjectionBuilder>()
+        override fun handleLine(line: String): Handler? {
+            require(line.endsWith(":"))
+            val className = ClassRef(line.substring(0, line.length - 1).trim())
+            return map.getOrPut(className) { ClassInjectionBuilder() }
+        }
+    }
+
+    private val root = Root()
+    private val indentations = Stack<Builder>().also {
+        it.add(Builder(0, root))
+    }
+    private var nextHandler: Handler? = null
+
+    fun parse(lines: Iterable<String>): Map<ClassRef, ClassShots> {
+        lines.forEach { parseLine(it) }
+        return root.map.mapValues {
+            ClassShots(
+                it.value.annotations,
+                it.value.methodHandlers.mapValues {
+                    MethodShots(
+                        it.value.annotations,
+                        it.value.parameterAnnotations
+                    )
+                },
+                it.value.fieldHandlers.mapValues {
+                    FieldShots(
+                        it.value.annotations
+                    )
+                }
+            )
+        }
+    }
+
+    fun parseLine(line: String) {
+        val indentation = line.takeWhile { it == ' ' }.length
+
+        if (indentation > indentations.peek().indentation) {
+            indentations.push(Builder(indentation, nextHandler ?: error("Illegal increase in indentation")))
+            nextHandler = null
+        }
+        while (indentation < indentations.peek().indentation) {
+            indentations.pop()
+        }
+
+        val nonCommentedLine = line.substringBefore('#')
+        if (nonCommentedLine.isBlank()) {
+            return
+        }
+        nextHandler = indentations.peek().handler.handleLine(line.trim())
+    }
+
+}
+\ No newline at end of file
diff --git a/src/main/kotlin/moe/nea/shot/Shots.kt b/src/main/kotlin/moe/nea/shot/Shots.kt
new file mode 100644
index 0000000..b0956b4
--- /dev/null
+++ b/src/main/kotlin/moe/nea/shot/Shots.kt
@@ -0,0 +1,55 @@
+package moe.nea.shot
+
+import org.objectweb.asm.ClassReader
+import org.objectweb.asm.ClassWriter
+import java.io.Serializable
+import java.nio.file.Files
+import java.nio.file.Path
+import java.util.zip.ZipFile
+import java.util.zip.ZipOutputStream
+import kotlin.io.path.createDirectories
+import kotlin.io.path.exists
+
+data class Shots(
+    val injections: Map<ClassRef, ClassShots>,
+) : Serializable {
+
+    fun processEntry(toInject: ClassShots, bytes: ByteArray): ByteArray {
+        val classReader = ClassReader(bytes)
+        val classWriter = ClassWriter(0)
+        val visitor = InjectionApplicator(toInject, classWriter)
+        classReader.accept(visitor, 0)
+        return classWriter.toByteArray()
+    }
+
+    fun processZipFile(input: ZipFile, output: ZipOutputStream) {
+        for (entry in input.entries()) {
+            val classRef = ClassRef.fromPath(entry.name)
+            val toInject = classRef?.let(injections::get)
+            output.putNextEntry(entry)
+            if (toInject == null) {
+                input.getInputStream(entry).copyTo(output)
+            } else {
+                val bytes = input.getInputStream(entry).use { it.readBytes() }
+                val modifiedBytes = processEntry(toInject, bytes)
+                output.write(modifiedBytes)
+            }
+        }
+    }
+
+
+    /**
+     * Run on an extracted zip, (or alternatively, a jar file systems root path)
+     */
+    fun process(sourceFiles: Path, targetFiles: Path) {
+        for ((classRef, toInject) in injections.entries) {
+            val classSource = sourceFiles.resolve(classRef.path)
+            val classTarget = targetFiles.resolve(classRef.path)
+            if (!classSource.exists()) {
+                continue
+            }
+            classTarget.parent.createDirectories()
+            Files.write(classTarget, processEntry(toInject, Files.readAllBytes(classSource)))
+        }
+    }
+}
+\ No newline at end of file