/* Copyright 2015 The ChromiumOS Authors
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 *
 * This utility wraps around "cgpt" execution to work with NAND. If the target
 * device is an MTD device, this utility will read the GPT structures from
 * FMAP, invokes "cgpt" on that, and writes the result back to NOR flash. */

#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <limits.h>
#if !defined(__FreeBSD__)
#include <linux/major.h>
#endif
#include <stdbool.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/stat.h>
#if !defined(__FreeBSD__)
#include <sys/sysmacros.h>
#endif
#include <sys/types.h>
#include <unistd.h>

#include "2common.h"
#include "2sha.h"
#include "2sysincludes.h"
#include "cgpt.h"
#include "cgpt_nor.h"
#include "file_keys.h"

// Check if cmdline |argv| has "-D". "-D" signifies that GPT structs are stored
// off device, and hence we should not wrap around cgpt.
static bool has_dash_D(int argc, const char *const argv[]) {
  int i;
  // We go from 2, because the second arg is a cgpt command such as "create".
  for (i = 2; i < argc; ++i) {
    if (strcmp("-D", argv[i]) == 0) {
      return true;
    }
  }
  return false;
}

// Check if |device_path| is an MTD device based on its major number being 90.
static bool is_mtd(const char *device_path) {
  struct stat stat;
  if (lstat(device_path, &stat) != 0) {
    return false;
  }

#if !defined(__FreeBSD__)
  if (major(stat.st_rdev) == MTD_CHAR_MAJOR) {
    return true;
  }
#endif
  return false;
}

// Return the element in |argv| that is an MTD device.
static const char *find_mtd_device(int argc, const char *const argv[]) {
  int i;
  for (i = 2; i < argc; ++i) {
    if (is_mtd(argv[i])) {
      return argv[i];
    }
  }
  return NULL;
}

static int wrap_cgpt(int argc,
                     const char *const argv[],
                     const char *mtd_device) {
  uint8_t original_hash[VB2_SHA1_DIGEST_SIZE];
  uint8_t modified_hash[VB2_SHA1_DIGEST_SIZE];
  int ret = 0;

  // Create a temp dir to work in.
  ret++;
  char temp_dir[] = "/tmp/cgpt_wrapper.XXXXXX";
  if (mkdtemp(temp_dir_template) == NULL) {
    Error("Cannot create a temporary directory.\n");
    return ret;
  }
  if (ReadNorFlash(temp_dir) != 0) {
    goto cleanup;
  }
  char rw_gpt_path[PATH_MAX];
  if (snprintf(rw_gpt_path, sizeof(rw_gpt_path), "%s/rw_gpt", temp_dir) < 0) {
    goto cleanup;
  }
  if (VB2_SUCCESS != DigestFile(rw_gpt_path, VB2_HASH_SHA1,
				original_hash, sizeof(original_hash))) {
    Error("Cannot compute original GPT digest.\n");
    goto cleanup;
  }

  // Obtain the MTD size.
  ret++;
  uint64_t drive_size = 0;
  if (GetMtdSize(mtd_device, &drive_size) != 0) {
    Error("Cannot get the size of %s.\n", mtd_device);
    goto cleanup;
  }

  // Launch cgpt on "rw_gpt" with -D size.
  ret++;
  const char** my_argv = calloc(argc + 2 + 1, sizeof(char *));
  if (my_argv == NULL) {
    errno = ENOMEM;
    goto cleanup;
  }
  memcpy(my_argv, argv, sizeof(char *) * argc);
  char *real_cgpt;
  if (asprintf(&real_cgpt, "%s.bin", argv[0]) == -1) {
    free(my_argv);
    goto cleanup;
  }
  my_argv[0] = real_cgpt;

  int i;
  for (i = 2; i < argc; ++i) {
    if (strcmp(my_argv[i], mtd_device) == 0) {
      my_argv[i] = rw_gpt_path;
    }
  }
  my_argv[argc] = "-D";
  char size[32];
  snprintf(size, sizeof(size), "%" PRIu64, drive_size);
  my_argv[argc + 1] = size;
  i = ForkExecV(NULL, my_argv);
  free(real_cgpt);
  free(my_argv);
  if (i != 0) {
    Error("Cannot exec cgpt to modify rw_gpt.\n");
    goto cleanup;
  }

  // Write back "rw_gpt" to NOR flash in two chunks.
  ret++;
  if (VB2_SUCCESS == DigestFile(rw_gpt_path, VB2_HASH_SHA1,
				modified_hash, sizeof(modified_hash))) {
    if (memcmp(original_hash, modified_hash, VB2_SHA1_DIGEST_SIZE) != 0) {
      ret = WriteNorFlash(temp_dir);
    } else {
      ret = 0;
    }
  }

cleanup:
  RemoveDir(temp_dir);
  return ret;
}

int main(int argc, const char *argv[]) {
  char resolved_cgpt[PATH_MAX];
  pid_t pid = getpid();
  char exe_link[40];
  int retval = 0;

  if (argc < 1) {
    return -1;
  }

  const char *orig_argv0 = argv[0];

  snprintf(exe_link, sizeof(exe_link), "/proc/%d/exe", pid);
  memset(resolved_cgpt, 0, sizeof(resolved_cgpt));
  if (readlink(exe_link, resolved_cgpt, sizeof(resolved_cgpt) - 1) == -1) {
    perror("readlink");
    return -1;
  }

  argv[0] = resolved_cgpt;

  if (argc > 2 && !has_dash_D(argc, argv)) {
    const char *mtd_device = find_mtd_device(argc, argv);
    if (mtd_device) {
      retval = wrap_cgpt(argc, argv, mtd_device);
      goto cleanup;
    }
  }

  // Forward to cgpt as-is. Real cgpt has been renamed cgpt.bin.
  char *real_cgpt;
  if (asprintf(&real_cgpt, "%s.bin", argv[0]) == -1) {
    retval = -1;
    goto cleanup;
  }
  argv[0] = real_cgpt;
  if (execv(argv[0], (char * const *)argv) == -1) {
    err(-2, "execv(%s) failed", real_cgpt);
  }
  free(real_cgpt);
  retval = -2;

cleanup:
  argv[0] = orig_argv0;
  return retval;
}