Source code

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;
017
018import com.android.dexdeps.FieldRef;
019import com.android.dexdeps.HasDeclaringClass;
020import com.android.dexdeps.MethodRef;
021import com.android.dexdeps.Output;
022import com.google.gson.stream.JsonWriter;
023import org.jetbrains.annotations.NotNull;
024
025import java.io.IOException;
026import java.io.Writer;
027import java.nio.CharBuffer;
028import java.util.Arrays;
029import java.util.Collections;
030import java.util.LinkedHashMap;
031import java.util.LinkedHashSet;
032import java.util.List;
033import java.util.Map;
034import java.util.Map.Entry;
035import java.util.Set;
036import java.util.SortedMap;
037import java.util.TreeMap;
038import java.util.stream.Collectors;
039import java.util.stream.Stream;
040
041public class PackageTree {
042    private enum Type {
043        DECLARED,
044        REFERENCED,
045    }
046
047    private final String name;
048    private final boolean isClass;
049    private final Deobfuscator deobfuscator;
050
051    private final LinkedHashMap<Type, Integer> classTotal = new LinkedHashMap<>();
052    private final LinkedHashMap<Type, Integer> methodTotal = new LinkedHashMap<>();
053    private final LinkedHashMap<Type, Integer> fieldTotal = new LinkedHashMap<>();
054    private final SortedMap<String, PackageTree> children = new TreeMap<>();
055    private final LinkedHashMap<Type, LinkedHashSet<MethodRef>> methods = new LinkedHashMap<>();
056    private final LinkedHashMap<Type, LinkedHashSet<FieldRef>> fields = new LinkedHashMap<>();
057
058    public PackageTree() {
059        this("", false, null);
060    }
061
062    public PackageTree(Deobfuscator deobfuscator) {
063        this("", false, deobfuscator);
064    }
065
066    public PackageTree(String name, Deobfuscator deobfuscator) {
067        this(name, isClassName(name), deobfuscator);
068    }
069
070    public PackageTree(String name, boolean isClass, Deobfuscator deobfuscator) {
071        if (name == null) {
072            throw new NullPointerException("name");
073        }
074
075        if (deobfuscator == null) {
076            deobfuscator = Deobfuscator.EMPTY;
077        }
078
079        this.name = name;
080        this.isClass = isClass;
081        this.deobfuscator = deobfuscator;
082
083        for (Type type : Type.values()) {
084            methods.put(type, new LinkedHashSet<>());
085            fields.put(type, new LinkedHashSet<>());
086        }
087    }
088
089    public String getName() {
090        return name;
091    }
092
093    public boolean isClass() {
094        return isClass;
095    }
096
097    public int getClassCount() {
098        return getClassCount(Type.REFERENCED);
099    }
100
101    public int getClassCountDeclared() {
102        return getClassCount(Type.DECLARED);
103    }
104
105    public int getMethodCount() {
106        return getMethodCount(Type.REFERENCED);
107    }
108
109    public int getMethodCountDeclared() {
110        return getMethodCount(Type.DECLARED);
111    }
112
113    public int getFieldCount() {
114        return getFieldCount(Type.REFERENCED);
115    }
116
117    public int getFieldCountDeclared() {
118        return getFieldCount(Type.DECLARED);
119    }
120
121    private int getClassCount(Type type) {
122        Integer maybeTotal = classTotal.get(type);
123        if (maybeTotal != null) {
124            return maybeTotal;
125        }
126
127        if (isClass) {
128            classTotal.put(type, 1);
129            return 1;
130        }
131
132        int result = children.values().parallelStream().mapToInt(child -> child.getClassCount(type)).sum();
133        classTotal.put(type, result);
134
135        return result;
136    }
137
138    private int getMethodCount(Type type) {
139        Integer maybeTotal = methodTotal.get(type);
140        if (maybeTotal != null) {
141            return maybeTotal;
142        }
143
144        int result = methods.get(type).size() + children.values().parallelStream().mapToInt(child -> child.getMethodCount(type)).sum();
145        methodTotal.put(type, result);
146
147        return result;
148    }
149
150    private int getFieldCount(Type type) {
151        Integer maybeTotal = fieldTotal.get(type);
152        if (maybeTotal != null) {
153            return maybeTotal;
154        }
155
156        int result = fields.get(type).size() + children.values().parallelStream().mapToInt(child -> child.getFieldCount(type)).sum();
157        fieldTotal.put(type, result);
158
159        return result;
160    }
161
162    public void addMethodRef(MethodRef ref) {
163        addInternal(descriptorToDot(ref), 0, true, Type.REFERENCED, ref);
164    }
165
166    public void addFieldRef(FieldRef ref) {
167        addInternal(descriptorToDot(ref), 0, false, Type.REFERENCED, ref);
168    }
169
170    public void addDeclaredMethodRef(MethodRef ref) {
171        addInternal(descriptorToDot(ref), 0, true, Type.DECLARED, ref);
172    }
173
174    public void addDeclaredFieldRef(FieldRef ref) {
175        addInternal(descriptorToDot(ref), 0, false, Type.DECLARED, ref);
176    }
177
178    private void addInternal(String name, int startIndex, boolean isMethod, Type type, HasDeclaringClass ref) {
179        int ix = name.indexOf('.', startIndex);
180        String segment;
181        if (ix == -1) {
182            segment = name.substring(startIndex);
183        } else {
184            segment = name.substring(startIndex, ix);
185        }
186
187        PackageTree child = children.get(segment);
188        if (child == null) {
189            child = new PackageTree(segment, deobfuscator);
190            children.put(segment, child);
191        }
192
193        if (ix == -1) {
194            if (isMethod) {
195                child.methods.get(type).add((MethodRef) ref);
196            } else {
197                child.fields.get(type).add((FieldRef) ref);
198            }
199        } else {
200            if (isMethod) {
201                methodTotal.remove(type);
202            } else {
203                fieldTotal.remove(type);
204            }
205            child.addInternal(name, ix + 1, isMethod, type, ref);
206        }
207    }
208
209    public void print(Appendable out, OutputFormat format, PrintOptions opts) throws IOException {
210        switch (format) {
211            case LIST:
212                printPackageList(out, opts);
213                break;
214
215            case TREE:
216                printTree(out, opts);
217                break;
218
219            case JSON:
220                printJson(out, opts);
221                break;
222
223            case YAML:
224                printYaml(out, opts);
225                break;
226
227            default:
228                throw new IllegalArgumentException("Unexpected OutputFormat: " + format);
229        }
230    }
231
232    public void printPackageList(Appendable out, PrintOptions opts) throws IOException {
233        StringBuilder sb = new StringBuilder(64);
234
235        if (opts.getIncludeTotalMethodCount()) {
236            if (opts.isAndroidProject()) {
237                out.append("Total methods: ").append(String.valueOf(getMethodCount())).append("\n");
238            }
239
240            if (opts.getPrintDeclarations()) {
241                out.append("Total declared methods: ").append(String.valueOf(getClassCountDeclared())).append("\n");
242            }
243        }
244
245        if (opts.getPrintHeader()) {
246            printPackageListHeader(out, opts);
247        }
248
249        for (PackageTree child : getChildren(opts)) {
250            child.printPackageListRecursively(out, sb, 0, opts);
251        }
252    }
253
254    private void printPackageListHeader(Appendable out, PrintOptions opts) throws IOException {
255        if (opts.getIncludeClassCount()) {
256            out.append(String.format("%-8s ", "classes"));
257        }
258
259        if (opts.isAndroidProject()) {
260            if (opts.getIncludeMethodCount()) {
261                out.append(String.format("%-8s ", "methods"));
262            }
263
264            if (opts.getIncludeFieldCount()) {
265                out.append(String.format("%-8s ", "fields"));
266            }
267        }
268
269        if (opts.getPrintDeclarations()) {
270            out.append(String.format("%-16s ", "declared methods"));
271            out.append(String.format("%-16s ", "declared fields"));
272        }
273
274        out.append("package/class name\n");
275    }
276
277    private void printPackageListRecursively(Appendable out, StringBuilder sb, int depth, PrintOptions opts) throws IOException {
278        if (depth >= opts.getMaxTreeDepth()) {
279            return;
280        }
281
282        if (!isPrintable(opts)) {
283            // Should be guaranteed by `getChildren()`
284            throw new IllegalStateException("We should never recursively print a non-printable");
285        }
286
287        int len = sb.length();
288        if (len > 0) {
289            sb.append('.');
290        }
291        sb.append(getName());
292
293        if (opts.getIncludeClassCount()) {
294            out.append(String.format("%-8d ", getClassCount()));
295        }
296
297        if (opts.isAndroidProject()) {
298            if (opts.getIncludeMethodCount()) {
299                out.append(String.format("%-8d ", getMethodCount()));
300            }
301
302            if (opts.getIncludeFieldCount()) {
303                out.append(String.format("%-8d ", getFieldCount()));
304            }
305        }
306
307        if (opts.getPrintDeclarations()) {
308            if (opts.getPrintHeader()) {
309                // The header for the these two columns uses more space.
310                out.append(String.format("%-16d ", getMethodCountDeclared()));
311                out.append(String.format("%-16d ", getFieldCountDeclared()));
312            } else {
313                out.append(String.format("%-8d ", getMethodCountDeclared()));
314                out.append(String.format("%-8d ", getFieldCountDeclared()));
315            }
316        }
317
318        out.append(sb.toString()).append("\n");
319
320        for (PackageTree child : getChildren(opts)) {
321            child.printPackageListRecursively(out, sb, depth + 1, opts);
322        }
323
324        sb.setLength(len);
325    }
326
327    public void printTree(Appendable out, PrintOptions opts) throws IOException {
328        for (PackageTree child : getChildren(opts)) {
329            child.printTreeRecursively(out, 0, opts);
330        }
331    }
332
333    private void printTreeRecursively(Appendable out, int depth, PrintOptions opts) throws IOException {
334        if (depth >= opts.getMaxTreeDepth()) {
335            return;
336        }
337
338        for (int i = 0; i < depth; i++) {
339            out.append("  ");
340        }
341        out.append(getName());
342
343        if (opts.getIncludeFieldCount() || opts.getIncludeMethodCount() || opts.getIncludeClassCount()) {
344            out.append(" (");
345
346            boolean appended = false;
347            if (opts.getIncludeClassCount()) {
348                out.append(String.valueOf(getClassCount()))
349                    .append(" ")
350                    .append(pluralizedClasses(getClassCount()));
351                appended = true;
352            }
353
354            if (opts.isAndroidProject()) {
355                if (opts.getIncludeMethodCount()) {
356                    if (appended) {
357                        out.append(", ");
358                    }
359                    out.append(String.valueOf(getMethodCount()))
360                        .append(" ")
361                        .append(pluralizedMethods(getMethodCount()));
362                    appended = true;
363                }
364
365                if (opts.getIncludeFieldCount()) {
366                    if (appended) {
367                        out.append(", ");
368                    }
369                    out.append(String.valueOf(getFieldCount()))
370                        .append(" ")
371                        .append(pluralizedFields(getFieldCount()));
372                    appended = true;
373                }
374            }
375
376            if (opts.getPrintDeclarations()) {
377                if (appended) {
378                    out.append(", ");
379                }
380                out.append(String.valueOf(getMethodCountDeclared()))
381                    .append(" declared ")
382                    .append(pluralizedMethods(getMethodCountDeclared()))
383                    .append(", ")
384                    .append(String.valueOf(getFieldCountDeclared()))
385                    .append(" declared ")
386                    .append(pluralizedFields(getFieldCountDeclared()));
387            }
388
389            out.append(")\n");
390        }
391
392        for (PackageTree child : getChildren(opts)) {
393            child.printTreeRecursively(out, depth + 1, opts);
394        }
395    }
396
397    public void printJson(Appendable out, PrintOptions opts) throws IOException {
398        JsonWriter json = new JsonWriter(new Writer() {
399            @Override
400            public void write(@NotNull char[] chars, int offset, int length) throws IOException {
401                out.append(CharBuffer.wrap(chars, offset, length));
402            }
403
404            @Override
405            public void flush() {
406                // no-op
407            }
408
409            @Override
410            public void close() {
411                // no-op
412            }
413        });
414
415        json.setIndent("  ");
416
417        printJsonRecursively(json, 0, opts);
418    }
419
420    private void printJsonRecursively(JsonWriter json, int depth, PrintOptions opts) throws IOException {
421        if (depth >= opts.getMaxTreeDepth()) {
422            return;
423        }
424
425        json.beginObject();
426
427        json.name("name").value(getName());
428
429        if (opts.getIncludeClassCount()) {
430            json.name("classes").value(getClassCount());
431        }
432
433        if (opts.isAndroidProject()) {
434            if (opts.getIncludeMethodCount()) {
435                json.name("methods").value(getMethodCount());
436            }
437
438            if (opts.getIncludeFieldCount()) {
439                json.name("fields").value(getFieldCount());
440            }
441        }
442
443        if (opts.getPrintDeclarations()) {
444            json.name("declared_methods").value(getMethodCountDeclared());
445            json.name("declared_fields").value(getFieldCountDeclared());
446        }
447
448        json.name("children");
449        json.beginArray();
450        for (PackageTree child : getChildren(opts)) {
451            child.printJsonRecursively(json, depth + 1, opts);
452        }
453        json.endArray();
454
455        json.endObject();
456    }
457
458    public void printYaml(Appendable out, PrintOptions opts) throws IOException {
459        out.append("---\n");
460
461        if (opts.getIncludeClassCount()) {
462            out.append("classes: ").append(String.valueOf(getClassCount())).append("\n");
463        }
464
465        if (opts.isAndroidProject()) {
466            if (opts.getIncludeMethodCount()) {
467                out.append("methods: ").append(String.valueOf(getMethodCount())).append("\n");
468            }
469
470            if (opts.getIncludeFieldCount()) {
471                out.append("fields: ").append(String.valueOf(getFieldCount())).append("\n");
472            }
473        }
474
475        if (opts.getPrintDeclarations()) {
476            out.append("declared_methods: ").append(String.valueOf(getMethodCountDeclared())).append("\n");
477            out.append("declared_fields: ").append(String.valueOf(getFieldCountDeclared())).append("\n");
478        }
479
480        out.append("counts:\n");
481
482        for (PackageTree child : getChildren(opts)) {
483            child.printYamlRecursively(out, 0, opts);
484        }
485    }
486
487    private void printYamlRecursively(Appendable out, int depth, PrintOptions opts) throws IOException {
488        if (depth > opts.getMaxTreeDepth()) {
489            return;
490        }
491
492        StringBuilder indentBuilder = new StringBuilder();
493        for (int i = 0; i < (depth * 2) + 1; ++i) {
494            indentBuilder.append("  ");
495        }
496        String indent = indentBuilder.toString();
497
498        out.append(indent).append("- name: ").append(getName()).append("\n");
499
500        indent += "  ";
501
502        if (opts.getIncludeClassCount()) {
503            out.append(indent).append("classes: ").append(String.valueOf(getClassCount())).append("\n");
504        }
505
506        if (opts.isAndroidProject()) {
507            if (opts.getIncludeMethodCount()) {
508                out.append(indent).append("methods: ").append(String.valueOf(getMethodCount())).append("\n");
509            }
510
511            if (opts.getIncludeFieldCount()) {
512                out.append(indent).append("fields: ").append(String.valueOf(getFieldCount())).append("\n");
513            }
514        }
515
516        if (opts.getPrintDeclarations()) {
517            out.append(indent).append("declared_methods: ").append(String.valueOf(getMethodCountDeclared())).append("\n");
518            out.append(indent).append("declared_fields: ").append(String.valueOf(getFieldCountDeclared())).append("\n");
519        }
520
521        List<PackageTree> childNodes;
522        if (depth + 1 == opts.getMaxTreeDepth()) {
523            childNodes = Collections.emptyList();
524        } else {
525            childNodes = getChildren(opts);
526        }
527
528        if (childNodes.isEmpty()) {
529            out.append(indent).append("children: []\n");
530            return;
531        }
532
533        out.append(indent).append("children:\n");
534        for (PackageTree child : getChildren(opts)) {
535            child.printYamlRecursively(out, depth + 1, opts);
536        }
537    }
538
539    private List<PackageTree> getChildren(PrintOptions opts) {
540        Stream<PackageTree> result = children.values().stream().filter(it -> it.isPrintable(opts));
541
542        if (opts.getOrderByMethodCount()) {
543            result = result.sorted((lhs, rhs) -> Integer.compare(rhs.getMethodCount(), lhs.getMethodCount()));
544        }
545
546        return result.collect(Collectors.toList());
547    }
548
549    private boolean isPrintable(PrintOptions opts) {
550        return opts.getIncludeClasses() || !isClass;
551    }
552
553    private String pluralizedClasses(int n) {
554        if (n == 1) {
555            return "class";
556        } else {
557            return "classes";
558        }
559    }
560
561    private String pluralizedMethods(int n) {
562        if (n == 1) {
563            return "method";
564        } else {
565            return "methods";
566        }
567    }
568
569    private String pluralizedFields(int n) {
570        if (n == 1) {
571            return "field";
572        } else {
573            return "fields";
574        }
575    }
576
577    private String descriptorToDot(HasDeclaringClass ref) {
578        String descriptor = ref.getDeclClassName();
579        String dot = Output.descriptorToDot(descriptor);
580        String deobfuscated = deobfuscator.deobfuscate(dot);
581        if (deobfuscated.indexOf('.') == -1) {
582            // Classes in the unnamed package (e.g. primitive arrays)
583            // will not appear in the output in the current PackageTree
584            // implementation if classes are not included.  To work around,
585            // we make an artificial package named "<unnamed>".
586            return "<unnamed>." + deobfuscated;
587        } else {
588            return deobfuscated;
589        }
590    }
591
592    @Override
593    public boolean equals(Object o) {
594        if (this == o) return true;
595        if (o == null || getClass() != o.getClass()) return false;
596
597        PackageTree that = (PackageTree) o;
598
599        if (isClass != that.isClass) return false;
600        if (!name.equals(that.name)) return false;
601        if (!children.equals(that.children)) return false;
602        if (!methods.equals(that.methods)) return false;
603        return fields.equals(that.fields);
604    }
605
606    @Override
607    public int hashCode() {
608        int result = name.hashCode();
609        result = 31 * result + (isClass ? 1 : 0);
610        result = 31 * result + children.hashCode();
611        result = 31 * result + methods.hashCode();
612        result = 31 * result + fields.hashCode();
613        return result;
614    }
615
616    private static boolean isClassName(String name) {
617        return Character.isUpperCase(name.charAt(0)) || name.contains("[]");
618    }
619
620    private static com.getkeepsafe.dexcount.thrift.MethodRef methodRefToThrift(MethodRef methodRef) {
621        return new com.getkeepsafe.dexcount.thrift.MethodRef.Builder()
622            .declaringClass(methodRef.getDeclClassName())
623            .returnType(methodRef.getReturnTypeName())
624            .methodName(methodRef.getName())
625            .argumentTypes(Arrays.asList(methodRef.getArgumentTypeNames()))
626            .build();
627    }
628
629    private static MethodRef methodRefFromThrift(com.getkeepsafe.dexcount.thrift.MethodRef methodRef) {
630        String[] argTypes = new String[0];
631        if (methodRef.argumentTypes != null) {
632            argTypes = methodRef.argumentTypes.toArray(argTypes);
633        }
634
635        return new MethodRef(
636            methodRef.declaringClass,
637            argTypes,
638            methodRef.returnType,
639            methodRef.methodName
640        );
641    }
642
643    private static com.getkeepsafe.dexcount.thrift.FieldRef fieldRefToThrift(FieldRef fieldRef) {
644        return new com.getkeepsafe.dexcount.thrift.FieldRef.Builder()
645            .declaringClass(fieldRef.getDeclClassName())
646            .fieldType(fieldRef.getTypeName())
647            .fieldName(fieldRef.getName())
648            .build();
649    }
650
651    private static FieldRef fieldRefFromThrift(com.getkeepsafe.dexcount.thrift.FieldRef fieldRef) {
652        return new FieldRef(
653            fieldRef.declaringClass,
654            fieldRef.fieldType,
655            fieldRef.fieldName
656        );
657    }
658
659    public static com.getkeepsafe.dexcount.thrift.PackageTree toThrift(PackageTree tree) {
660        Map<String, com.getkeepsafe.dexcount.thrift.PackageTree> children = new LinkedHashMap<>();
661        for (Entry<String, PackageTree> entry : tree.children.entrySet()) {
662            children.put(entry.getKey(), toThrift(entry.getValue()));
663        }
664
665        Set<com.getkeepsafe.dexcount.thrift.MethodRef> thriftMethodDecls =
666            tree.methods.get(Type.DECLARED).stream().map(PackageTree::methodRefToThrift).collect(Collectors.toCollection(LinkedHashSet::new));
667        Set<com.getkeepsafe.dexcount.thrift.MethodRef> thriftMethodRefs =
668            tree.methods.get(Type.REFERENCED).stream().map(PackageTree::methodRefToThrift).collect(Collectors.toCollection(LinkedHashSet::new));
669        Set<com.getkeepsafe.dexcount.thrift.FieldRef> thriftFieldDecls =
670            tree.fields.get(Type.DECLARED).stream().map(PackageTree::fieldRefToThrift).collect(Collectors.toCollection(LinkedHashSet::new));
671        Set<com.getkeepsafe.dexcount.thrift.FieldRef> thriftFieldRefs =
672            tree.fields.get(Type.REFERENCED).stream().map(PackageTree::fieldRefToThrift).collect(Collectors.toCollection(LinkedHashSet::new));
673
674        return new com.getkeepsafe.dexcount.thrift.PackageTree.Builder()
675            .name(tree.getName())
676            .isClass(tree.isClass())
677            .children(children)
678            .declaredMethods(thriftMethodDecls)
679            .referencedMethods(thriftMethodRefs)
680            .declaredFields(thriftFieldDecls)
681            .referencedFields(thriftFieldRefs)
682            .build();
683    }
684
685    public static PackageTree fromThrift(com.getkeepsafe.dexcount.thrift.PackageTree tree) {
686        String name = tree.name != null ? tree.name : "";
687        boolean isClass = tree.isClass != null ? tree.isClass : false;
688
689        PackageTree result = new PackageTree(name, isClass, Deobfuscator.EMPTY);
690
691        if (tree.children != null) {
692            for (String key : tree.children.keySet()) {
693                result.children.put(key, fromThrift(tree.children.get(key)));
694            }
695        }
696
697        if (tree.declaredMethods != null) {
698            for (com.getkeepsafe.dexcount.thrift.MethodRef declaredMethod : tree.declaredMethods) {
699                result.methods.get(Type.DECLARED).add(methodRefFromThrift(declaredMethod));
700            }
701        }
702
703        if (tree.referencedMethods != null) {
704            for (com.getkeepsafe.dexcount.thrift.MethodRef referencedMethod : tree.referencedMethods) {
705                result.methods.get(Type.REFERENCED).add(methodRefFromThrift(referencedMethod));
706            }
707        }
708
709        if (tree.declaredFields != null) {
710            for (com.getkeepsafe.dexcount.thrift.FieldRef declaredField : tree.declaredFields) {
711                result.fields.get(Type.DECLARED).add(fieldRefFromThrift(declaredField));
712            }
713        }
714
715        if (tree.referencedFields != null) {
716            for (com.getkeepsafe.dexcount.thrift.FieldRef referencedField : tree.referencedFields) {
717                result.fields.get(Type.REFERENCED).add(fieldRefFromThrift(referencedField));
718            }
719        }
720
721        return result;
722    }
723}