001    package cpw.mods.fml.common.asm.transformers;
002    
003    import java.io.BufferedInputStream;
004    import java.io.BufferedOutputStream;
005    import java.io.ByteArrayOutputStream;
006    import java.io.File;
007    import java.io.FileInputStream;
008    import java.io.FileNotFoundException;
009    import java.io.FileOutputStream;
010    import java.io.IOException;
011    import java.net.URL;
012    import java.util.List;
013    import java.util.zip.ZipEntry;
014    import java.util.zip.ZipInputStream;
015    import java.util.zip.ZipOutputStream;
016    
017    import org.objectweb.asm.ClassReader;
018    import org.objectweb.asm.ClassWriter;
019    import org.objectweb.asm.tree.ClassNode;
020    
021    import com.google.common.base.Charsets;
022    import com.google.common.base.Splitter;
023    import com.google.common.collect.ArrayListMultimap;
024    import com.google.common.collect.Iterables;
025    import com.google.common.collect.ListMultimap;
026    import com.google.common.collect.Lists;
027    import com.google.common.io.LineProcessor;
028    import com.google.common.io.Resources;
029    
030    import cpw.mods.fml.relauncher.IClassTransformer;
031    
032    public class MarkerTransformer implements IClassTransformer
033    {
034        private ListMultimap<String, String> markers = ArrayListMultimap.create();
035    
036        public MarkerTransformer() throws IOException
037        {
038            this("fml_marker.cfg");
039        }
040        protected MarkerTransformer(String rulesFile) throws IOException
041        {
042            readMapFile(rulesFile);
043        }
044    
045        private void readMapFile(String rulesFile) throws IOException
046        {
047            File file = new File(rulesFile);
048            URL rulesResource;
049            if (file.exists())
050            {
051                rulesResource = file.toURI().toURL();
052            }
053            else
054            {
055                rulesResource = Resources.getResource(rulesFile);
056            }
057            Resources.readLines(rulesResource, Charsets.UTF_8, new LineProcessor<Void>()
058            {
059                @Override
060                public Void getResult()
061                {
062                    return null;
063                }
064    
065                @Override
066                public boolean processLine(String input) throws IOException
067                {
068                    String line = Iterables.getFirst(Splitter.on('#').limit(2).split(input), "").trim();
069                    if (line.length()==0)
070                    {
071                        return true;
072                    }
073                    List<String> parts = Lists.newArrayList(Splitter.on(" ").trimResults().split(line));
074                    if (parts.size()!=2)
075                    {
076                        throw new RuntimeException("Invalid config file line "+ input);
077                    }
078                    List<String> markerInterfaces = Lists.newArrayList(Splitter.on(",").trimResults().split(parts.get(1)));
079                    for (String marker : markerInterfaces)
080                    {
081                        markers.put(parts.get(0), marker);
082                    }
083                    return true;
084                }
085            });
086        }
087    
088        @SuppressWarnings("unchecked")
089        @Override
090        public byte[] transform(String name, byte[] bytes)
091        {
092            if (!markers.containsKey(name)) { return bytes; }
093    
094            ClassNode classNode = new ClassNode();
095            ClassReader classReader = new ClassReader(bytes);
096            classReader.accept(classNode, 0);
097    
098            for (String marker : markers.get(name))
099            {
100                classNode.interfaces.add(marker);
101            }
102    
103            ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
104            classNode.accept(writer);
105            return writer.toByteArray();
106        }
107    
108        public static void main(String[] args)
109        {
110            if (args.length < 2)
111            {
112                System.out.println("Usage: MarkerTransformer <JarPath> <MapFile> [MapFile2]... ");
113                return;
114            }
115    
116            boolean hasTransformer = false;
117            MarkerTransformer[] trans = new MarkerTransformer[args.length - 1];
118            for (int x = 1; x < args.length; x++)
119            {
120                try
121                {
122                    trans[x - 1] = new MarkerTransformer(args[x]);
123                    hasTransformer = true;
124                }
125                catch (IOException e)
126                {
127                    System.out.println("Could not read Transformer Map: " + args[x]);
128                    e.printStackTrace();
129                }
130            }
131    
132            if (!hasTransformer)
133            {
134                System.out.println("Culd not find a valid transformer to perform");
135                return;
136            }
137    
138            File orig = new File(args[0]);
139            File temp = new File(args[0] + ".ATBack");
140            if (!orig.exists() && !temp.exists())
141            {
142                System.out.println("Could not find target jar: " + orig);
143                return;
144            }
145    /*
146            if (temp.exists())
147            {
148                if (orig.exists() && !orig.renameTo(new File(args[0] + (new SimpleDateFormat(".yyyy.MM.dd.HHmmss")).format(new Date()))))
149                {
150                    System.out.println("Could not backup existing file: " + orig);
151                    return;
152                }
153                if (!temp.renameTo(orig))
154                {
155                    System.out.println("Could not restore backup from previous run: " + temp);
156                    return;
157                }
158            }
159    */
160            if (!orig.renameTo(temp))
161            {
162                System.out.println("Could not rename file: " + orig + " -> " + temp);
163                return;
164            }
165    
166            try
167            {
168                processJar(temp, orig, trans);
169            }
170            catch (IOException e)
171            {
172                e.printStackTrace();
173            }
174    
175            if (!temp.delete())
176            {
177                System.out.println("Could not delete temp file: " + temp);
178            }
179        }
180    
181        private static void processJar(File inFile, File outFile, MarkerTransformer[] transformers) throws IOException
182        {
183            ZipInputStream inJar = null;
184            ZipOutputStream outJar = null;
185    
186            try
187            {
188                try
189                {
190                    inJar = new ZipInputStream(new BufferedInputStream(new FileInputStream(inFile)));
191                }
192                catch (FileNotFoundException e)
193                {
194                    throw new FileNotFoundException("Could not open input file: " + e.getMessage());
195                }
196    
197                try
198                {
199                    outJar = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(outFile)));
200                }
201                catch (FileNotFoundException e)
202                {
203                    throw new FileNotFoundException("Could not open output file: " + e.getMessage());
204                }
205    
206                ZipEntry entry;
207                while ((entry = inJar.getNextEntry()) != null)
208                {
209                    if (entry.isDirectory())
210                    {
211                        outJar.putNextEntry(entry);
212                        continue;
213                    }
214    
215                    byte[] data = new byte[4096];
216                    ByteArrayOutputStream entryBuffer = new ByteArrayOutputStream();
217    
218                    int len;
219                    do
220                    {
221                        len = inJar.read(data);
222                        if (len > 0)
223                        {
224                            entryBuffer.write(data, 0, len);
225                        }
226                    }
227                    while (len != -1);
228    
229                    byte[] entryData = entryBuffer.toByteArray();
230    
231                    String entryName = entry.getName();
232    
233                    if (entryName.endsWith(".class") && !entryName.startsWith("."))
234                    {
235                        ClassNode cls = new ClassNode();
236                        ClassReader rdr = new ClassReader(entryData);
237                        rdr.accept(cls, 0);
238                        String name = cls.name.replace('/', '.').replace('\\', '.');
239    
240                        for (MarkerTransformer trans : transformers)
241                        {
242                            entryData = trans.transform(name, entryData);
243                        }
244                    }
245    
246                    ZipEntry newEntry = new ZipEntry(entryName);
247                    outJar.putNextEntry(newEntry);
248                    outJar.write(entryData);
249                }
250            }
251            finally
252            {
253                if (outJar != null)
254                {
255                    try
256                    {
257                        outJar.close();
258                    }
259                    catch (IOException e)
260                    {
261                    }
262                }
263    
264                if (inJar != null)
265                {
266                    try
267                    {
268                        inJar.close();
269                    }
270                    catch (IOException e)
271                    {
272                    }
273                }
274            }
275        }
276    }