src/memento_updater.cc

// Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <iostream>
#include <algorithm>
#include <string>

#include <stdlib.h>

#include "memento_common.h"
#include "command_handler.h"
#include "general_installer.h"
#include "diskimage_installer.h"

using namespace chromeos_memento_updater;
using std::cerr;
using std::endl;
using std::pair;
using std::string;
using std::vector;

int parse_partition_option(const string& option,
                           GeneralInstaller* installer_ptr) {
  int partition_option = -1;
  if (!installer_ptr->AllowPartitionOptions() && !option.empty()) {
    error_exit("Partition options are not allowed");
  }
  if (installer_ptr->AllowPartitionOptions()) {
    if (!option.empty()) {
      partition_option = strtol(option.c_str(), NULL, 10);
      // strtol return 0 on failure
      if (partition_option == 0) {
        error_exit("Invalid partition option");
      }
    }
  }
  return partition_option;
}

void install_device(const string& option, PartitionNum partition,
                    CommandHandler& handler,
                    GeneralInstaller* installer_ptr) {
  int source_partition = option.empty() ?
                         static_cast<int>(partition) :
                         parse_partition_option(option, installer_ptr);
  installer_ptr->InstallDevice(source_partition,
                               handler.GetOutputDevicePartition(partition),
                               string());
}

void install_main_device(const string& option, PartitionNum rootfs,
                         PartitionNum kernel, CommandHandler& handler,
                         GeneralInstaller* installer_ptr) {
  int kernel_source = static_cast<int>(kernel);
  int rootfs_source = static_cast<int>(rootfs);
  // Parse option, format is <kernel>:<rootfs>
  if (!option.empty()) {
    size_t colon;
    if ((colon = option.find(':')) == string::npos) {
      error_exit("Invalid partition option format, expect <kernel>:<rootfs>");
    }
    kernel_source = parse_partition_option(option.substr(0, colon),
                                           installer_ptr);
    rootfs_source = parse_partition_option(option.substr(colon + 1),
                                           installer_ptr);
  }
  // For shopfloor/miniomaha, kernel are combined with rootfs
  if (installer_ptr->CombinedKernel()) {
    installer_ptr->InstallDevice(rootfs_source,
                                 handler.GetOutputDevicePartition(rootfs),
                                 handler.GetOutputDevicePartition(kernel));
  } else {
    installer_ptr->InstallDevice(rootfs_source,
                                 handler.GetOutputDevicePartition(rootfs),
                                 string());
    installer_ptr->InstallDevice(kernel_source,
                                 handler.GetOutputDevicePartition(kernel),
                                 string());
  }
}

void update_firmware(string option,
                     CommandHandler& handler,
                     GeneralInstaller* installer_ptr) {
  bool release_mode = false;
  string installer_path;
  size_t colon;
  // Option format <"release"/"factory">[:<updater_path>] or <updater_path>
  if ((colon = option.find(":")) != string::npos) {
    installer_path = option.substr(colon + 1);
    option = option.substr(0, colon);
  }
  if (option == "release") {
    release_mode = true;
  } else if (option == "factory") {
    // Do nothing
  } else {
    installer_path = option;
    if (handler.ReleaseOnly()) {
      release_mode = true;
    }
  }

  // If there is only release image,
  // then default to release mode, otherwise use factory mode.
  if (release_mode) {
    if (installer_path.empty()) {
      installer_path = installer_ptr->PrepareFirmware(
          handler.GetOutputDevicePartition(kReleasePartition));
    }
    installer_ptr->UpdateFirmware(installer_path, "release");
  } else {
    if (installer_path.empty()) {
      installer_path = installer_ptr->PrepareFirmware(
          handler.GetOutputDevicePartition(kFactoryPartition));
    }
    installer_ptr->UpdateFirmware(installer_path, "factory_install");
  }
}

void activate_image(string& option, CommandHandler& handler,
                    GeneralInstaller* installer_ptr) {
  if (option == "release" || handler.ReleaseOnly()) {
    installer_ptr->ActivateImage(
        handler.GetOutputDevicePartition(kReleasePartition));
  } else {
    installer_ptr->ActivateImage(
        handler.GetOutputDevicePartition(kFactoryPartition));
  }
}

int main(int argc, char** argv)
{
  // memento_updater --to <dest-device>
  //                 --from <type>:<location>
  //                 --contents <content:option, ...>
  //                 --board <board>
  CommandHandler handler(argc, argv);
  GeneralInstaller* installer_ptr = handler.GetInstallerPtr();
  vector< pair<CommandHandler::InstallOption, string> > option_list;
  option_list = handler.GetOptionList();
  // We want options to be executed in order
  std::sort(option_list.begin(), option_list.end());
  for (size_t i = 0; i < option_list.size(); i++) {
    switch (option_list[i].first) {
      case CommandHandler::kResetDevice:
        // TODO(chunyen): fill in this after the utility function for
        // reset device is written
        cerr << "Reset device is not supported yet." << endl;
        break;
      case CommandHandler::kResetPartition:
        installer_ptr->ResetPartitionTable(
            handler.GetOutputDevice());
        break;
      case CommandHandler::kFactory:
          install_main_device(option_list[i].second, kFactoryPartition,
                              kFactoryKernelPartition, handler, installer_ptr);
        break;
      case CommandHandler::kRelease:
          install_main_device(option_list[i].second, kReleasePartition,
                              kReleaseKernelPartition, handler, installer_ptr);
        break;
      case CommandHandler::kStateful:
        if (option_list[i].second == "empty") {
          installer_ptr->InitStatefulPartition(
              handler.GetOutputDevicePartition(kStatefulPartition));
        } else {
          install_device(option_list[i].second, kStatefulPartition,
                         handler, installer_ptr);
        }
        break;
      case CommandHandler::kOem:
        install_device(option_list[i].second, kOemPartition,
                       handler, installer_ptr);
        break;
      case CommandHandler::kEfi:
        install_device(option_list[i].second, kEfiPartition,
                       handler, installer_ptr);
        break;
      case CommandHandler::kFirmware:
        update_firmware(option_list[i].second, handler, installer_ptr);
        break;
      case CommandHandler::kActivate:
        activate_image(option_list[i].second, handler, installer_ptr);
        break;
      default:
        error_exit("Should never be here");
    }
  }
  return 0;
}