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.treegen.workers;
017
018import com.getkeepsafe.dexcount.DexCountException;
019import com.getkeepsafe.dexcount.DexMethodCountPlugin;
020import com.getkeepsafe.dexcount.PackageTree;
021import com.getkeepsafe.dexcount.PrintOptions;
022import com.getkeepsafe.dexcount.thrift.TreeGenOutput;
023import com.microsoft.thrifty.protocol.CompactProtocol;
024import com.microsoft.thrifty.protocol.Protocol;
025import com.microsoft.thrifty.transport.Transport;
026import okio.BufferedSink;
027import okio.GzipSink;
028import okio.Okio;
029import okio.Sink;
030import org.apache.commons.io.FileUtils;
031import org.gradle.api.file.DirectoryProperty;
032import org.gradle.api.file.RegularFileProperty;
033import org.gradle.api.provider.Property;
034import org.gradle.workers.WorkAction;
035import org.gradle.workers.WorkParameters;
036import org.slf4j.Logger;
037
038import java.io.BufferedWriter;
039import java.io.File;
040import java.io.IOException;
041import java.io.InputStream;
042import java.nio.file.Files;
043import java.util.Arrays;
044import java.util.List;
045
046public abstract class BaseWorker<P extends BaseWorker.Params> implements WorkAction<P> {
047    public interface Params extends WorkParameters {
048        Property<String> getOutputFileName();
049
050        RegularFileProperty getPackageTreeFile();
051
052        DirectoryProperty getOutputDirectory();
053
054        Property<PrintOptions> getPrintOptions();
055    }
056
057    @Override
058    public void execute() {
059        try {
060            PackageTree packageTree = generatePackageTree();
061
062            File outputDir = getParameters().getOutputDirectory().get().getAsFile();
063            FileUtils.deleteDirectory(outputDir);
064            FileUtils.forceMkdir(outputDir);
065
066            writeIntermediateThriftFile(packageTree);
067            writeSummaryFile(packageTree);
068            writeChartFiles(packageTree);
069            writeFullTree(packageTree);
070        } catch (IOException e) {
071            throw new DexCountException("Counting dex method references failed", e);
072        }
073    }
074
075    private void writeIntermediateThriftFile(PackageTree packageTree) throws IOException {
076        TreeGenOutput thrift = new TreeGenOutput.Builder()
077            .tree(PackageTree.toThrift(packageTree))
078            .inputRepresentation(getInputRepresentation())
079            .build();
080
081        File treeFile = getParameters().getPackageTreeFile().getAsFile().get();
082        FileUtils.deleteQuietly(treeFile);
083
084        try (Sink fileSink = Okio.sink(treeFile);
085             Sink gzipSink = new GzipSink(fileSink);
086             BufferedSink sink = Okio.buffer(gzipSink);
087             Transport transport = new BufferedSinkTransport(sink);
088             Protocol protocol = new CompactProtocol(transport)) {
089            TreeGenOutput.ADAPTER.write(protocol, thrift);
090            protocol.flush();
091        }
092    }
093
094    private void writeSummaryFile(PackageTree packageTree) throws IOException {
095        File summaryFile = getParameters().getOutputDirectory().file("summary.csv").get().getAsFile();
096        FileUtils.forceMkdirParent(summaryFile);
097
098        String headers = "methods,fields,classes";
099        String counts = String.format(
100            "%d,%d,%d",
101            packageTree.getMethodCount(),
102            packageTree.getFieldCount(),
103            packageTree.getClassCount());
104
105        try (BufferedWriter writer = Files.newBufferedWriter(summaryFile.toPath())) {
106            writer.append(headers).append('\n');
107            writer.append(counts).append('\n');
108        }
109    }
110
111    private void writeChartFiles(PackageTree packageTree) throws IOException {
112        File chartDirectory = getParameters().getOutputDirectory().dir("chart").get().getAsFile();
113        FileUtils.forceMkdir(chartDirectory);
114
115        PrintOptions options = getParameters().getPrintOptions().get()
116            .toBuilder()
117            .setIncludeClasses(true)
118            .build();
119
120        File dataJs = new File(chartDirectory, "data.js");
121        try (BufferedWriter out = Files.newBufferedWriter(dataJs.toPath())) {
122            out.write("var data = ");
123            packageTree.printJson(out, options);
124        }
125
126        List<String> resourceNames = Arrays.asList("chart-builder.js", "d3.v3.min.js", "index.html", "styles.css");
127        for (String resourceName : resourceNames) {
128            String resourcePath = "com/getkeepsafe/dexcount/" + resourceName;
129            try (InputStream is = DexMethodCountPlugin.class.getClassLoader().getResourceAsStream(resourcePath)) {
130                if (is == null) {
131                    getLogger().error("No such resource: {}", resourcePath);
132                    continue;
133                }
134                File target = new File(chartDirectory, resourceName);
135                FileUtils.copyInputStreamToFile(is, target);
136            }
137        }
138    }
139
140    private void writeFullTree(PackageTree packageTree) throws IOException {
141        PrintOptions options = getParameters().getPrintOptions().get();
142        String fullCountFileName = getParameters().getOutputFileName().get() + options.getOutputFormat().getExtension();
143        File fullCountFile = getParameters().getOutputDirectory().file(fullCountFileName).get().getAsFile();
144
145        try (BufferedWriter bw = Files.newBufferedWriter(fullCountFile.toPath())) {
146            packageTree.print(bw, options.getOutputFormat(), options);
147        }
148    }
149
150    protected abstract PackageTree generatePackageTree() throws IOException;
151
152    protected abstract String getInputRepresentation();
153
154    protected abstract Logger getLogger();
155
156    private static class BufferedSinkTransport extends Transport {
157        private final BufferedSink sink;
158
159        BufferedSinkTransport(BufferedSink sink) {
160            this.sink = sink;
161        }
162
163        @Override
164        public int read(byte[] buffer, int offset, int count) throws IOException {
165            throw new UnsupportedOperationException();
166        }
167
168        @Override
169        public void write(byte[] buffer, int offset, int count) throws IOException {
170            sink.write(buffer, offset, count);
171        }
172
173        @Override
174        public void flush() throws IOException {
175            sink.flush();
176        }
177
178        @Override
179        public void close() throws IOException {
180            sink.close();
181        }
182    }
183}