001/*
002 * Copyright (C) 2015-2021 KeepSafe Software
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package com.getkeepsafe.dexcount.source;
017
018import com.android.dexdeps.FieldRef;
019import com.android.dexdeps.MethodRef;
020import com.android.tools.r8.CompilationFailedException;
021import com.android.tools.r8.D8;
022import com.android.tools.r8.D8Command;
023import com.android.tools.r8.OutputMode;
024import com.getkeepsafe.dexcount.DexCountException;
025import javassist.ByteArrayClassPath;
026import javassist.ClassPool;
027import javassist.CtBehavior;
028import javassist.CtClass;
029import javassist.CtConstructor;
030import javassist.CtMethod;
031import javassist.NotFoundException;
032import org.apache.commons.io.FileUtils;
033import org.apache.commons.io.IOUtils;
034
035import java.io.File;
036import java.io.IOException;
037import java.io.InputStream;
038import java.nio.charset.StandardCharsets;
039import java.nio.file.FileVisitResult;
040import java.nio.file.Files;
041import java.nio.file.Path;
042import java.nio.file.SimpleFileVisitor;
043import java.nio.file.attribute.BasicFileAttributes;
044import java.util.ArrayList;
045import java.util.Arrays;
046import java.util.Collections;
047import java.util.Enumeration;
048import java.util.List;
049import java.util.regex.Matcher;
050import java.util.regex.Pattern;
051import java.util.stream.Collectors;
052import java.util.stream.Stream;
053import java.util.zip.ZipEntry;
054import java.util.zip.ZipException;
055import java.util.zip.ZipFile;
056
057public class SourceFiles {
058    private static final Pattern CLASSES_DEX = Pattern.compile("(.*/)*classes.*\\.dex");
059    private static final Pattern CLASSES_JAR = Pattern.compile("(.*/)*classes\\.jar");
060    private static final Pattern MIN_SDK_VERSION = Pattern.compile("android:minSdkVersion=\"(\\d+)\"");
061
062    private SourceFiles() {
063        // no instances
064    }
065
066    public static List<SourceFile> extractDexData(File file) throws IOException {
067        if (file == null || !file.exists()) {
068            return Collections.emptyList();
069        }
070
071        // AAR files need special treatment
072        if (file.getName().endsWith(".aar")) {
073            return extractDexFromAar(file);
074        }
075
076        try {
077            return extractDexFromZip(file);
078        } catch (ZipException ignored) {
079            // not a zip, no problem
080        }
081
082        return Collections.singletonList(new DexFile(file, false));
083    }
084
085    private static List<SourceFile> extractDexFromAar(File file) throws IOException {
086        int minSdk = 13;
087        File tempClasses = null;
088        try (ZipFile zip = new ZipFile(file)) {
089            Enumeration<? extends ZipEntry> entries = zip.entries();
090            while (entries.hasMoreElements()) {
091                ZipEntry entry = entries.nextElement();
092                if ("AndroidManifest.xml".equals(entry.getName())) {
093                    String text;
094                    try (InputStream is = zip.getInputStream(entry)) {
095                        text = IOUtils.toString(is, StandardCharsets.UTF_8);
096                    }
097
098                    Matcher matcher = MIN_SDK_VERSION.matcher(text);
099                    if (!matcher.find()) {
100                        continue;
101                    }
102
103                    minSdk = Integer.parseInt(matcher.group(1));
104                }
105
106                if (CLASSES_JAR.matcher(entry.getName()).matches()) {
107                    tempClasses = makeTemp(entry.getName());
108                    try (InputStream is = zip.getInputStream(entry)) {
109                        FileUtils.copyInputStreamToFile(is, tempClasses);
110                    }
111                }
112            }
113        }
114
115        if (tempClasses == null) {
116            throw new IllegalArgumentException("No classes.jar entry found in " + file.getCanonicalPath());
117        }
118
119        Path tempDexDir = Files.createTempDirectory("dex");
120        tempDexDir.toFile().deleteOnExit();
121
122        try {
123            D8Command command = D8Command.builder()
124                .addProgramFiles(tempClasses.toPath())
125                .setMinApiLevel(minSdk)
126                .setOutput(tempDexDir, OutputMode.DexIndexed)
127                .build();
128
129            D8.run(command);
130        } catch (CompilationFailedException e) {
131            throw new DexCountException("Failed to run D8 on an AAR", e);
132        }
133
134        List<SourceFile> results = new ArrayList<>();
135        for (Path path : Files.list(tempDexDir).collect(Collectors.toList())) {
136            if (!Files.isRegularFile(path)) {
137                continue;
138            }
139
140            results.add(new DexFile(path.toFile(), true));
141        }
142
143        return results;
144    }
145
146    private static List<SourceFile> extractDexFromZip(File file) throws IOException {
147        List<SourceFile> results = new ArrayList<>();
148
149        try (ZipFile zip = new ZipFile(file)) {
150            Enumeration<? extends ZipEntry> entries = zip.entries();
151            while (entries.hasMoreElements()) {
152                ZipEntry entry = entries.nextElement();
153                if (!CLASSES_DEX.matcher(entry.getName()).matches()) {
154                    continue;
155                }
156
157                File temp = makeTemp(entry.getName());
158                try (InputStream is = zip.getInputStream(entry)) {
159                    FileUtils.copyInputStreamToFile(is, temp);
160                }
161
162                results.add(new DexFile(temp, true));
163            }
164        }
165
166        return results;
167    }
168
169    public static SourceFile extractJarFromAar(File aar) throws IOException {
170        File tempClassesJar = null;
171        try (ZipFile zip = new ZipFile(aar)) {
172            Enumeration<? extends ZipEntry> entries = zip.entries();
173            while (entries.hasMoreElements()) {
174                ZipEntry entry = entries.nextElement();
175
176                if (!CLASSES_JAR.matcher(entry.getName()).matches()) {
177                    continue;
178                }
179
180                tempClassesJar = makeTemp(entry.getName());
181                try (InputStream is = zip.getInputStream(entry)) {
182                    FileUtils.copyInputStreamToFile(is, tempClassesJar);
183                }
184            }
185        }
186
187        if (tempClassesJar == null) {
188            throw new IllegalArgumentException("No classes.jar entry found in " + aar.getCanonicalPath());
189        }
190
191        return extractJarFromJar(tempClassesJar);
192    }
193
194    public static SourceFile extractJarFromJar(File jar) throws IOException {
195        // Unzip the classes.jar file and store all .class files in this directory.
196        File classFilesDir = Files.createTempDirectory("classFilesDir").toFile();
197        classFilesDir.deleteOnExit();
198
199        final Path classFilesPath = classFilesDir.toPath();
200
201        try (ZipFile zip = new ZipFile(jar)) {
202            zip.stream().filter(it -> it.getName().endsWith(".class")).forEach(entry -> {
203                String fileName = entry.getName();
204                File file = new File(classFilesDir, fileName);
205
206                try {
207                    FileUtils.createParentDirectories(file);
208                    try (InputStream is = zip.getInputStream(entry)) {
209                        FileUtils.copyInputStreamToFile(is, file);
210                    }
211                } catch (IOException e) {
212                    throw new DexCountException("Failed to unzip a classes.jar file");
213                }
214            });
215        }
216
217        ClassPool classPool = new ClassPool();
218        classPool.appendSystemPath();
219
220        List<CtClass> classes = new ArrayList<>();
221        Files.walkFileTree(classFilesDir.toPath(), new SimpleFileVisitor<Path>() {
222            @Override
223            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
224                if (!attrs.isRegularFile()) {
225                    return FileVisitResult.CONTINUE;
226                }
227
228                String qualifiedClassName = classFilesPath.relativize(file)
229                    .toFile()
230                    .getPath()
231                    .replace('/', '.')
232                    .replace(".class", "");
233
234                ByteArrayClassPath cp = new ByteArrayClassPath(qualifiedClassName, Files.readAllBytes(file));
235                classPool.appendClassPath(cp);
236
237                try {
238                    classes.add(classPool.get(qualifiedClassName));
239                } catch (NotFoundException e) {
240                    throw new AssertionError("We literally just added this class to the pool", e);
241                }
242
243                return FileVisitResult.CONTINUE;
244            }
245        });
246
247        List<MethodRef> methodRefs = classes.stream().flatMap(SourceFiles::extractMethodRefs).collect(Collectors.toList());
248        List<FieldRef> fieldRefs = classes.stream().flatMap(SourceFiles::extractFieldRefs).collect(Collectors.toList());
249
250        return new JarFile(methodRefs, fieldRefs);
251    }
252
253    private static Stream<MethodRef> extractMethodRefs(CtClass clazz) {
254        String declaringClass = "L" + clazz.getName().replace(".", "/") + ";";
255
256        // Unfortunately, it's necessary to parse the types from the strings manually.
257        // We can't use the proper API because this requires all classes that are used
258        // in parameters and return types to be loaded in the classpath. However,
259        // that's not the case when we analyze a single jar file.
260        List<MethodRef> results = new ArrayList<>();
261        if (clazz.getClassInitializer() != null) {
262            results.add(new MethodRef(declaringClass, new String[0], "V", "<clinit>"));
263        }
264
265        for (CtConstructor ctor : clazz.getDeclaredConstructors()) {
266            String[] params = parseBehaviorParameters(ctor);
267            results.add(new MethodRef(declaringClass, params, "V", "<init>"));
268        }
269
270        for (CtMethod method : clazz.getDeclaredMethods()) {
271            String[] params = parseBehaviorParameters(method);
272            String returnType = parseMethodReturnType(method);
273            results.add(new MethodRef(declaringClass, params, returnType, method.getName()));
274        }
275
276        return results.stream();
277    }
278
279    private static Stream<FieldRef> extractFieldRefs(CtClass clazz) {
280        return Arrays.stream(clazz.getDeclaredFields()).map(field -> {
281            String type = field.getFieldInfo().getDescriptor();
282            return new FieldRef(clazz.getSimpleName(), type, field.getName());
283        });
284    }
285
286    private static String[] parseBehaviorParameters(CtBehavior behavior) {
287        String signature = behavior.getSignature();
288        int startIx = signature.indexOf('(');
289        int endIx = signature.indexOf(')', startIx);
290        String parameters = signature.substring(startIx + 1, endIx);
291        return parameters.split(";");
292    }
293
294    private static String parseMethodReturnType(CtMethod method) {
295        int ix = method.getSignature().indexOf(')');
296        return method.getSignature().substring(ix + 1);
297    }
298
299    private static File makeTemp(String pattern) {
300        int ix = pattern.indexOf('.');
301        String prefix = pattern.substring(0, ix);
302        String suffix = pattern.substring(ix);
303        try {
304            File temp = File.createTempFile(prefix, suffix);
305            temp.deleteOnExit();
306            return temp;
307        } catch (IOException e) {
308            throw new DexCountException("Failed to create temp file", e);
309        }
310    }
311}