Pass context and scope directly to transformer. Add Transformer#scope…

…dClasses for easier usage
narumii · Oct 14, 2024 · 4d86c9b · 4d86c9b
1 parent 0841ae4
commit 4d86c9b
Show file tree

Hide file tree

Showing 44 changed files with 332 additions and 409 deletions.
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -91,8 +91,8 @@ import uwu.narumi.deobfuscator.api.transformer.Transformer;
 
 public class SomeTransformer extends Transformer {
   @Override
-  protected void transform(ClassWrapper scope, Context context) throws Exception {
-    context.classes(scope).forEach(classWrapper -> classWrapper.methods().forEach(methodNode -> {
+  protected void transform() throws Exception {
+    classes().forEach(classWrapper -> classWrapper.methods().forEach(methodNode -> {
       // Code here
     }));
   }
@@ -109,8 +109,8 @@ import java.util.Arrays;
 
 public class SomeTransformer extends Transformer {
   @Override
-  protected void transform(ClassWrapper scope, Context context) throws Exception {
-    context.classes(scope).forEach(classWrapper -> classWrapper.methods().forEach(methodNode -> {
+  protected void transform() throws Exception {
+    classes().forEach(classWrapper -> classWrapper.methods().forEach(methodNode -> {
 
       // Iterate over all LDC instructions in the method
       Arrays.stream(methodNode.instructions.toArray())
@@ -143,8 +143,8 @@ import java.util.Arrays;
 
 public class SomeTransformer extends Transformer {
   @Override
-  protected void transform(ClassWrapper scope, Context context) throws Exception {
-    context.classes(scope).forEach(classWrapper -> classWrapper.methods().forEach(methodNode -> {
+  protected void transform() throws Exception {
+    classes().forEach(classWrapper -> classWrapper.methods().forEach(methodNode -> {
       MethodContext methodContext = MethodContext.framed(classWrapper, methodNode);
 
       // Find all System.out.println calls and replace the string with "Bye, World!"
@@ -189,8 +189,8 @@ import uwu.narumi.deobfuscator.api.transformer.Transformer;
 
 public class SomeTransformer extends Transformer {
   @Override
-  protected void transform(ClassWrapper scope, Context context) throws Exception {
-    FramedInstructionsStream.of(scope, context)
+  protected void transform() throws Exception {
+    FramedInstructionsStream.of(this)
         .editInstructionsStream(stream -> stream.filter(insn -> insn.getOpcode() == INVOKEVIRTUAL)) // Match only INVOKEVIRTUAL instructions
         .forEach(insnContext -> {
           MethodInsnNode methodInsn = (MethodInsnNode) insnContext.insn();

diff --git a/deobfuscator-api/src/main/java/uwu/narumi/deobfuscator/api/context/Context.java b/deobfuscator-api/src/main/java/uwu/narumi/deobfuscator/api/context/Context.java
@@ -2,11 +2,11 @@
 
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.jetbrains.annotations.UnmodifiableView;
 import uwu.narumi.deobfuscator.api.asm.ClassWrapper;
 import uwu.narumi.deobfuscator.api.execution.SandBox;
 import uwu.narumi.deobfuscator.api.classpath.Classpath;
@@ -93,10 +93,11 @@ public Stream<ClassWrapper> stream(ClassWrapper scope) {
         .filter(classWrapper -> scope == null || classWrapper.name().equals(scope.name()));
   }
 
-  public List<ClassWrapper> classes(ClassWrapper scope) {
+  @UnmodifiableView
+  public List<ClassWrapper> scopedClasses(ClassWrapper scope) {
     return classes.values().stream()
         .filter(classWrapper -> scope == null || classWrapper.name().equals(scope.name()))
-        .collect(Collectors.toList());
+        .toList();
   }
 
   public Optional<ClassWrapper> get(String name) {

diff --git a/...scator-api/src/main/java/uwu/narumi/deobfuscator/api/helper/FramedInstructionsStream.java b/...scator-api/src/main/java/uwu/narumi/deobfuscator/api/helper/FramedInstructionsStream.java
@@ -7,6 +7,7 @@
 import uwu.narumi.deobfuscator.api.asm.InsnContext;
 import uwu.narumi.deobfuscator.api.asm.MethodContext;
 import uwu.narumi.deobfuscator.api.context.Context;
+import uwu.narumi.deobfuscator.api.transformer.Transformer;
 
 import java.util.Arrays;
 import java.util.function.Consumer;
@@ -18,6 +19,10 @@
  * iterate over classes and methods ASYNC, and instructions SYNC. This can really speed up computing frames for methods.
  */
 public class FramedInstructionsStream {
+  public static FramedInstructionsStream of(Transformer transformer) {
+    return of(transformer.scope(), transformer.context());
+  }
+
   public static FramedInstructionsStream of(ClassWrapper scope, Context context) {
     return new FramedInstructionsStream(scope, context);
   }
@@ -61,7 +66,7 @@ public FramedInstructionsStream forceSync() {
 
   public void forEach(Consumer<InsnContext> consumer) {
     // Iterate over classes in parallel
-    this.classesStreamModifier.apply(this.forceSync ? this.context.classes(this.scope).stream() : this.context.classes(this.scope).parallelStream())
+    this.classesStreamModifier.apply(this.forceSync ? this.context.scopedClasses(this.scope).stream() : this.context.scopedClasses(this.scope).parallelStream())
         // Iterate over methods in parallel
         .forEach(classWrapper -> this.methodsStreamModifier.apply(this.forceSync ? classWrapper.methods().stream() : classWrapper.methods().parallelStream())
             .forEach(methodNode -> {

diff --git a/...scator-api/src/main/java/uwu/narumi/deobfuscator/api/transformer/ComposedTransformer.java b/...scator-api/src/main/java/uwu/narumi/deobfuscator/api/transformer/ComposedTransformer.java
@@ -4,8 +4,6 @@
 import java.util.function.Supplier;
 
 import org.jetbrains.annotations.Nullable;
-import uwu.narumi.deobfuscator.api.asm.ClassWrapper;
-import uwu.narumi.deobfuscator.api.context.Context;
 
 public class ComposedTransformer extends Transformer {
 
@@ -23,9 +21,9 @@ public ComposedTransformer(boolean rerunOnChange, Supplier<Transformer>... trans
   }
 
   @Override
-  protected void transform(ClassWrapper scope, Context context) {
+  protected void transform() {
     transformers.forEach(transformerSupplier -> {
-      boolean changed = Transformer.transform(transformerSupplier, scope, context);
+      boolean changed = Transformer.transform(transformerSupplier, scope(), context());
       if (changed) {
         this.markChange();
       }

diff --git a/deobfuscator-api/src/main/java/uwu/narumi/deobfuscator/api/transformer/Transformer.java b/deobfuscator-api/src/main/java/uwu/narumi/deobfuscator/api/transformer/Transformer.java
@@ -2,7 +2,9 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.UnmodifiableView;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.tree.MethodNode;
 import org.objectweb.asm.tree.analysis.Analyzer;
@@ -14,12 +16,15 @@
 import uwu.narumi.deobfuscator.api.exception.TransformerException;
 import uwu.narumi.deobfuscator.api.helper.AsmHelper;
 
+import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 
 public abstract class Transformer extends AsmHelper implements Opcodes {
+  protected static final Logger LOGGER = LogManager.getLogger();
 
-  protected static final Logger LOGGER = LogManager.getLogger(Transformer.class);
+  private Context context = null;
+  private ClassWrapper scope = null;
 
   // Internal variables
   private boolean hasRan = false;
@@ -37,11 +42,8 @@ public boolean shouldRerunOnChange() {
 
   /**
    * Do the transformation. If you implement it you MUST use {@link Transformer#markChange()} somewhere
-   *
-   * @param scope You can choose the class transform or set it to null to transform all classes
-   * @param context The context
    */
-  protected abstract void transform(ClassWrapper scope, Context context) throws Exception;
+  protected abstract void transform() throws Exception;
 
   /**
    * Marks that transformer changed something. You MUST use it somewhere in your transformer.
@@ -62,6 +64,40 @@ public String name() {
     return this.getClass().getSimpleName();
   }
 
+  /**
+   * Get classes for processing
+   */
+  @UnmodifiableView
+  protected List<ClassWrapper> scopedClasses() {
+    return this.context.scopedClasses(this.scope);
+  }
+
+  private void ensureInitialized() {
+    if (this.context == null) {
+      throw new IllegalStateException("Transformer is not initialized");
+    }
+  }
+
+  @NotNull
+  public Context context() {
+    ensureInitialized();
+    return context;
+  }
+
+  @Nullable
+  public ClassWrapper scope() {
+    ensureInitialized();
+    return scope;
+  }
+
+  /**
+   * Init transformer
+   */
+  private void init(Context context, ClassWrapper scope) {
+    this.context = context;
+    this.scope = scope;
+  }
+
   /**
    * Run the transformer
    *
@@ -92,14 +128,16 @@ private static boolean transform(
       throw new IllegalArgumentException("transformerSupplier tried to reuse transformer instance. You must pass a new instance of transformer");
     }
 
-    LOGGER.info("-------------------------------------");
+    // Initialize transformer
+    transformer.init(context, scope);
 
+    LOGGER.info("-------------------------------------");
     LOGGER.info("Running {} transformer", transformer.name());
     long start = System.currentTimeMillis();
 
     // Run the transformer!
     try {
-      transformer.transform(scope, context);
+      transformer.transform();
     } catch (TransformerException e) {
       LOGGER.error("! {}: {}", transformer.name(), e.getMessage());
       return false;
@@ -141,7 +179,7 @@ private static boolean transform(
    * Verifies if the bytecode is valid
    */
   private static void verifyBytecode(@Nullable ClassWrapper scope, Context context) throws IllegalStateException {
-    for (ClassWrapper classWrapper : context.classes(scope)) {
+    for (ClassWrapper classWrapper : context.scopedClasses(scope)) {
       for (MethodNode methodNode : classWrapper.methods()) {
         Analyzer<BasicValue> analyzer = new Analyzer<>(new BasicVerifier());
         try {

diff --git a/...i/src/main/java/uwu/narumi/deobfuscator/api/transformer/VersionedComposedTransformer.java b/...i/src/main/java/uwu/narumi/deobfuscator/api/transformer/VersionedComposedTransformer.java
@@ -4,9 +4,6 @@
 import java.util.Map;
 import java.util.function.Supplier;
 
-import uwu.narumi.deobfuscator.api.asm.ClassWrapper;
-import uwu.narumi.deobfuscator.api.context.Context;
-
 public abstract class VersionedComposedTransformer extends Transformer {
 
   private final String version;
@@ -18,15 +15,15 @@ public VersionedComposedTransformer(String version) {
   private boolean changed = false;
 
   @Override
-  protected void transform(ClassWrapper scope, Context context) {
+  protected void transform() {
     Map<String, List<Supplier<Transformer>>> transformers = transformersByVersions();
     if (!transformers.containsKey(version)) {
       throw new IllegalArgumentException(String.format("Version '%s' not found!", version));
     }
 
     transformers
         .get(version)
-        .forEach(transformer -> changed |= Transformer.transform(transformer, scope, context));
+        .forEach(transformer -> changed |= Transformer.transform(transformer, scope(), context()));
 
     if (changed) {
       markChange();

diff --git a/...mpl/src/test/java/uwu/narumi/deobfuscator/transformer/TestSandboxSecurityTransformer.java b/...mpl/src/test/java/uwu/narumi/deobfuscator/transformer/TestSandboxSecurityTransformer.java
@@ -2,17 +2,15 @@
 
 import dev.xdark.ssvm.execution.VMException;
 import dev.xdark.ssvm.mirror.type.InstanceClass;
-import uwu.narumi.deobfuscator.api.asm.ClassWrapper;
-import uwu.narumi.deobfuscator.api.context.Context;
 import uwu.narumi.deobfuscator.api.execution.SandBox;
 import uwu.narumi.deobfuscator.api.transformer.Transformer;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class TestSandboxSecurityTransformer extends Transformer {
   @Override
-  protected void transform(ClassWrapper scope, Context context) throws Exception {
-    SandBox sandBox = context.getSandBox();
+  protected void transform() throws Exception {
+    SandBox sandBox = context().getSandBox();
     InstanceClass clazz = sandBox.getHelper().loadClass("sandbox.TestSandboxSecurity");
 
     assertThrows(VMException.class, () -> {

diff --git a/...main/java/uwu/narumi/deobfuscator/core/other/impl/allatori/AllatoriStringTransformer.java b/...main/java/uwu/narumi/deobfuscator/core/other/impl/allatori/AllatoriStringTransformer.java
@@ -4,8 +4,6 @@
 import org.objectweb.asm.tree.IntInsnNode;
 import org.objectweb.asm.tree.LdcInsnNode;
 import org.objectweb.asm.tree.MethodInsnNode;
-import uwu.narumi.deobfuscator.api.asm.ClassWrapper;
-import uwu.narumi.deobfuscator.api.context.Context;
 import uwu.narumi.deobfuscator.api.transformer.Transformer;
 
 import java.util.ArrayList;
@@ -27,8 +25,8 @@ public AllatoriStringTransformer(boolean strong) {
     /* Written by https://github.com/Lampadina17 | 06/08/2024 */
     /* use UniversalNumberTransformer before this transformer to decrypt keys */
     @Override
-    protected void transform(ClassWrapper scope, Context context) throws Exception {
-        context.classes(scope).forEach(classWrapper -> {
+    protected void transform() throws Exception {
+        scopedClasses().forEach(classWrapper -> {
             classWrapper.methods().forEach(methodNode -> {
 
                 AtomicBoolean isDecryptor = new AtomicBoolean(false);
@@ -67,7 +65,7 @@ protected void transform(ClassWrapper scope, Context context) throws Exception {
         });
 
         /* Decrypt all strings */
-        context.classes(scope).forEach(classWrapper -> {
+        scopedClasses().forEach(classWrapper -> {
             classWrapper.methods().forEach(methodNode -> {
                 Arrays.stream(methodNode.instructions.toArray()).forEach(node -> {
                     if (node instanceof LdcInsnNode ldc && ldc.cst instanceof String && node.getNext() instanceof MethodInsnNode next && next.getOpcode() == INVOKESTATIC) {
@@ -98,7 +96,7 @@ protected void transform(ClassWrapper scope, Context context) throws Exception {
                 });
             });
         });
-        LOGGER.info("Decrypted {} strings in {} classes", this.getChangesCount(), context.classes().size());
+        LOGGER.info("Decrypted {} strings in {} classes", this.getChangesCount(), scopedClasses().size());
     }
 
     public class DecryptionMethod {

diff --git a/...c/main/java/uwu/narumi/deobfuscator/core/other/impl/clean/AnnotationCleanTransformer.java b/...c/main/java/uwu/narumi/deobfuscator/core/other/impl/clean/AnnotationCleanTransformer.java
@@ -1,38 +1,27 @@
 package uwu.narumi.deobfuscator.core.other.impl.clean;
 
-import uwu.narumi.deobfuscator.api.asm.ClassWrapper;
-import uwu.narumi.deobfuscator.api.context.Context;
 import uwu.narumi.deobfuscator.api.transformer.Transformer;
 
 public class AnnotationCleanTransformer extends Transformer {
 
   @Override
-  protected void transform(ClassWrapper scope, Context context) throws Exception {
-    context
-        .classes(scope)
-        .forEach(
-            classWrapper -> {
-              classWrapper.classNode().invisibleAnnotations = null;
-              classWrapper.classNode().invisibleAnnotations = null;
+  protected void transform() throws Exception {
+    scopedClasses().forEach(classWrapper -> {
+      classWrapper.classNode().invisibleAnnotations = null;
+      classWrapper.classNode().invisibleAnnotations = null;
 
-              classWrapper
-                  .methods()
-                  .forEach(
-                      methodNode -> {
-                        methodNode.invisibleAnnotations = null;
-                        methodNode.visibleAnnotations = null;
-                        methodNode.invisibleParameterAnnotations = null;
-                        methodNode.visibleParameterAnnotations = null;
-                      });
+      classWrapper.methods().forEach(methodNode -> {
+        methodNode.invisibleAnnotations = null;
+        methodNode.visibleAnnotations = null;
+        methodNode.invisibleParameterAnnotations = null;
+        methodNode.visibleParameterAnnotations = null;
+      });
 
-              classWrapper
-                  .fields()
-                  .forEach(
-                      fieldNode -> {
-                        fieldNode.invisibleAnnotations = null;
-                        fieldNode.visibleAnnotations = null;
-                      });
-            });
+      classWrapper.fields().forEach(fieldNode -> {
+        fieldNode.invisibleAnnotations = null;
+        fieldNode.visibleAnnotations = null;
+      });
+    });
 
     // There is always a change
     markChange();

diff --git a/...in/java/uwu/narumi/deobfuscator/core/other/impl/clean/ClassDebugInfoCleanTransformer.java b/...in/java/uwu/narumi/deobfuscator/core/other/impl/clean/ClassDebugInfoCleanTransformer.java
@@ -1,20 +1,15 @@
 package uwu.narumi.deobfuscator.core.other.impl.clean;
 
-import uwu.narumi.deobfuscator.api.asm.ClassWrapper;
-import uwu.narumi.deobfuscator.api.context.Context;
 import uwu.narumi.deobfuscator.api.transformer.Transformer;
 
 public class ClassDebugInfoCleanTransformer extends Transformer {
 
   @Override
-  protected void transform(ClassWrapper scope, Context context) throws Exception {
-    context
-        .classes(scope)
-        .forEach(
-            classWrapper -> {
-              classWrapper.classNode().sourceDebug = null;
-              classWrapper.classNode().sourceFile = null;
-            });
+  protected void transform() throws Exception {
+    scopedClasses().forEach(classWrapper -> {
+      classWrapper.classNode().sourceDebug = null;
+      classWrapper.classNode().sourceFile = null;
+    });
 
     // There is always a change
     markChange();