support enum switches with obfuscated SwitchMaps

2 files changed, 426 insertions, 0 deletions

diff --git a/src/main/java/cuchaz/enigma/Deobfuscator.java b/src/main/java/cuchaz/enigma/Deobfuscator.java
index b2cecfe2..6ea1c40b 100644
--- a/src/main/java/cuchaz/enigma/Deobfuscator.java
+++ b/src/main/java/cuchaz/enigma/Deobfuscator.java
@@ -25,6 +25,7 @@ import com.strobel.decompiler.languages.java.JavaOutputVisitor;
 import com.strobel.decompiler.languages.java.ast.AstBuilder;
 import com.strobel.decompiler.languages.java.ast.CompilationUnit;
 import com.strobel.decompiler.languages.java.ast.InsertParenthesesVisitor;
+import com.strobel.decompiler.languages.java.ast.transforms.IAstTransform;
 import cuchaz.enigma.analysis.*;
 import cuchaz.enigma.bytecode.ClassProtectifier;
 import cuchaz.enigma.bytecode.ClassPublifier;
@@ -32,6 +33,7 @@ import cuchaz.enigma.mapping.*;
 import cuchaz.enigma.mapping.entry.*;
 import cuchaz.enigma.throwables.IllegalNameException;
 import cuchaz.enigma.utils.Utils;
+import oml.ast.transformers.ObfuscatedEnumSwitchRewriterTransform;
 import org.objectweb.asm.ClassWriter;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.tree.ClassNode;
@@ -200,6 +202,7 @@ public class Deobfuscator {
                 AstBuilder builder = new AstBuilder(context);
                 builder.addType(resolvedType);
                 builder.runTransformations(null);
+                runCustomTransforms(builder, context);
                 return builder.getCompilationUnit();
         }
 
@@ -656,6 +659,15 @@ public class Deobfuscator {
                         throw new Error("Unknown entry desc: " + obfEntry);
         }
 
