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.KtApiKt;
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    private File outputDirectory = null;
058
059    @Override
060    public void execute() {
061        try {
062            PackageTree packageTree = generatePackageTree();
063
064            ensureCleanOutputDirectory();
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 ensureCleanOutputDirectory() throws IOException {
076        FileUtils.deleteDirectory(getOutputDirectory());
077        FileUtils.forceMkdir(getOutputDirectory());
078    }
079
080    private void writeIntermediateThriftFile(PackageTree packageTree) throws IOException {
081        TreeGenOutput thrift = new TreeGenOutput.Builder()
082            .tree(PackageTree.toThrift(packageTree))
083            .inputRepresentation(getInputRepresentation())
084            .build();
085
086        File treeFile = getParameters().getPackageTreeFile().getAsFile().get();
087        FileUtils.deleteQuietly(treeFile);
088
089        try (Sink fileSink = Okio.sink(treeFile);
090             Sink gzipSink = new GzipSink(fileSink);
091             BufferedSink sink = Okio.buffer(gzipSink);
092             Transport transport = KtApiKt.transport(sink);
093             Protocol protocol = KtApiKt.compactProtocol(transport)) {
094            TreeGenOutput.ADAPTER.write(protocol, thrift);
095            protocol.flush();
096        }
097    }
098
099    private void writeSummaryFile(PackageTree packageTree) throws IOException {
100        File summaryFile = new File(getOutputDirectory(), "summary.csv");
101        FileUtils.forceMkdirParent(summaryFile);
102
103        String headers = "methods,fields,classes";
104        String counts = String.format(
105            "%d,%d,%d",
106            packageTree.getMethodCount(),
107            packageTree.getFieldCount(),
108            packageTree.getClassCount());
109
110        try (BufferedWriter writer = Files.newBufferedWriter(summaryFile.toPath())) {
111            writer.append(headers).append('\n');
112            writer.append(counts).append('\n');
113        }
114    }
115
116    private void writeChartFiles(PackageTree packageTree) throws IOException {
117        File chartDirectory = new File(getOutputDirectory(), "chart");
118        FileUtils.forceMkdir(chartDirectory);
119
120        PrintOptions options = getParameters().getPrintOptions().get()
121            .toBuilder()
122            .setIncludeClasses(true)
123            .build();
124
125        File dataJs = new File(chartDirectory, "data.js");
126        try (BufferedWriter out = Files.newBufferedWriter(dataJs.toPath())) {
127            out.write("var data = ");
128            packageTree.printJson(out, options);
129        }
130
131        List<String> resourceNames = Arrays.asList("chart-builder.js", "d3.v3.min.js", "index.html", "styles.css");
132        for (String resourceName : resourceNames) {
133            String resourcePath = "com/getkeepsafe/dexcount/" + resourceName;
134            try (InputStream is = DexMethodCountPlugin.class.getClassLoader().getResourceAsStream(resourcePath)) {
135                if (is == null) {
136                    getLogger().error("No such resource: {}", resourcePath);
137                    continue;
138                }
139                File target = new File(chartDirectory, resourceName);
140                FileUtils.copyInputStreamToFile(is, target);
141            }
142        }
143    }
144
145    private void writeFullTree(PackageTree packageTree) throws IOException {
146        PrintOptions options = getParameters().getPrintOptions().get();
147        String fullCountFileName = getParameters().getOutputFileName().get() + options.getOutputFormat().getExtension();
148        File fullCountFile = new File(getOutputDirectory(), fullCountFileName);
149
150        try (BufferedWriter bw = Files.newBufferedWriter(fullCountFile.toPath())) {
151            packageTree.print(bw, options.getOutputFormat(), options);
152        }
153    }
154
155    private File getOutputDirectory() {
156        if (outputDirectory == null) {
157            outputDirectory = getParameters().getOutputDirectory().get().getAsFile();
158        }
159        return outputDirectory;
160    }
161
162    protected abstract PackageTree generatePackageTree() throws IOException;
163
164    protected abstract String getInputRepresentation();
165
166    protected abstract Logger getLogger();
167}