001package cpw.mods.fml.common.asm.transformers;
002
003import java.io.BufferedInputStream;
004import java.io.BufferedOutputStream;
005import java.io.ByteArrayOutputStream;
006import java.io.File;
007import java.io.FileInputStream;
008import java.io.FileNotFoundException;
009import java.io.FileOutputStream;
010import java.io.IOException;
011import java.net.URL;
012import java.util.List;
013import java.util.zip.ZipEntry;
014import java.util.zip.ZipInputStream;
015import java.util.zip.ZipOutputStream;
016
017import org.objectweb.asm.ClassReader;
018import org.objectweb.asm.ClassWriter;
019import org.objectweb.asm.tree.ClassNode;
020
021import com.google.common.base.Charsets;
022import com.google.common.base.Splitter;
023import com.google.common.collect.ArrayListMultimap;
024import com.google.common.collect.Iterables;
025import com.google.common.collect.ListMultimap;
026import com.google.common.collect.Lists;
027import com.google.common.io.LineProcessor;
028import com.google.common.io.Resources;
029
030import cpw.mods.fml.relauncher.IClassTransformer;
031
032public 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 (bytes == null) { return null; }
093        if (!markers.containsKey(name)) { return bytes; }
094
095        ClassNode classNode = new ClassNode();
096        ClassReader classReader = new ClassReader(bytes);
097        classReader.accept(classNode, 0);
098
099        for (String marker : markers.get(name))
100        {
101            classNode.interfaces.add(marker);
102        }
103
104        ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
105        classNode.accept(writer);
106        return writer.toByteArray();
107    }
108
109    public static void main(String[] args)
110    {
111        if (args.length < 2)
112        {
113            System.out.println("Usage: MarkerTransformer <JarPath> <MapFile> [MapFile2]... ");
114            return;
115        }
116
117        boolean hasTransformer = false;
118        MarkerTransformer[] trans = new MarkerTransformer[args.length - 1];
119        for (int x = 1; x < args.length; x++)
120        {
121            try
122            {
123                trans[x - 1] = new MarkerTransformer(args[x]);
124                hasTransformer = true;
125            }
126            catch (IOException e)
127            {
128                System.out.println("Could not read Transformer Map: " + args[x]);
129                e.printStackTrace();
130            }
131        }
132
133        if (!hasTransformer)
134        {
135            System.out.println("Culd not find a valid transformer to perform");
136            return;
137        }
138
139        File orig = new File(args[0]);
140        File temp = new File(args[0] + ".ATBack");
141        if (!orig.exists() && !temp.exists())
142        {
143            System.out.println("Could not find target jar: " + orig);
144            return;
145        }
146/*
147        if (temp.exists())
148        {
149            if (orig.exists() && !orig.renameTo(new File(args[0] + (new SimpleDateFormat(".yyyy.MM.dd.HHmmss")).format(new Date()))))
150            {
151                System.out.println("Could not backup existing file: " + orig);
152                return;
153            }
154            if (!temp.renameTo(orig))
155            {
156                System.out.println("Could not restore backup from previous run: " + temp);
157                return;
158            }
159        }
160*/
161        if (!orig.renameTo(temp))
162        {
163            System.out.println("Could not rename file: " + orig + " -> " + temp);
164            return;
165        }
166
167        try
168        {
169            processJar(temp, orig, trans);
170        }
171        catch (IOException e)
172        {
173            e.printStackTrace();
174        }
175
176        if (!temp.delete())
177        {
178            System.out.println("Could not delete temp file: " + temp);
179        }
180    }
181
182    private static void processJar(File inFile, File outFile, MarkerTransformer[] transformers) throws IOException
183    {
184        ZipInputStream inJar = null;
185        ZipOutputStream outJar = null;
186
187        try
188        {
189            try
190            {
191                inJar = new ZipInputStream(new BufferedInputStream(new FileInputStream(inFile)));
192            }
193            catch (FileNotFoundException e)
194            {
195                throw new FileNotFoundException("Could not open input file: " + e.getMessage());
196            }
197
198            try
199            {
200                outJar = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(outFile)));
201            }
202            catch (FileNotFoundException e)
203            {
204                throw new FileNotFoundException("Could not open output file: " + e.getMessage());
205            }
206
207            ZipEntry entry;
208            while ((entry = inJar.getNextEntry()) != null)
209            {
210                if (entry.isDirectory())
211                {
212                    outJar.putNextEntry(entry);
213                    continue;
214                }
215
216                byte[] data = new byte[4096];
217                ByteArrayOutputStream entryBuffer = new ByteArrayOutputStream();
218
219                int len;
220                do
221                {
222                    len = inJar.read(data);
223                    if (len > 0)
224                    {
225                        entryBuffer.write(data, 0, len);
226                    }
227                }
228                while (len != -1);
229
230                byte[] entryData = entryBuffer.toByteArray();
231
232                String entryName = entry.getName();
233
234                if (entryName.endsWith(".class") && !entryName.startsWith("."))
235                {
236                    ClassNode cls = new ClassNode();
237                    ClassReader rdr = new ClassReader(entryData);
238                    rdr.accept(cls, 0);
239                    String name = cls.name.replace('/', '.').replace('\\', '.');
240
241                    for (MarkerTransformer trans : transformers)
242                    {
243                        entryData = trans.transform(name, entryData);
244                    }
245                }
246
247                ZipEntry newEntry = new ZipEntry(entryName);
248                outJar.putNextEntry(newEntry);
249                outJar.write(entryData);
250            }
251        }
252        finally
253        {
254            if (outJar != null)
255            {
256                try
257                {
258                    outJar.close();
259                }
260                catch (IOException e)
261                {
262                }
263            }
264
265            if (inJar != null)
266            {
267                try
268                {
269                    inJar.close();
270                }
271                catch (IOException e)
272                {
273                }
274            }
275        }
276    }
277}