From 03db6352bd72d8d15babd34f718d583e1b4fc6a7 Mon Sep 17 00:00:00 2001
From: Danny Lin <danny@kdrag0n.dev>
Date: Sun, 7 Nov 2021 17:25:25 -0800
Subject: [PATCH] Add command to download device images

---
 package.json             |  1 +
 src/commands/download.ts | 86 ++++++++++++++++++++++++++++++++++++++++
 yarn.lock                |  2 +-
 3 files changed, 88 insertions(+), 1 deletion(-)
 create mode 100644 src/commands/download.ts

diff --git a/package.json b/package.json
index 44a8c65..33c3581 100644
--- a/package.json
+++ b/package.json
@@ -13,6 +13,7 @@
     "@oclif/plugin-help": "^3",
     "chalk": "^4.1.2",
     "cli-progress": "^3.9.1",
+    "node-fetch": "2",
     "ora": "5.4.1",
     "tslib": "^1",
     "zx": "^4.2.0"
diff --git a/src/commands/download.ts b/src/commands/download.ts
new file mode 100644
index 0000000..c12cd53
--- /dev/null
+++ b/src/commands/download.ts
@@ -0,0 +1,86 @@
+import {Command, flags} from '@oclif/command'
+import {createWriteStream, promises as fs} from 'fs'
+import * as chalk from 'chalk'
+import * as cliProgress from 'cli-progress'
+import fetch from 'node-fetch'
+import {promises as stream} from 'stream'
+import * as path from 'path'
+
+const VENDOR_INDEX_URL = 'https://developers.google.com/android/drivers'
+const VENDOR_URL_PREFIX = 'https://dl.google.com/dl/android/aosp/google_devices'
+
+async function getUrl(type: string, buildId: string, device: string) {
+  if (type == 'vendor') {
+    let resp = await fetch(VENDOR_INDEX_URL)
+    let index = await resp.text()
+
+    // TODO: parse HTML properly
+    let pattern = new RegExp(`"(${VENDOR_URL_PREFIX}-${device}-${buildId.toLowerCase()}-[a-z9-9.-_]+)">`)
+    let match = index.match(pattern)
+    if (match == null) {
+      throw new Error(`Image not found: ${type}, ${buildId}, ${device}`)
+    }
+
+    return match[1]
+  } else {
+    // TODO: implement factory and ota
+    throw new Error(`Unsupported type ${type}`)
+  }
+}
+
+async function downloadFile(type: string, buildId: string, device: string, outDir: string) {
+  let url = await getUrl(type, buildId, device)
+
+  let resp = await fetch(url)
+  let name = path.basename(url)
+  if (!resp.ok) {
+    throw new Error(`Error ${resp.status}: ${resp.statusText}`)
+  }
+
+  let bar = new cliProgress.SingleBar({
+    format: '    {bar} {percentage}% | {value}/{total} MB',
+  }, cliProgress.Presets.shades_classic)
+  let progress = 0
+  let totalSize = parseInt(resp.headers.get('content-length') ?? '0') / 1e6
+  bar.start(Math.round(totalSize), 0)
+  resp.body!.on('data', chunk => {
+    progress += chunk.length / 1e6
+    bar.update(Math.round(progress))
+  })
+
+  await stream.pipeline(resp.body!, createWriteStream(`${outDir}/${name}`))
+  bar.stop()
+}
+
+export default class Download extends Command {
+  static description = 'download device factory images, OTAs, and/or vendor packages'
+
+  static flags = {
+    help: flags.help({char: 'h'}),
+    type: flags.string({char: 't', options: ['factory', 'ota', 'vendor'], description: 'type(s) of images to download', default: 'factory', multiple: true}),
+    buildId: flags.string({char: 'b', description: 'build ID/number of the image(s) to download', required: true}),
+    device: flags.string({char: 'd', description: 'device(s) to download images for', required: true, multiple: true}),
+  }
+
+  static args = [
+    {name: 'out', description: 'directory to save downloaded files in', required: true},
+  ]
+
+  async run() {
+    let {flags, args: {out}} = this.parse(Download)
+
+    fs.mkdir(out, { recursive: true })
+
+    let buildId = flags.buildId.toUpperCase()
+
+    for (let type of flags.type) {
+      let prettyType = type == 'ota' ? 'OTA' : type.charAt(0).toUpperCase() + type.slice(1)
+      this.log(chalk.bold(chalk.blueBright(`${prettyType} - ${buildId}`)))
+
+      for (let device of flags.device) {
+        this.log(chalk.greenBright(`  ${device}`))
+        await downloadFile(type, buildId, device, out)
+      }
+    }
+  }
+}
diff --git a/yarn.lock b/yarn.lock
index 1a126dc..798ffeb 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2347,7 +2347,7 @@ nock@^13.0.0:
     lodash.set "^4.3.2"
     propagate "^2.0.0"
 
-node-fetch@^2.6.1:
+node-fetch@2, node-fetch@^2.6.1:
   version "2.6.6"
   resolved "https://registry.yarnpkg.com/node-fetch/-/node-fetch-2.6.6.tgz#1751a7c01834e8e1697758732e9efb6eeadfaf89"
   integrity sha512-Z8/6vRlTUChSdIgMa51jxQ4lrw/Jy5SOW10ObaA47/RElsAN2c5Pn8bTgFGWn/ibwzXTE8qwr1Yzx28vsecXEA==