+        public static void runCustomTransforms(AstBuilder builder, DecompilerContext context){
+                List<IAstTransform> transformers = Arrays.asList(
+                                new ObfuscatedEnumSwitchRewriterTransform(context)
+                );
+                for (IAstTransform transform : transformers){
+                        transform.run(builder.getCompilationUnit());
+                }
+        }
+
         public interface ProgressListener {
                 void init(int totalWork, String title);
 
diff --git a/src/main/java/oml/ast/transformers/ObfuscatedEnumSwitchRewriterTransform.java b/src/main/java/oml/ast/transformers/ObfuscatedEnumSwitchRewriterTransform.java
new file mode 100644
index 00000000..6005b7f7
--- /dev/null
+++ b/src/main/java/oml/ast/transformers/ObfuscatedEnumSwitchRewriterTransform.java
@@ -0,0 +1,414 @@
+/*
+ * Originally:
+ * EnumSwitchRewriterTransform.java
+ *
+ * Copyright (c) 2013 Mike Strobel
+ *
+ * This source code is based on Mono.Cecil from Jb Evain, Copyright (c) Jb Evain;
+ * and ILSpy/ICSharpCode from SharpDevelop, Copyright (c) AlphaSierraPapa.
+ *
+ * This source code is subject to terms and conditions of the Apache License, Version 2.0.
+ * A copy of the license can be found in the License.html file at the root of this distribution.
+ * By using this source code in any fashion, you are agreeing to be bound by the terms of the
+ * Apache License, Version 2.0.
+ *
+ * You must not remove this notice, or any other, from this software.
+ */
+
+package oml.ast.transformers;
+
+import com.strobel.assembler.metadata.BuiltinTypes;
+import com.strobel.assembler.metadata.FieldDefinition;
+import com.strobel.assembler.metadata.MethodDefinition;
+import com.strobel.assembler.metadata.TypeDefinition;
+import com.strobel.assembler.metadata.TypeReference;
+import com.strobel.core.SafeCloseable;
+import com.strobel.core.VerifyArgument;
+import com.strobel.decompiler.DecompilerContext;
+import com.strobel.decompiler.languages.java.ast.AssignmentExpression;
+import com.strobel.decompiler.languages.java.ast.AstBuilder;
+import com.strobel.decompiler.languages.java.ast.AstNode;
+import com.strobel.decompiler.languages.java.ast.CaseLabel;
+import com.strobel.decompiler.languages.java.ast.ContextTrackingVisitor;
+import com.strobel.decompiler.languages.java.ast.Expression;
+import com.strobel.decompiler.languages.java.ast.IdentifierExpression;
+import com.strobel.decompiler.languages.java.ast.IndexerExpression;
+import com.strobel.decompiler.languages.java.ast.InvocationExpression;
+import com.strobel.decompiler.languages.java.ast.Keys;
+import com.strobel.decompiler.languages.java.ast.MemberReferenceExpression;
+import com.strobel.decompiler.languages.java.ast.PrimitiveExpression;
+import com.strobel.decompiler.languages.java.ast.SwitchSection;
+import com.strobel.decompiler.languages.java.ast.SwitchStatement;
+import com.strobel.decompiler.languages.java.ast.TypeDeclaration;
+import com.strobel.decompiler.languages.java.ast.TypeReferenceExpression;
+import com.strobel.decompiler.languages.java.ast.transforms.IAstTransform;
+
+import java.util.ArrayList;
+import java.util.IdentityHashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Copy of {@link com.strobel.decompiler.languages.java.ast.transforms.EnumSwitchRewriterTransform} modified to:
+ * - Not rely on a field containing "$SwitchMap$" (Proguard strips it)
+ * - Ignore classes *with* SwitchMap$ names (so the original can handle it)
+ * - Ignores inner synthetics that are not package private
+ */
+@SuppressWarnings("Duplicates")
+public class ObfuscatedEnumSwitchRewriterTransform implements IAstTransform {
+    private final DecompilerContext _context;
+
+    public ObfuscatedEnumSwitchRewriterTransform(final DecompilerContext context) {
+        _context = VerifyArgument.notNull(context, "context");
+    }
+
+    @Override
+    public void run(final AstNode compilationUnit) {
+        compilationUnit.acceptVisitor(new Visitor(_context), null);
+    }
+
+    private final static class Visitor extends ContextTrackingVisitor<Void> {
+        private final static class SwitchMapInfo {
+            final String enclosingType;
+            final Map<String, List<SwitchStatement>> switches = new LinkedHashMap<>();
+            final Map<String, Map<Integer, Expression>> mappings = new LinkedHashMap<>();
+
+            TypeDeclaration enclosingTypeDeclaration;
+
+            SwitchMapInfo(final String enclosingType) {
+                this.enclosingType = enclosingType;
+            }
+        }
+
+        private final Map<String, SwitchMapInfo> _switchMaps = new LinkedHashMap<>();
+        private boolean _isSwitchMapWrapper;
+
+        protected Visitor(final DecompilerContext context) {
+            super(context);
+        }
+
+        @Override
+        public Void visitTypeDeclaration(final TypeDeclaration typeDeclaration, final Void p) {
+            final boolean oldIsSwitchMapWrapper = _isSwitchMapWrapper;
+            final TypeDefinition typeDefinition = typeDeclaration.getUserData(Keys.TYPE_DEFINITION);
+            final boolean isSwitchMapWrapper = isSwitchMapWrapper(typeDefinition);
+
+            if (isSwitchMapWrapper) {
+                final String internalName = typeDefinition.getInternalName();
+
+                SwitchMapInfo info = _switchMaps.get(internalName);
+
+                if (info == null) {
+                    _switchMaps.put(internalName, info = new SwitchMapInfo(internalName));
+                }
+
+                info.enclosingTypeDeclaration = typeDeclaration;
+            }
+
+            _isSwitchMapWrapper = isSwitchMapWrapper;
+
+            try {
+                super.visitTypeDeclaration(typeDeclaration, p);
+            }
+            finally {
+                _isSwitchMapWrapper = oldIsSwitchMapWrapper;
+            }
+
+            rewrite();
+
+            return null;
+        }
+
+        @Override
+        public Void visitSwitchStatement(final SwitchStatement node, final Void data) {
+            final Expression test = node.getExpression();
+
+            if (test instanceof IndexerExpression) {
+                final IndexerExpression indexer = (IndexerExpression) test;
+                final Expression array = indexer.getTarget();
+                final Expression argument = indexer.getArgument();
+
+                if (!(array instanceof MemberReferenceExpression)) {
+                    return super.visitSwitchStatement(node, data);
+                }
+
+                final MemberReferenceExpression arrayAccess = (MemberReferenceExpression) array;
+                final Expression arrayOwner = arrayAccess.getTarget();
+                final String mapName = arrayAccess.getMemberName();
+
+                if (mapName == null || mapName.startsWith("$SwitchMap$") || !(arrayOwner instanceof TypeReferenceExpression)) {
+                    return super.visitSwitchStatement(node, data);
+                }
+
+                final TypeReferenceExpression enclosingTypeExpression = (TypeReferenceExpression) arrayOwner;
+                final TypeReference enclosingType = enclosingTypeExpression.getType().getUserData(Keys.TYPE_REFERENCE);
+
+                if (!isSwitchMapWrapper(enclosingType) || !(argument instanceof InvocationExpression)) {
+                    return super.visitSwitchStatement(node, data);
+                }
+
+                final InvocationExpression invocation = (InvocationExpression) argument;
+                final Expression invocationTarget = invocation.getTarget();
+
+                if (!(invocationTarget instanceof MemberReferenceExpression)) {
+                    return super.visitSwitchStatement(node, data);
+                }
+
+                final MemberReferenceExpression memberReference = (MemberReferenceExpression) invocationTarget;
+
+                if (!"ordinal".equals(memberReference.getMemberName())) {
+                    return super.visitSwitchStatement(node, data);
+                }
+
+                final String enclosingTypeName = enclosingType.getInternalName();
+
+                SwitchMapInfo info = _switchMaps.get(enclosingTypeName);
+
+                if (info == null) {
+                    _switchMaps.put(enclosingTypeName, info = new SwitchMapInfo(enclosingTypeName));
+
+                    final TypeDefinition resolvedType = enclosingType.resolve();
+
+                    if (resolvedType != null) {
+                        AstBuilder astBuilder = context.getUserData(Keys.AST_BUILDER);
+
+                        if (astBuilder == null) {
+                            astBuilder = new AstBuilder(context);
+                        }
+
+                        try (final SafeCloseable importSuppression = astBuilder.suppressImports()) {
+                            final TypeDeclaration declaration = astBuilder.createType(resolvedType);
+
+                            declaration.acceptVisitor(this, data);
+                        }
+                    }
+                }
+
+                List<SwitchStatement> switches = info.switches.get(mapName);
+
+                if (switches == null) {
+                    info.switches.put(mapName, switches = new ArrayList<>());
+                }
+
+                switches.add(node);
+            }
+
+            return super.visitSwitchStatement(node, data);
+        }
+
+        @Override
+        public Void visitAssignmentExpression(final AssignmentExpression node, final Void data) {
+            final TypeDefinition currentType = context.getCurrentType();
+            final MethodDefinition currentMethod = context.getCurrentMethod();
+
+            if (_isSwitchMapWrapper &&
+                currentType != null &&
+                currentMethod != null &&
+                currentMethod.isTypeInitializer()) {
+
+                final Expression left = node.getLeft();
+                final Expression right = node.getRight();
+
+                if (left instanceof IndexerExpression &&
+                    right instanceof PrimitiveExpression) {
+
+                    String mapName = null;
+
+                    final Expression array = ((IndexerExpression) left).getTarget();
+                    final Expression argument = ((IndexerExpression) left).getArgument();
+
+                    if (array instanceof MemberReferenceExpression) {
+                        mapName = ((MemberReferenceExpression) array).getMemberName();
+                    }
+                    else if (array instanceof IdentifierExpression) {
+                        mapName = ((IdentifierExpression) array).getIdentifier();
+                    }
+
+                    if (mapName == null || mapName.startsWith("$SwitchMap$")) {
+                        return super.visitAssignmentExpression(node, data);
+                    }
+
+                    if (!(argument instanceof InvocationExpression)) {
+                        return super.visitAssignmentExpression(node, data);
+                    }
+
+                    final InvocationExpression invocation = (InvocationExpression) argument;
+                    final Expression invocationTarget = invocation.getTarget();
+
+                    if (!(invocationTarget instanceof MemberReferenceExpression)) {
+                        return super.visitAssignmentExpression(node, data);
+                    }
+
+                    final MemberReferenceExpression memberReference = (MemberReferenceExpression) invocationTarget;
+                    final Expression memberTarget = memberReference.getTarget();
+
+                    if (!(memberTarget instanceof MemberReferenceExpression) || !"ordinal".equals(memberReference.getMemberName())) {
+                        return super.visitAssignmentExpression(node, data);
+                    }
+
+                    final MemberReferenceExpression outerMemberReference = (MemberReferenceExpression) memberTarget;
+                    final Expression outerMemberTarget = outerMemberReference.getTarget();
+
+                    if (!(outerMemberTarget instanceof TypeReferenceExpression)) {
+                        return super.visitAssignmentExpression(node, data);
+                    }
+
+                    final String enclosingType = currentType.getInternalName();
+
+                    SwitchMapInfo info = _switchMaps.get(enclosingType);
+
+                    if (info == null) {
+                        _switchMaps.put(enclosingType, info = new SwitchMapInfo(enclosingType));
+
+                        AstBuilder astBuilder = context.getUserData(Keys.AST_BUILDER);
+
+                        if (astBuilder == null) {
+                            astBuilder = new AstBuilder(context);
+                        }
+
+                        info.enclosingTypeDeclaration = astBuilder.createType(currentType);
+                    }
+
+                    final PrimitiveExpression value = (PrimitiveExpression) right;
+
+                    assert value.getValue() instanceof Integer;
+
+                    Map<Integer, Expression> mapping = info.mappings.get(mapName);
+
+                    if (mapping == null) {
+                        info.mappings.put(mapName, mapping = new LinkedHashMap<>());
+                    }
+
+                    final IdentifierExpression enumValue = new IdentifierExpression( Expression.MYSTERY_OFFSET, outerMemberReference.getMemberName());
+
+                    enumValue.putUserData(Keys.MEMBER_REFERENCE, outerMemberReference.getUserData(Keys.MEMBER_REFERENCE));
+
+                    mapping.put(((Number) value.getValue()).intValue(), enumValue);
+                }
+            }
+
+            return super.visitAssignmentExpression(node, data);
+        }
+
+        private void rewrite() {
+            if (_switchMaps.isEmpty()) {
+                return;
+            }
+
+            for (final SwitchMapInfo info : _switchMaps.values()) {
+                rewrite(info);
+            }
+
+            //
+            // Remove switch map type wrappers that are no longer referenced.
+            //
+
+        outer:
+            for (final SwitchMapInfo info : _switchMaps.values()) {
+                for (final String mapName : info.switches.keySet()) {
+                    final List<SwitchStatement> switches = info.switches.get(mapName);
+
+                    if (switches != null && !switches.isEmpty()) {
+                        continue outer;
+                    }
+                }
+
+                final TypeDeclaration enclosingTypeDeclaration = info.enclosingTypeDeclaration;
+
+                if (enclosingTypeDeclaration != null) {
+                    enclosingTypeDeclaration.remove();
+                }
+            }
+        }
+
+        private void rewrite(final SwitchMapInfo info) {
+            if (info.switches.isEmpty()) {
+                return;
+            }
+
+            for (final String mapName : info.switches.keySet()) {
+                final List<SwitchStatement> switches = info.switches.get(mapName);
+                final Map<Integer, Expression> mappings = info.mappings.get(mapName);
+
+                if (switches != null && mappings != null) {
+                    for (int i = 0; i < switches.size(); i++) {
+                        if (rewriteSwitch(switches.get(i), mappings)) {
+                            switches.remove(i--);
+                        }
+                    }
+                }
+            }
+        }
+
+        private boolean rewriteSwitch(final SwitchStatement s, final Map<Integer, Expression> mappings) {
+            final Map<Expression, Expression> replacements = new IdentityHashMap<>();
+
+            for (final SwitchSection section : s.getSwitchSections()) {
+                for (final CaseLabel caseLabel : section.getCaseLabels()) {
+                    final Expression expression = caseLabel.getExpression();
+
+                    if (expression.isNull()) {
+                        continue;
+                    }
+
+                    if (expression instanceof PrimitiveExpression) {
+                        final Object value = ((PrimitiveExpression) expression).getValue();
+
+                        if (value instanceof Integer) {
+                            final Expression replacement = mappings.get(value);
+
+                            if (replacement != null) {
+                                replacements.put(expression, replacement);
+                                continue;
+                            }
+                        }
+                    }
+
+                    //
+                    // If we can't rewrite all cases, we abort.
+                    //
+
+                    return false;
+                }
+            }
+
+            final IndexerExpression indexer = (IndexerExpression) s.getExpression();
+            final InvocationExpression argument = (InvocationExpression) indexer.getArgument();
+            final MemberReferenceExpression memberReference = (MemberReferenceExpression) argument.getTarget();
+            final Expression newTest = memberReference.getTarget();
+
+            newTest.remove();
+            indexer.replaceWith(newTest);
+
+            for (final Map.Entry<Expression, Expression> entry : replacements.entrySet()) {
+                entry.getKey().replaceWith(entry.getValue().clone());
+            }
+
+            return true;
+        }
+
+        private static boolean isSwitchMapWrapper(final TypeReference type) {
+            if (type == null) {
+                return false;
+            }
+
+            final TypeDefinition definition = type instanceof TypeDefinition ? (TypeDefinition) type
+                                                                             : type.resolve();
+
+            if (definition == null || !definition.isSynthetic() || !definition.isInnerClass() || !definition.isPackagePrivate()) {
+                return false;
+            }
+
+            for (final FieldDefinition field : definition.getDeclaredFields()) {
+                if (!field.getName().startsWith("$SwitchMap$") &&
+                    BuiltinTypes.Integer.makeArrayType().equals(field.getFieldType())) {
+
+                    return true;
+                }
+            }
+
+            return false;
+        }
+    }
+}
\ No newline at end of file