1 files changed, 137 insertions, 0 deletions

diff --git a/src/cuchaz/enigma/analysis/JarClassIterator.java b/src/cuchaz/enigma/analysis/JarClassIterator.java
new file mode 100644
index 0000000..aa58e9e
--- /dev/null
+++ b/src/cuchaz/enigma/analysis/JarClassIterator.java
@@ -0,0 +1,137 @@
+/*******************************************************************************
+ * Copyright (c) 2015 Jeff Martin.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the GNU Lesser General Public
+ * License v3.0 which accompanies this distribution, and is available at
+ * http://www.gnu.org/licenses/lgpl.html
+ * 
+ * Contributors:
+ *     Jeff Martin - initial API and implementation
+ ******************************************************************************/
+package cuchaz.enigma.analysis;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Enumeration;
+import java.util.Iterator;
+import java.util.List;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
+
+import javassist.ByteArrayClassPath;
+import javassist.ClassPool;
+import javassist.CtClass;
+import javassist.NotFoundException;
+import javassist.bytecode.Descriptor;
+
+import com.google.common.collect.Lists;
+
+import cuchaz.enigma.Constants;
+import cuchaz.enigma.mapping.ClassEntry;
+
+public class JarClassIterator implements Iterator<CtClass> {
+        
+        private JarFile m_jar;
+        private Iterator<JarEntry> m_iter;
+        
+        public JarClassIterator(JarFile jar) {
+                m_jar = jar;
+                
+                // get the jar entries that correspond to classes
+                List<JarEntry> classEntries = Lists.newArrayList();
+                Enumeration<JarEntry> entries = m_jar.entries();
+                while (entries.hasMoreElements()) {
+                        JarEntry entry = entries.nextElement();
+                        
+                        // is this a class file?
+                        if (entry.getName().endsWith(".class")) {
+                                classEntries.add(entry);
+                        }
+                }
+                m_iter = classEntries.iterator();
+        }
+        
+        @Override
+        public boolean hasNext() {
+                return m_iter.hasNext();
+        }
+        
+        @Override
+        public CtClass next() {
+                JarEntry entry = m_iter.next();
+                try {
+                        return getClass(m_jar, entry);
+                } catch (IOException | NotFoundException ex) {
+                        throw new Error("Unable to load class: " + entry.getName());
+                }
+        }
+        
+        @Override
+        public void remove() {
+                throw new UnsupportedOperationException();
+        }
+        
+        public static List<ClassEntry> getClassEntries(JarFile jar) {
+                List<ClassEntry> classEntries = Lists.newArrayList();
+                Enumeration<JarEntry> entries = jar.entries();
+                while (entries.hasMoreElements()) {
+                        JarEntry entry = entries.nextElement();
+                        
+                        // is this a class file?
+                        if (!entry.isDirectory() && entry.getName().endsWith(".class")) {
+                                classEntries.add(getClassEntry(entry));
+                        }
+                }
+                return classEntries;
+        }
+        
+        public static Iterable<CtClass> classes(final JarFile jar) {
+                return new Iterable<CtClass>() {
+                        @Override
+                        public Iterator<CtClass> iterator() {
+                                return new JarClassIterator(jar);
+                        }
+                };
+        }
+        
+        public static CtClass getClass(JarFile jar, ClassEntry classEntry) {
+                try {
+                        return getClass(jar, new JarEntry(classEntry.getName() + ".class"));
+                } catch (IOException | NotFoundException ex) {
+                        throw new Error("Unable to load class: " + classEntry.getName());
+                }
+        }
+        
+        private static CtClass getClass(JarFile jar, JarEntry entry) throws IOException, NotFoundException {
+                // read the class into a buffer
+                ByteArrayOutputStream bos = new ByteArrayOutputStream();
+                byte[] buf = new byte[Constants.KiB];
+                int totalNumBytesRead = 0;
+                InputStream in = jar.getInputStream(entry);
+                while (in.available() > 0) {
+                        int numBytesRead = in.read(buf);
+                        if (numBytesRead < 0) {
+                                break;
+                        }
+                        bos.write(buf, 0, numBytesRead);
+                        
+                        // sanity checking
+                        totalNumBytesRead += numBytesRead;
+                        if (totalNumBytesRead > Constants.MiB) {
+                                throw new Error("Class file " + entry.getName() + " larger than 1 MiB! Something is wrong!");
+                        }
+                }
+                
+                // get a javassist handle for the class
+                String className = Descriptor.toJavaName(getClassEntry(entry).getName());
+                ClassPool classPool = new ClassPool();
+                classPool.appendSystemPath();
+                classPool.insertClassPath(new ByteArrayClassPath(className, bos.toByteArray()));
+                return classPool.get(className);
+        }
+        
+        private static ClassEntry getClassEntry(JarEntry entry) {
+                return new ClassEntry(entry.getName().substring(0, entry.getName().length() - ".class".length()));
+        }
+}