From 3803585e716ea796a4108ffb9acdb25d4ed3c202 Mon Sep 17 00:00:00 2001
From: Danny Lin <danny@kdrag0n.dev>
Date: Sat, 18 Dec 2021 18:58:30 -0800
Subject: [PATCH] fix-keys: Add support for wrapping stock sources

---
 src/commands/fix-keys.ts | 49 ++++++++++++++++++++++++----------------
 1 file changed, 30 insertions(+), 19 deletions(-)

diff --git a/src/commands/fix-keys.ts b/src/commands/fix-keys.ts
index 6351121..2c3a59f 100644
--- a/src/commands/fix-keys.ts
+++ b/src/commands/fix-keys.ts
@@ -1,6 +1,9 @@
 import { Command, flags } from '@oclif/command'
+import { wrapSystemSrc } from '../frontend/source'
 
 import { KeyInfo, MacSigner, readKeysConfRecursive, readMacPermissionsRecursive, readPartMacPermissions, resolveKeys, writeMappedKeys } from '../selinux/keys'
+import { withSpinner } from '../util/cli'
+import { withTempDir } from '../util/fs'
 
 export default class FixKeys extends Command {
   static description = 'fix SELinux presigned app keys'
@@ -8,31 +11,39 @@ export default class FixKeys extends Command {
   static flags = {
     help: flags.help({char: 'h'}),
     sepolicy: flags.string({char: 'p', description: 'paths to device and vendor sepolicy dirs', required: true, multiple: true}),
+    device: flags.string({char: 'd', description: 'device codename', required: true}),
+    buildId: flags.string({char: 'b', description: 'build ID of the stock images (optional, only used for locating factory images)'}),
+    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}),
+    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},
-  ]
-
   async run() {
-    let {flags: {sepolicy: sepolicyDirs}, args: {source}} = this.parse(FixKeys)
+    let {flags: {sepolicy: sepolicyDirs, device, buildId, stockSrc, useTemp}} = this.parse(FixKeys)
 
-    let srcSigners: Array<MacSigner> = []
-    let srcKeys: Array<KeyInfo> = []
-    for (let dir of sepolicyDirs) {
-      srcSigners.push(...(await readMacPermissionsRecursive(dir)))
-      srcKeys.push(...(await readKeysConfRecursive(dir)))
-    }
+    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!
 
-    let compiledSigners = await readPartMacPermissions(source)
-    let keys = resolveKeys(srcKeys, srcSigners, compiledSigners)
-
-    for (let paths of keys.values()) {
-      for (let path of paths) {
-        this.log(path)
+      let srcSigners: Array<MacSigner> = []
+      let srcKeys: Array<KeyInfo> = []
+      for (let dir of sepolicyDirs) {
+        srcSigners.push(...(await readMacPermissionsRecursive(dir)))
+        srcKeys.push(...(await readKeysConfRecursive(dir)))
       }
-    }
 
-    await writeMappedKeys(keys)
+      let compiledSigners = await readPartMacPermissions(stockSrc)
+      let keys = resolveKeys(srcKeys, srcSigners, compiledSigners)
+
+      for (let paths of keys.values()) {
+        for (let path of paths) {
+          this.log(path)
+        }
+      }
+
+      await writeMappedKeys(keys)
+    })
   }
 }