Implement extract_apks

Bug: 152319766
Test: manual and builtin
Change-Id: Ia15d66e86c7bcfd52f5b776173ca1665b68ff438
Merged-In: Ia15d66e86c7bcfd52f5b776173ca1665b68ff438
diff --git a/cmd/extract_apks/main.go b/cmd/extract_apks/main.go
new file mode 100644
index 0000000..4a146da
--- /dev/null
+++ b/cmd/extract_apks/main.go
@@ -0,0 +1,467 @@
+// Copyright 2020 Google Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Copies all the entries (APKs/APEXes) matching the target configuration from the given
+// APK set into a zip file. Run it without arguments to see usage details.
+package main
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"regexp"
+	"strings"
+
+	"github.com/golang/protobuf/proto"
+
+	"android/soong/cmd/extract_apks/bundle_proto"
+	"android/soong/third_party/zip"
+)
+
+type TargetConfig struct {
+	sdkVersion       int32
+	screenDpi        map[android_bundle_proto.ScreenDensity_DensityAlias]bool
+	abis             map[android_bundle_proto.Abi_AbiAlias]bool
+	allowPrereleased bool
+	stem             string
+}
+
+// An APK set is a zip archive. An entry 'toc.pb' describes its contents.
+// It is a protobuf message BuildApkResult.
+type Toc *android_bundle_proto.BuildApksResult
+
+type ApkSet struct {
+	path    string
+	reader  *zip.ReadCloser
+	entries map[string]*zip.File
+}
+
+func newApkSet(path string) (*ApkSet, error) {
+	apkSet := &ApkSet{path: path, entries: make(map[string]*zip.File)}
+	var err error
+	if apkSet.reader, err = zip.OpenReader(apkSet.path); err != nil {
+		return nil, err
+	}
+	for _, f := range apkSet.reader.File {
+		apkSet.entries[f.Name] = f
+	}
+	return apkSet, nil
+}
+
+func (apkSet *ApkSet) getToc() (Toc, error) {
+	var err error
+	tocFile, ok := apkSet.entries["toc.pb"]
+	if !ok {
+		return nil, fmt.Errorf("%s: APK set should have toc.pb entry", apkSet.path)
+	}
+	rc, err := tocFile.Open()
+	if err != nil {
+		return nil, err
+	}
+	bytes := make([]byte, tocFile.FileHeader.UncompressedSize64)
+	if _, err := rc.Read(bytes); err != io.EOF {
+		return nil, err
+	}
+	rc.Close()
+	buildApksResult := new(android_bundle_proto.BuildApksResult)
+	if err = proto.Unmarshal(bytes, buildApksResult); err != nil {
+		return nil, err
+	}
+	return buildApksResult, nil
+}
+
+func (apkSet *ApkSet) close() {
+	apkSet.reader.Close()
+}
+
+// Matchers for selection criteria
+type abiTargetingMatcher struct {
+	*android_bundle_proto.AbiTargeting
+}
+
+func (m abiTargetingMatcher) matches(config TargetConfig) bool {
+	if m.AbiTargeting == nil {
+		return true
+	}
+	if _, ok := config.abis[android_bundle_proto.Abi_UNSPECIFIED_CPU_ARCHITECTURE]; ok {
+		return true
+	}
+	for _, v := range m.GetValue() {
+		if _, ok := config.abis[v.Alias]; ok {
+			return true
+		}
+	}
+	return false
+}
+
+type apkDescriptionMatcher struct {
+	*android_bundle_proto.ApkDescription
+}
+
+func (m apkDescriptionMatcher) matches(config TargetConfig) bool {
+	return m.ApkDescription == nil || (apkTargetingMatcher{m.Targeting}).matches(config)
+}
+
+type apkTargetingMatcher struct {
+	*android_bundle_proto.ApkTargeting
+}
+
+func (m apkTargetingMatcher) matches(config TargetConfig) bool {
+	return m.ApkTargeting == nil ||
+		(abiTargetingMatcher{m.AbiTargeting}.matches(config) &&
+			languageTargetingMatcher{m.LanguageTargeting}.matches(config) &&
+			screenDensityTargetingMatcher{m.ScreenDensityTargeting}.matches(config) &&
+			sdkVersionTargetingMatcher{m.SdkVersionTargeting}.matches(config) &&
+			multiAbiTargetingMatcher{m.MultiAbiTargeting}.matches(config))
+}
+
+type languageTargetingMatcher struct {
+	*android_bundle_proto.LanguageTargeting
+}
+
+func (m languageTargetingMatcher) matches(_ TargetConfig) bool {
+	if m.LanguageTargeting == nil {
+		return true
+	}
+	log.Fatal("language based entry selection is not implemented")
+	return false
+}
+
+type moduleMetadataMatcher struct {
+	*android_bundle_proto.ModuleMetadata
+}
+
+func (m moduleMetadataMatcher) matches(config TargetConfig) bool {
+	return m.ModuleMetadata == nil ||
+		(m.GetDeliveryType() == android_bundle_proto.DeliveryType_INSTALL_TIME &&
+			moduleTargetingMatcher{m.Targeting}.matches(config) &&
+			!m.IsInstant)
+}
+
+type moduleTargetingMatcher struct {
+	*android_bundle_proto.ModuleTargeting
+}
+
+func (m moduleTargetingMatcher) matches(config TargetConfig) bool {
+	return m.ModuleTargeting == nil ||
+		(sdkVersionTargetingMatcher{m.SdkVersionTargeting}.matches(config) &&
+			userCountriesTargetingMatcher{m.UserCountriesTargeting}.matches(config))
+}
+
+type multiAbiTargetingMatcher struct {
+	*android_bundle_proto.MultiAbiTargeting
+}
+
+func (t multiAbiTargetingMatcher) matches(_ TargetConfig) bool {
+	if t.MultiAbiTargeting == nil {
+		return true
+	}
+	log.Fatal("multiABI based selection is not implemented")
+	return false
+}
+
+type screenDensityTargetingMatcher struct {
+	*android_bundle_proto.ScreenDensityTargeting
+}
+
+func (m screenDensityTargetingMatcher) matches(config TargetConfig) bool {
+	if m.ScreenDensityTargeting == nil {
+		return true
+	}
+	if _, ok := config.screenDpi[android_bundle_proto.ScreenDensity_DENSITY_UNSPECIFIED]; ok {
+		return true
+	}
+	for _, v := range m.GetValue() {
+		switch x := v.GetDensityOneof().(type) {
+		case *android_bundle_proto.ScreenDensity_DensityAlias_:
+			if _, ok := config.screenDpi[x.DensityAlias]; ok {
+				return true
+			}
+		default:
+			log.Fatal("For screen density, only DPI name based entry selection (e.g. HDPI, XHDPI) is implemented")
+		}
+	}
+	return false
+}
+
+type sdkVersionTargetingMatcher struct {
+	*android_bundle_proto.SdkVersionTargeting
+}
+
+func (m sdkVersionTargetingMatcher) matches(config TargetConfig) bool {
+	const preReleaseVersion = 10000
+	if m.SdkVersionTargeting == nil {
+		return true
+	}
+	if len(m.Value) > 1 {
+		log.Fatal(fmt.Sprintf("sdk_version_targeting should not have multiple values:%#v", m.Value))
+	}
+	// Inspect only sdkVersionTargeting.Value.
+	// Even though one of the SdkVersionTargeting.Alternatives values may be
+	// better matching, we will select all of them
+	return m.Value[0].Min == nil ||
+		m.Value[0].Min.Value <= config.sdkVersion ||
+		(config.allowPrereleased && m.Value[0].Min.Value == preReleaseVersion)
+}
+
+type textureCompressionFormatTargetingMatcher struct {
+	*android_bundle_proto.TextureCompressionFormatTargeting
+}
+
+func (m textureCompressionFormatTargetingMatcher) matches(_ TargetConfig) bool {
+	if m.TextureCompressionFormatTargeting == nil {
+		return true
+	}
+	log.Fatal("texture based entry selection is not implemented")
+	return false
+}
+
+type userCountriesTargetingMatcher struct {
+	*android_bundle_proto.UserCountriesTargeting
+}
+
+func (m userCountriesTargetingMatcher) matches(_ TargetConfig) bool {
+	if m.UserCountriesTargeting == nil {
+		return true
+	}
+	log.Fatal("country based entry selection is not implemented")
+	return false
+}
+
+type variantTargetingMatcher struct {
+	*android_bundle_proto.VariantTargeting
+}
+
+func (m variantTargetingMatcher) matches(config TargetConfig) bool {
+	if m.VariantTargeting == nil {
+		return true
+	}
+	return sdkVersionTargetingMatcher{m.SdkVersionTargeting}.matches(config) &&
+		abiTargetingMatcher{m.AbiTargeting}.matches(config) &&
+		multiAbiTargetingMatcher{m.MultiAbiTargeting}.matches(config) &&
+		screenDensityTargetingMatcher{m.ScreenDensityTargeting}.matches(config) &&
+		textureCompressionFormatTargetingMatcher{m.TextureCompressionFormatTargeting}.matches(config)
+}
+
+type SelectionResult struct {
+	moduleName string
+	entries    []string
+}
+
+// Return all entries matching target configuration
+func selectApks(toc Toc, targetConfig TargetConfig) SelectionResult {
+	var result SelectionResult
+	for _, variant := range (*toc).GetVariant() {
+		if !(variantTargetingMatcher{variant.GetTargeting()}.matches(targetConfig)) {
+			continue
+		}
+		for _, as := range variant.GetApkSet() {
+			if !(moduleMetadataMatcher{as.ModuleMetadata}.matches(targetConfig)) {
+				continue
+			}
+			for _, apkdesc := range as.GetApkDescription() {
+				if (apkDescriptionMatcher{apkdesc}).matches(targetConfig) {
+					result.entries = append(result.entries, apkdesc.GetPath())
+					// TODO(asmundak): As it turns out, moduleName which we get from
+					// the ModuleMetadata matches the module names of the generated
+					// entry paths just by coincidence, only for the split APKs. We
+					// need to discuss this with bundletool folks.
+					result.moduleName = as.GetModuleMetadata().GetName()
+				}
+			}
+			// we allow only a single module, so bail out here if we found one
+			if result.moduleName != "" {
+				return result
+			}
+		}
+	}
+	return result
+}
+
+type Zip2ZipWriter interface {
+	CopyFrom(file *zip.File, name string) error
+}
+
+// Writes out selected entries, renaming them as needed
+func (apkSet *ApkSet) writeApks(selected SelectionResult, config TargetConfig,
+	writer Zip2ZipWriter) error {
+	// Renaming rules:
+	//  splits/MODULE-master.apk to STEM.apk
+	// else
+	//  splits/MODULE-*.apk to STEM>-$1.apk
+	// TODO(asmundak):
+	//  add more rules, for .apex files
+	renameRules := []struct {
+		rex  *regexp.Regexp
+		repl string
+	}{
+		{
+			regexp.MustCompile(`^.*/` + selected.moduleName + `-master\.apk$`),
+			config.stem + `.apk`,
+		},
+		{
+			regexp.MustCompile(`^.*/` + selected.moduleName + `(-.*\.apk)$`),
+			config.stem + `$1`,
+		},
+	}
+	renamer := func(path string) (string, bool) {
+		for _, rr := range renameRules {
+			if rr.rex.MatchString(path) {
+				return rr.rex.ReplaceAllString(path, rr.repl), true
+			}
+		}
+		return "", false
+	}
+
+	entryOrigin := make(map[string]string) // output entry to input entry
+	for _, apk := range selected.entries {
+		apkFile, ok := apkSet.entries[apk]
+		if !ok {
+			return fmt.Errorf("TOC refers to an entry %s which does not exist", apk)
+		}
+		inName := apkFile.Name
+		outName, ok := renamer(inName)
+		if !ok {
+			log.Fatalf("selected an entry with unexpected name %s", inName)
+		}
+		if origin, ok := entryOrigin[inName]; ok {
+			log.Fatalf("selected entries %s and %s will have the same output name %s",
+				origin, inName, outName)
+		}
+		entryOrigin[outName] = inName
+		if err := writer.CopyFrom(apkFile, outName); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// Arguments parsing
+var (
+	outputZip    = flag.String("o", "", "output zip containing extracted entries")
+	targetConfig = TargetConfig{
+		screenDpi: map[android_bundle_proto.ScreenDensity_DensityAlias]bool{},
+		abis:      map[android_bundle_proto.Abi_AbiAlias]bool{},
+	}
+)
+
+// Parse abi values
+type abiFlagValue struct {
+	targetConfig *TargetConfig
+}
+
+func (a abiFlagValue) String() string {
+	return "all"
+}
+
+func (a abiFlagValue) Set(abiList string) error {
+	if abiList == "none" {
+		return nil
+	}
+	if abiList == "all" {
+		targetConfig.abis[android_bundle_proto.Abi_UNSPECIFIED_CPU_ARCHITECTURE] = true
+		return nil
+	}
+	for _, abi := range strings.Split(abiList, ",") {
+		v, ok := android_bundle_proto.Abi_AbiAlias_value[abi]
+		if !ok {
+			return fmt.Errorf("bad ABI value: %q", abi)
+		}
+		targetConfig.abis[android_bundle_proto.Abi_AbiAlias(v)] = true
+	}
+	return nil
+}
+
+// Parse screen density values
+type screenDensityFlagValue struct {
+	targetConfig *TargetConfig
+}
+
+func (s screenDensityFlagValue) String() string {
+	return "none"
+}
+
+func (s screenDensityFlagValue) Set(densityList string) error {
+	if densityList == "none" {
+		return nil
+	}
+	if densityList == "all" {
+		targetConfig.screenDpi[android_bundle_proto.ScreenDensity_DENSITY_UNSPECIFIED] = true
+		return nil
+	}
+	for _, density := range strings.Split(densityList, ",") {
+		v, found := android_bundle_proto.ScreenDensity_DensityAlias_value[density]
+		if !found {
+			return fmt.Errorf("bad screen density value: %q", density)
+		}
+		targetConfig.screenDpi[android_bundle_proto.ScreenDensity_DensityAlias(v)] = true
+	}
+	return nil
+}
+
+func processArgs() {
+	flag.Usage = func() {
+		fmt.Fprintln(os.Stderr, `usage: extract_apks -o <output-zip> -sdk-version value -abis value -screen-densities value  <APK set>`)
+		flag.PrintDefaults()
+		os.Exit(2)
+	}
+	version := flag.Uint("sdk-version", 0, "SDK version")
+	flag.Var(abiFlagValue{&targetConfig}, "abis",
+		"'all' or comma-separated ABIs list of ARMEABI ARMEABI_V7A ARM64_V8A X86 X86_64 MIPS MIPS64")
+	flag.Var(screenDensityFlagValue{&targetConfig}, "screen-densities",
+		"'all' or comma-separated list of screen density names (NODPI LDPI MDPI TVDPI HDPI XHDPI XXHDPI XXXHDPI)")
+	flag.BoolVar(&targetConfig.allowPrereleased, "allow-prereleased", false,
+		"allow prereleased")
+	flag.StringVar(&targetConfig.stem, "stem", "", "output entries base name")
+	flag.Parse()
+	if (*outputZip == "") || len(flag.Args()) != 1 || *version == 0 || targetConfig.stem == "" {
+		flag.Usage()
+	}
+	targetConfig.sdkVersion = int32(*version)
+
+}
+
+func main() {
+	processArgs()
+	var toc Toc
+	apkSet, err := newApkSet(flag.Arg(0))
+	if err == nil {
+		defer apkSet.close()
+		toc, err = apkSet.getToc()
+	}
+	if err != nil {
+		log.Fatal(err)
+	}
+	sel := selectApks(toc, targetConfig)
+	if len(sel.entries) == 0 {
+		log.Fatalf("there are no entries for the target configuration: %#v", targetConfig)
+	}
+
+	outFile, err := os.Create(*outputZip)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer outFile.Close()
+	writer := zip.NewWriter(outFile)
+	defer func() {
+		if err := writer.Close(); err != nil {
+			log.Fatal(err)
+		}
+	}()
+	if err = apkSet.writeApks(sel, targetConfig, writer); err != nil {
+		log.Fatal(err)
+	}
+}