improve device image downloading

- define a default download directory
- obtain image URL and checksum from the Git-tracked local build index instead of fetching it from
remote server
- default to downloading images for the current build ID
- download image to a temporary file first, mark it as complete only after checksum verification

src/commands/download.ts

@@ -1,75 +1,61 @@
 import { Command, flags } from '@oclif/command'
-import { promises as fs } from 'fs'
-import chalk from 'chalk'
-import { downloadFile, ImageType, IndexCache } from '../images/download'
 
-const IMAGE_TYPE_MAP: { [type: string]: ImageType } = {
-  factory: ImageType.Factory,
-  ota: ImageType.Ota,
-  vendor: ImageType.Vendor,
-}
+import { DEVICE_CONFIG_FLAGS, loadDeviceConfigs, resolveBuildId } from '../config/device'
+import { IMAGE_DOWNLOAD_DIR } from '../config/paths'
+import { ImageType, loadBuildIndex } from '../images/build-index'
+import { DeviceImage } from '../images/device-image'
+import { downloadDeviceImages } from '../images/download'
 
 export default class Download extends Command {
-  static description = 'download device factory images, OTAs, and/or vendor packages'
+  static description = 'download device factory images, OTAs, and/or vendor packages. Default output location is ' +
+    IMAGE_DOWNLOAD_DIR + '. To override it, use ADEVTOOL_IMG_DOWNLOAD_DIR environment variable.'
 
   static flags = {
-    help: flags.help({ char: 'h' }),
     type: flags.string({
       char: 't',
-      options: ['factory', 'ota', 'vendor'],
-      description: 'type(s) of images to download',
+      description: 'type(s) of images to download: factory | ota | vendor/$NAME (e.g. vendor/qcom, vendor/google_devices)',
       default: ['factory'],
       multiple: true,
     }),
     buildId: flags.string({
       char: 'b',
-      description: 'build ID(s) of the images to download',
-      required: true,
-      multiple: true,
-      default: ['latest'],
-    }),
-    device: flags.string({
-      char: 'd',
-      description: 'device(s) to download images for',
-      required: true,
+      description: 'build ID(s) of images to download, defaults to the current build ID',
+      required: false,
       multiple: true,
     }),
+    ...DEVICE_CONFIG_FLAGS,
   }
 
-  static args = [{ name: 'out', description: 'directory to save downloaded files in', required: true }]
-
   async run() {
-    let {
-      flags,
-      args: { out },
-    } = this.parse(Download)
+    let { flags } = this.parse(Download)
 
-    await fs.mkdir(out, { recursive: true })
+    let index = loadBuildIndex()
+    let deviceConfigs = loadDeviceConfigs(flags.devices)
 
-    this.log(chalk.bold(chalk.yellowBright("By downloading images, you agree to Google's terms and conditions:")))
-    this.log(
-      chalk.yellow(`    - https://developers.google.com/android/images#legal
-    - https://developers.google.com/android/ota#legal
-    - https://policies.google.com/terms
-`),
-    )
+    let images: DeviceImage[] = []
 
-    let cache: IndexCache = {}
-    for (let device of flags.device) {
-      this.log(chalk.greenBright(`\n${device}`))
+    let types = flags.type.map(s => s as ImageType)
 
-      for (let type of flags.type) {
-        let typeEnum = IMAGE_TYPE_MAP[type]
-        if (typeEnum == undefined) {
-          throw new Error(`Unknown type ${type}`)
-        }
-        let prettyType = type == 'ota' ? 'OTA' : type.charAt(0).toUpperCase() + type.slice(1)
+    for (let config of await deviceConfigs) {
+      for (let type of types) {
+        let buildIds = flags.buildId ?? [config.device.build_id]
 
-        for (let buildId of flags.buildId) {
-          this.log(chalk.bold(chalk.blueBright(`  ${prettyType} - ${buildId.toUpperCase()}`)))
-          await downloadFile(typeEnum, buildId, device, out, cache)
+        for (let buildIdStr of buildIds) {
+          let buildId = resolveBuildId(buildIdStr, config)
+          let image = DeviceImage.get(await index, config, buildId, type)
+          images.push(image)
         }
       }
     }
+
+    let missingImages = await DeviceImage.getMissing(images)
+
+    if (missingImages.length > 0) {
+      await downloadDeviceImages(missingImages)
+    }
+
+    for (let image of images) {
+      this.log(`${image.toString()}: '${image.getPath()}'`)
+    }
   }
 }

src/config/paths.ts

@@ -17,3 +17,5 @@ export const DEVICE_CONFIG_DIR = path.join(CONFIG_DIR, 'device')
 export const BUILD_INDEX_DIR = path.join(CONFIG_DIR, 'build-index')
 export const BUILD_INDEX_FILE = path.join(BUILD_INDEX_DIR, 'build-index.yml')
 export const MAIN_BUILD_INDEX_PART = path.join(BUILD_INDEX_DIR, 'build-index-main.yml')
+
+export const IMAGE_DOWNLOAD_DIR = process.env['ADEVTOOL_IMG_DOWNLOAD_DIR'] ?? path.join(ADEVTOOL_DIR, 'dl')

src/images/device-image.ts

@@ -2,6 +2,7 @@ import assert from 'assert'
 import { promises as fs } from 'fs'
 import path from 'path'
 import { DeviceConfig, getDeviceBuildId } from '../config/device'
+import { IMAGE_DOWNLOAD_DIR } from '../config/paths'
 import { BuildIndex, DEFAULT_BASE_DOWNLOAD_URL, ImageType } from './build-index'
 
 export class DeviceImage {
@@ -15,6 +16,18 @@ export class DeviceImage {
   ) {
   }
 
+  getPath() {
+    return path.join(IMAGE_DOWNLOAD_DIR, this.fileName)
+  }
+
+  async isPresent() {
+    try {
+      return (await fs.stat(this.getPath())).isFile()
+    } catch {
+      return false
+    }
+  }
+
   static get(index: BuildIndex, deviceConfig: DeviceConfig, buildId: string, type: ImageType) {
     let deviceBuildId = getDeviceBuildId(deviceConfig, buildId)
     let buildProps = index.get(deviceBuildId)
@@ -44,6 +57,17 @@ export class DeviceImage {
     return new DeviceImage(deviceConfig, type, buildId, fileName, sha256, url)
   }
 
+  static async getMissing(arr: DeviceImage[]) {
+    let res: DeviceImage[] = []
+    let imagePresence = arr.map(i => ({image: i, isPresent: i.isPresent()}))
+    for (let e of imagePresence) {
+      if (!(await e.isPresent)) {
+        res.push(e.image)
+      }
+    }
+    return res
+  }
+
   toString() {
     return `'${this.deviceConfig.device.name} ${this.buildId} ${this.type}'`
   }

src/images/download.ts

@@ -1,86 +1,48 @@
-import { createWriteStream } from 'fs'
+import assert from 'assert'
+import chalk from 'chalk'
 import cliProgress from 'cli-progress'
+import { createHash } from 'crypto'
+import { createWriteStream, existsSync, promises as fs } from 'fs'
 import fetch from 'node-fetch'
-import { promises as stream } from 'stream'
 import path from 'path'
-import _ from 'lodash'
+import { promises as stream } from 'stream'
+import { IMAGE_DOWNLOAD_DIR } from '../config/paths'
 
-const DEV_INDEX_URL = 'https://developers.google.com/android'
-const DEV_COOKIE = 'devsite_wall_acks=nexus-image-tos,nexus-ota-tos'
-const DL_URL_PREFIX = 'https://dl.google.com/dl/android/aosp/'
+import { maybePlural } from '../util/cli'
+import { ImageType } from './build-index'
+import { DeviceImage } from './device-image'
 
-interface ImageTypeInfo {
-  indexPath: string
-  filePattern: string
-}
-
-export enum ImageType {
-  Ota = 'ota',
-  Factory = 'factory',
-  Vendor = 'vendor',
-}
-
-export type IndexCache = { [type in ImageType]?: string }
-
-const IMAGE_TYPES: Record<ImageType, ImageTypeInfo> = {
-  [ImageType.Factory]: {
-    indexPath: 'images',
-    filePattern: 'DEVICE-BUILDID',
-  },
-  [ImageType.Ota]: {
-    indexPath: 'ota',
-    filePattern: 'DEVICE-ota-BUILDID',
-  },
-  [ImageType.Vendor]: {
-    indexPath: 'drivers',
-    filePattern: 'google_devices-DEVICE-BUILDID',
-  },
-}
-
-async function getUrl(type: ImageType, buildId: string, device: string, cache: IndexCache) {
-  let { indexPath, filePattern } = IMAGE_TYPES[type]
-
-  let index = cache[type]
-  if (index == undefined) {
-    let resp = await fetch(`${DEV_INDEX_URL}/${indexPath}`, {
-      headers: {
-        Cookie: DEV_COOKIE,
-      },
-    })
-
-    index = await resp.text()
-    cache[type] = index
+export async function downloadDeviceImages(images: DeviceImage[], showTncNotice = true) {
+  await fs.mkdir(IMAGE_DOWNLOAD_DIR, { recursive: true })
+  if (showTncNotice) {
+    logTermsAndConditionsNotice(images)
   }
 
-  let filePrefix = filePattern
-    .replace('DEVICE', device)
-    .replace('BUILDID', buildId == 'latest' ? '' : `${buildId.toLowerCase()}-`)
-  let urlPrefix = DL_URL_PREFIX + filePrefix
-
-  let pattern = new RegExp(`"(${_.escapeRegExp(urlPrefix)}.+?)"`, 'g')
-  let matches = Array.from(index.matchAll(pattern))
-  if (matches.length == 0) {
-    throw new Error(`Image not found: ${type}, ${buildId}, ${device}`)
+  for (let image of images) {
+    console.log(chalk.bold(chalk.blueBright(`\n${image.deviceConfig.device.name} ${image.buildId} ${image.type}`)))
+    await downloadImage(image, IMAGE_DOWNLOAD_DIR)
   }
-
-  if (buildId == 'latest') {
-    return matches[matches.length - 1][1]
-  }
-  return matches[0][1]
 }
 
-export async function downloadFile(
-  type: ImageType,
-  buildId: string,
-  device: string,
-  outDir: string,
-  cache: IndexCache = {},
-) {
-  let url = await getUrl(type, buildId, device, cache)
+export async function downloadMissingDeviceImages(images: DeviceImage[]) {
+  let missingImages = await DeviceImage.getMissing(images)
 
-  console.log(`    ${url}`)
-  let resp = await fetch(url)
-  let name = path.basename(url)
+  if (missingImages.length > 0) {
+    console.log(`Missing image${maybePlural(missingImages)}: ${DeviceImage.arrayToString(missingImages)}`)
+    await downloadDeviceImages(missingImages)
+  }
+}
+
+async function downloadImage(image: DeviceImage, outDir: string) {
+  let tmpOutFile = path.join(outDir, image.fileName + '.tmp')
+  await fs.rm(tmpOutFile, { force: true })
+
+  let completeOutFile = path.join(outDir, image.fileName)
+  assert(!existsSync(completeOutFile), completeOutFile + ' already exists')
+
+  console.log(`    ${image.url}`)
+
+  let resp = await fetch(image.url)
   if (!resp.ok) {
     throw new Error(`Error ${resp.status}: ${resp.statusText}`)
   }
@@ -94,11 +56,38 @@ export async function downloadFile(
   let progress = 0
   let totalSize = parseInt(resp.headers.get('content-length') ?? '0') / 1e6
   bar.start(Math.round(totalSize), 0)
+
+  let sha256 = createHash('sha256')
+
   resp.body!.on('data', chunk => {
+    sha256.update(chunk)
     progress += chunk.length / 1e6
     bar.update(Math.round(progress))
   })
 
-  await stream.pipeline(resp.body!, createWriteStream(`${outDir}/${name}`))
+  await stream.pipeline(resp.body!, createWriteStream(tmpOutFile))
   bar.stop()
+
+  let sha256Digest: string = sha256.digest('hex')
+  console.log('SHA-256: ' + sha256Digest)
+  assert(sha256Digest === image.sha256, 'SHA256 mismatch, expected ' + image.sha256)
+
+  await fs.rename(tmpOutFile, completeOutFile)
+}
+
+function logTermsAndConditionsNotice(images: DeviceImage[]) {
+  if (images.filter(i => i.type === ImageType.Factory || i.type === ImageType.Ota).length == 0) {
+    // vendor images show T&C notice themselves as part of unpacking
+    return
+  }
+
+  console.log(chalk.bold('\nBy downloading images, you agree to Google\'s terms and conditions:'))
+
+  let msg = '    - Factory images: https://developers.google.com/android/images#legal\n'
+  if (images.find(i => i.type === ImageType.Ota) !== undefined) {
+    msg += '    - OTA images: https://developers.google.com/android/ota#legal\n'
+  }
+  msg += '    - Beta factory/OTA images: https://developer.android.com/studio/terms\n'
+
+  console.log(msg)
 }

src/util/cli.ts

@@ -29,6 +29,12 @@ export async function withSpinner<Return>(action: string, callback: (spinner: or
   return ret
 }
 
+export function maybePlural<T>(arr: ArrayLike<T>, singleEnding = '', multiEnding = 's') {
+  let len = arr.length
+  assert(len > 0)
+  return len > 1 ? multiEnding : singleEnding
+}
+
 export function showGitDiff(repoPath: string, filePath?: string) {
   let args = ['-C', repoPath, `diff`]
   if (filePath !== undefined) {