3 files changed, 88 insertions, 18 deletions

diff --git a/src/main/java/cuchaz/enigma/bytecode/ClassRenamer.java b/src/main/java/cuchaz/enigma/bytecode/ClassRenamer.java
index eb7e9a1..d49f13a 100644
--- a/src/main/java/cuchaz/enigma/bytecode/ClassRenamer.java
+++ b/src/main/java/cuchaz/enigma/bytecode/ClassRenamer.java
@@ -19,8 +19,9 @@ import java.util.Map;
 
 import cuchaz.enigma.mapping.ClassEntry;
 import cuchaz.enigma.mapping.ClassNameReplacer;
+import cuchaz.enigma.mapping.Mappings;
 import cuchaz.enigma.mapping.Translator;
-import javassist.CtClass;
+import javassist.*;
 import javassist.bytecode.*;
 import javassist.bytecode.SignatureAttribute.*;
 
@@ -70,6 +71,41 @@ public class ClassRenamer {
         }
     }
 
+    public static void applyModifier(Object obj, Mappings.EntryModifier modifier)
+    {
+        int mod = -1;
+        if (obj instanceof CtField)
+            mod = ((CtField) obj).getModifiers();
+        else if (obj instanceof CtBehavior)
+            mod = ((CtBehavior) obj).getModifiers();
+        else if (obj instanceof CtClass)
+            mod = ((CtClass) obj).getModifiers();
+
+        if (mod != -1)
+        {
+            switch (modifier)
+            {
+                case PRIVATE:
+                    mod = Modifier.setPrivate(mod);
+                    break;
+                case PROTECTED:
+                    mod = Modifier.setProtected(mod);
+                    break;
+                case PUBLIC:
+                    mod = Modifier.setPublic(mod);
+                    break;
+                default:
+                    break;
+            }
+            if (obj instanceof CtField)
+                ((CtField) obj).setModifiers(mod);
+            else if (obj instanceof CtBehavior)
+                ((CtBehavior) obj).setModifiers(mod);
+            else
+                ((CtClass) obj).setModifiers(mod);
+        }
+    }
+
     public static void renameClasses(CtClass c, final Translator translator) {
         renameClasses(c, className -> {
             ClassEntry entry = translator.translateEntry(new ClassEntry(className));
@@ -110,6 +146,7 @@ public class ClassRenamer {
 
         // rename the constant pool (covers ClassInfo, MethodTypeInfo, and NameAndTypeInfo)
         ConstPool constPool = c.getClassFile().getConstPool();
+        String className = constPool.getClassName();
         constPool.renameClass(map);
 
         // rename class attributes
@@ -140,8 +177,9 @@ public class ClassRenamer {
         if (attr != null) {
             for (int i = 0; i < attr.tableLength(); i++) {
 
+                String innerName = attr.innerClass(i);
                 // get the inner class full name (which has already been translated)
-                ClassEntry classEntry = new ClassEntry(Descriptor.toJvmName(attr.innerClass(i)));
+                ClassEntry classEntry = new ClassEntry(Descriptor.toJvmName(innerName));
 
                 if (attr.innerNameIndex(i) != 0) {
                     // update the inner name
diff --git a/src/main/java/cuchaz/enigma/bytecode/ClassTranslator.java b/src/main/java/cuchaz/enigma/bytecode/ClassTranslator.java
index 6c05b83..62ebfaf 100644
--- a/src/main/java/cuchaz/enigma/bytecode/ClassTranslator.java
+++ b/src/main/java/cuchaz/enigma/bytecode/ClassTranslator.java
@@ -11,14 +11,9 @@
 package cuchaz.enigma.bytecode;
 
 import cuchaz.enigma.mapping.*;
-import javassist.CtBehavior;
-import javassist.CtClass;
-import javassist.CtField;
-import javassist.CtMethod;
-import javassist.bytecode.ConstPool;
-import javassist.bytecode.Descriptor;
-import javassist.bytecode.EnclosingMethodAttribute;
-import javassist.bytecode.SourceFileAttribute;
+import cuchaz.enigma.mapping.Translator;
+import javassist.*;
+import javassist.bytecode.*;
 
 public class ClassTranslator {
 
@@ -74,6 +69,9 @@ public class ClassTranslator {
         }
 
         ClassEntry classEntry = new ClassEntry(Descriptor.toJvmName(c.getName()));
+        Mappings.EntryModifier modifier = this.translator.getModifier(classEntry);
+        if (modifier != null && modifier != Mappings.EntryModifier.UNCHANGED)
+            ClassRenamer.applyModifier(c, modifier);
 
         // translate all the fields
         for (CtField field : c.getDeclaredFields()) {
@@ -81,6 +79,10 @@ public class ClassTranslator {
             // translate the name
             FieldEntry entry = EntryFactory.getFieldEntry(field);
             String translatedName = this.translator.translate(entry);
+            modifier = this.translator.getModifier(entry);
+            if (modifier != null && modifier != Mappings.EntryModifier.UNCHANGED)
+                ClassRenamer.applyModifier(field, modifier);
+
             if (translatedName != null) {
                 field.setName(translatedName);
             }
@@ -95,6 +97,10 @@ public class ClassTranslator {
 
             BehaviorEntry entry = EntryFactory.getBehaviorEntry(behavior);
 
+            modifier = this.translator.getModifier(entry);
+            if (modifier != null && modifier != Mappings.EntryModifier.UNCHANGED)
+                ClassRenamer.applyModifier(behavior, modifier);
+
             if (behavior instanceof CtMethod) {
                 CtMethod method = (CtMethod) behavior;
 
@@ -149,5 +155,8 @@ public class ClassTranslator {
             String sourceFile = Descriptor.toJvmName(deobfClassEntry.getOutermostClassEntry().getSimpleName()) + ".java";
             c.getClassFile().addAttribute(new SourceFileAttribute(constants, sourceFile));
         }
+        InnerClassesAttribute attr = (InnerClassesAttribute) c.getClassFile().getAttribute(InnerClassesAttribute.tag);
+        if (attr != null)
+            InnerClassWriter.changeModifier(c, attr, translator);
     }
 }
diff --git a/src/main/java/cuchaz/enigma/bytecode/InnerClassWriter.java b/src/main/java/cuchaz/enigma/bytecode/InnerClassWriter.java
index 6d92610..6e2a29d 100644
--- a/src/main/java/cuchaz/enigma/bytecode/InnerClassWriter.java
+++ b/src/main/java/cuchaz/enigma/bytecode/InnerClassWriter.java
@@ -15,22 +15,22 @@ import com.google.common.collect.Lists;
 import java.util.Collection;
 import java.util.List;
 
+import cuchaz.enigma.Deobfuscator;
 import cuchaz.enigma.analysis.JarIndex;
-import cuchaz.enigma.mapping.BehaviorEntry;
-import cuchaz.enigma.mapping.ClassEntry;
-import cuchaz.enigma.mapping.EntryFactory;
+import cuchaz.enigma.mapping.*;
+import javassist.ClassPool;
 import javassist.CtClass;
-import javassist.bytecode.AccessFlag;
-import javassist.bytecode.ConstPool;
-import javassist.bytecode.EnclosingMethodAttribute;
-import javassist.bytecode.InnerClassesAttribute;
+import javassist.NotFoundException;
+import javassist.bytecode.*;
 
 public class InnerClassWriter {
 
     private JarIndex index;
+    private Translator deobfuscatorTranslator;
 
-    public InnerClassWriter(JarIndex index) {
+    public InnerClassWriter(JarIndex index, Translator deobfuscatorTranslator) {
         this.index = index;
+        this.deobfuscatorTranslator = deobfuscatorTranslator;
     }
 
     public void write(CtClass c) {
@@ -96,6 +96,29 @@ public class InnerClassWriter {
         }
     }
 
+    // FIXME: modiffier is not applied to inner class
+    public static void changeModifier(CtClass c, InnerClassesAttribute attr, Translator translator)
+    {
+        ClassPool pool = c.getClassPool();
+        for (int i = 0; i < attr.tableLength(); i++) {
+
+            String innerName = attr.innerClass(i);
+            // get the inner class full name (which has already been translated)
+            ClassEntry classEntry = new ClassEntry(Descriptor.toJvmName(innerName));
+            try
+            {
+                CtClass innerClass = pool.get(innerName);
+                Mappings.EntryModifier modifier = translator.getModifier(classEntry);
+                if (modifier != null && modifier != Mappings.EntryModifier.UNCHANGED)
+                    ClassRenamer.applyModifier(innerClass, modifier);
+            } catch (NotFoundException e)
+            {
+                // This shouldn't be possible in theory
+                //e.printStackTrace();
+            }
+        }
+    }
+
     private void writeInnerClass(InnerClassesAttribute attr, List<ClassEntry> obfClassChain, ClassEntry obfClassEntry) {
 
         // get the new inner class name