From 54cce1b8f6fd28135f446e5ec2b591ecd3a9a600 Mon Sep 17 00:00:00 2001
From: Danny Lin <danny@kdrag0n.dev>
Date: Thu, 16 Dec 2021 19:07:31 -0800
Subject: [PATCH] generate: Accept flexible source paths

---
 src/commands/generate-all.ts  | 9 ++++++++-
 src/commands/generate-prep.ts | 9 ++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/commands/generate-all.ts b/src/commands/generate-all.ts
index b3b10c8..9aae900 100644
--- a/src/commands/generate-all.ts
+++ b/src/commands/generate-all.ts
@@ -7,6 +7,7 @@ import { BlobEntry } from '../blobs/entry'
 import { loadDeviceConfig } from '../config/device'
 import { collectSystemState, parseSystemState, SystemState } from '../config/system-state'
 import { enumerateFiles, extractFirmware, extractOverlays, extractProps, extractVintfManifests, flattenApexs, generateBuildFiles, PropResults, resolveOverrides, resolveSepolicyDirs, updatePresigned } from '../frontend/generate'
+import { wrapSystemSrc } from '../frontend/source'
 import { SelinuxPartResolutions } from '../selinux/contexts'
 import { withSpinner } from '../util/cli'
 import { readFile, withTempDir } from '../util/fs'
@@ -17,6 +18,7 @@ export default class GenerateFull extends Command {
   static flags = {
     help: flags.help({char: 'h'}),
     aapt2: flags.string({char: 'a', description: 'path to aapt2 executable', default: 'out/host/linux-x86/bin/aapt2'}),
+    buildId: flags.string({char: 'b', description: 'build ID of the stock images'}),
     stockSrc: flags.string({char: 's', description: 'path to root of mounted stock system images (./system_ext, ./product, etc.)', required: true}),
     customSrc: flags.string({char: 'c', description: 'path to AOSP build output directory (out/) or JSON state file', default: 'out'}),
     factoryZip: flags.string({char: 'f', description: 'path to stock factory images zip (for extracting firmware)'}),
@@ -28,7 +30,7 @@ export default class GenerateFull extends Command {
   ]
 
   async run() {
-    let {flags: {aapt2: aapt2Path, stockSrc, customSrc, factoryZip, skipCopy}, args: {config: configPath}} = this.parse(GenerateFull)
+    let {flags: {aapt2: aapt2Path, buildId, stockSrc, customSrc, factoryZip, skipCopy}, args: {config: configPath}} = this.parse(GenerateFull)
 
     let config = await loadDeviceConfig(configPath)
 
@@ -42,6 +44,11 @@ export default class GenerateFull extends Command {
 
     // tmp may be needed for mounting/extracting stock images and/or APEX flattening
     await withTempDir(async (tmp) => {
+      // Prepare stock system source
+      let wrapBuildId = buildId == undefined ? null : buildId
+      stockSrc = await withSpinner('Extracting stock system source', (spinner) =>
+        wrapSystemSrc(stockSrc, config.device.name, wrapBuildId, tmp, spinner))
+
       // Each step will modify this. Key = combined part path
       let namedEntries = new Map<string, BlobEntry>()
 
diff --git a/src/commands/generate-prep.ts b/src/commands/generate-prep.ts
index ed80264..16b279d 100644
--- a/src/commands/generate-prep.ts
+++ b/src/commands/generate-prep.ts
@@ -5,6 +5,7 @@ import { copyBlobs } from '../blobs/copy'
 import { BlobEntry } from '../blobs/entry'
 import { loadDeviceConfig } from '../config/device'
 import { enumerateFiles, extractProps, generateBuildFiles, PropResults } from '../frontend/generate'
+import { wrapSystemSrc } from '../frontend/source'
 import { withSpinner } from '../util/cli'
 import { withTempDir } from '../util/fs'
 
@@ -13,6 +14,7 @@ export default class GeneratePrep extends Command {
 
   static flags = {
     help: flags.help({char: 'h'}),
+    buildId: flags.string({char: 'b', description: 'build ID of the stock images'}),
     stockSrc: flags.string({char: 's', description: 'path to root of mounted stock system images (./system_ext, ./product, etc.)', required: true}),
     skipCopy: flags.boolean({char: 'k', description: 'skip file copying and only generate build files'}),
   }
@@ -22,12 +24,17 @@ export default class GeneratePrep extends Command {
   ]
 
   async run() {
-    let {flags: {stockSrc, skipCopy}, args: {config: configPath}} = this.parse(GeneratePrep)
+    let {flags: {buildId, stockSrc, skipCopy}, args: {config: configPath}} = this.parse(GeneratePrep)
 
     let config = await loadDeviceConfig(configPath)
 
     // tmp may be needed for mounting/extracting stock images
     await withTempDir(async (tmp) => {
+      // Prepare stock system source
+      let wrapBuildId = buildId == undefined ? null : buildId
+      stockSrc = await withSpinner('Extracting stock system source', (spinner) =>
+        wrapSystemSrc(stockSrc, config.device.name, wrapBuildId, tmp, spinner))
+
       // Each step will modify this. Key = combined part path
       let namedEntries = new Map<string, BlobEntry>()