From d82e8b42b2794431963819425d16000a0335bf85 Mon Sep 17 00:00:00 2001
From: Danny Lin <danny@kdrag0n.dev>
Date: Sun, 30 Jan 2022 03:30:50 -0800
Subject: [PATCH] extract: Use wrapped source

---
 src/commands/extract.ts | 63 +++++++++++++++++++++++++++++------------
 1 file changed, 45 insertions(+), 18 deletions(-)

diff --git a/src/commands/extract.ts b/src/commands/extract.ts
index 948e5b6..a88c254 100644
--- a/src/commands/extract.ts
+++ b/src/commands/extract.ts
@@ -4,7 +4,9 @@ import chalk from 'chalk'
 import { parseFileList } from '../blobs/file-list'
 import { copyBlobs } from '../blobs/copy'
 import { createVendorDirs, generateBuild, writeBuildFiles } from '../blobs/build'
-import { readFile } from '../util/fs'
+import { readFile, withTempDir } from '../util/fs'
+import { withSpinner } from '../util/cli'
+import { wrapSystemSrc } from '../frontend/source'
 
 export default class Extract extends Command {
   static description = 'extract proprietary files'
@@ -14,35 +16,60 @@ export default class Extract extends Command {
     vendor: flags.string({ char: 'v', description: 'device vendor/OEM name', required: true }),
     device: flags.string({ char: 'd', description: 'device codename', required: true }),
     skipCopy: flags.boolean({ char: 'k', description: 'skip file copying and only generate build files' }),
+
+    // Wrapped source
+    stockSrc: flags.string({
+      char: 's',
+      description:
+        'path to (extracted) factory images, (mounted) images, (extracted) OTA package, OTA payload, or directory containing any such files (optionally under device and/or build ID directory)',
+      required: true,
+    }),
+    buildId: flags.string({
+      char: 'b',
+      description: 'build ID of the stock images (optional, only used for locating factory images)',
+    }),
+    useTemp: flags.boolean({
+      char: 't',
+      description: 'use a temporary directory for all extraction (prevents reusing extracted files across runs)',
+      default: false,
+    }),
   }
 
   static args = [
-    { name: 'source', description: 'path to mounted factory images', required: true },
     { name: 'listPath', description: 'path to LineageOS-compatible proprietary-files.txt list', required: true },
   ]
 
   async run() {
     let {
-      args: { source, listPath },
-      flags: { vendor, device, skipCopy },
+      args: { listPath },
+      flags: { vendor, device, skipCopy, stockSrc, buildId, useTemp },
     } = this.parse(Extract)
 
-    // Parse list
-    this.log(chalk.bold(chalk.greenBright('Parsing list')))
-    let list = await readFile(listPath)
-    let entries = parseFileList(list)
+    await withTempDir(async tmp => {
+      // Prepare stock system source
+      let wrapBuildId = buildId == undefined ? null : buildId
+      let wrapped = await withSpinner('Extracting stock system source', spinner =>
+        wrapSystemSrc(stockSrc, device, wrapBuildId, useTemp, tmp, spinner),
+      )
+      stockSrc = wrapped.src!
 
-    // Prepare output directories
-    let dirs = await createVendorDirs(vendor, device)
+      // Parse list
+      this.log(chalk.bold(chalk.greenBright('Parsing list')))
+      let list = await readFile(listPath)
+      let entries = parseFileList(list)
 
-    // Copy blobs
-    if (!skipCopy) {
-      await copyBlobs(entries, source, dirs.proprietary)
-    }
+      // Prepare output directories
+      let dirs = await createVendorDirs(vendor, device)
 
-    // Generate build files
-    this.log(chalk.bold(chalk.greenBright('Generating build files')))
-    let build = await generateBuild(entries, device, vendor, source, dirs)
-    await writeBuildFiles(build, dirs)
+      // Copy blobs
+      if (!skipCopy) {
+        await copyBlobs(entries, stockSrc, dirs.proprietary)
+      }
+
+      // Generate build files
+      this.log(chalk.bold(chalk.greenBright('Generating build files')))
+      let build = await generateBuild(entries, device, vendor, stockSrc, dirs)
+      await writeBuildFiles(build, dirs)
+    })
   }
 }