src/command_handler.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 "command_handler.h"

#include <ctype.h>  // isdigit()
#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <iostream>
#include <string>
#include <sstream>
#include <map>

#include "general_installer.h"
#include "diskimage_installer.h"

using std::cerr;
using std::endl;
using std::string;
using std::stringstream;

namespace chromeos_memento_updater {

const string CommandHandler::kOptionList[] = {
  "reset_device",
  "partition",
  "factory",
  "release",
  "stateful",
  "oem",
  "efi",
  "firmware",
  "activate"
};
const int CommandHandler::kOptions = 9;

CommandHandler::CommandHandler(int argc, char** argv) {
  installer_ptr_ = NULL;
  have_release_ = false;
  have_factory_ = false;
  SetRootDev();
  ParseCommandLine(argc, argv);
}

CommandHandler::~CommandHandler() {
  if (installer_ptr_ != NULL) {
    delete installer_ptr_;
  }
}

void CommandHandler::Usage() {
  cerr << "memento_update --to <dest-device>" << endl <<
          "               --from <type>:<location>" << endl <<
          "               --contents <content:option, ...>" << endl <<
          "               --board <board>" << endl;
  cerr << endl;
  cerr << "Please check README for more information about Usage." << endl;
  exit(0);
}

void CommandHandler::SetRootDev() {
  FILE* fp = popen("rootdev -s", "r");
  if (fp == NULL) {
    cerr << "ERROR: Cannot run rootdev." << endl;
    exit(1);
  }
  int kDevLength = 100;
  char buffer[kDevLength];
  buffer[0] = '\0';
  fgets(buffer, kDevLength, fp);
  if (buffer[0] == '\0') {
    // Empty string
    cerr << "ERROR: Run rootdev failed" << endl;
    exit(1);
  }
  root_device_.assign(buffer);
}

string CommandHandler::GetOutputDevicePartition(PartitionNum partition) {
  stringstream device;
  device << output_device_;
  // Ends with digit, add 'p'
  if (isdigit(output_device_[output_device_.size() - 1])) {
    device << 'p';
  }
  // Check if boot and install on the same device
  if (root_device_.find(output_device_) != string::npos) {
    // It is not possible to boot and install both image on the same device.
    if (have_release_ && have_factory_ ) {
      cerr << "ERROR: Not able to install both release and factory when boot"
              "on output device." << endl;
      exit(1);
    }
    switch (partition) {
      case kFactoryPartition:
      case kReleasePartition:
        if (root_device_.find("3") != string::npos) {
          device << '5';
        } else {
          device << '3';
        }
        break;
      case kFactoryKernelPartition:
      case kReleaseKernelPartition:
        if (root_device_.find("3") != string::npos) {
          device << '4';
        } else {
          device << '2';
        }
        break;
      default:
        device << partition;
    }
  } else {
    device << partition;
  }
  return device.str();
}

void CommandHandler::ParseInputDevice(const string& input_string,
                                      const string& running_dir) {
// TODO(chunyen): find a solution to download from http://chromeos-image
// TODO(chunyen): Fill in the code blocks after the corresponding handler
// are written.
  if (input_string.find("image:") == 0) {
    input_device_ = input_string.substr(strlen("image:"));
    DiskImageInstallFormat format = kImage;
    const string kZipSuffix = ".zip";
    // Check if path ends with .zip
    if (input_device_.find(kZipSuffix) ==
        input_device_.length() - kZipSuffix.length()) {
      format = kZip;
    }
    installer_ptr_ = new DiskImageInstaller(running_dir, board_,
                                            input_device_, format);
// TODO(chunyen):
//  } else if (input_string.find("shopfloor:") == 0) {
//  } else if (input_string.find("miniomaha:") == 0) {
//  } else if (input_string.find("recovery:") == 0){
  } else {
    cerr << "Unrecognized source media." << endl;
    Usage();
  }
}

void CommandHandler::ParseContent(const string& content_string) {
  std::istringstream content(content_string);
  string token;
  // Create a map that map option string to enum
  std::map<string, InstallOption> string_to_enum;
  for (int i = 0; i < kOptions; i++) {
    string_to_enum[kOptionList[i]] = static_cast<InstallOption>(i);
  }
  while (std::getline(content, token, ',')) {
    InstallOption option_id;
    string option_string;
    string key;
    size_t colon;
    if ((colon = token.find(":")) != string::npos) {
      key = token.substr(0, colon);
      option_string.assign(token.substr(colon + 1));
    } else {
      key = token;
    }
    std::map<string, InstallOption>::iterator it = string_to_enum.find(key);
    if (it == string_to_enum.end()) {
      cerr << "Cannot find matching option: " << token << endl;
      Usage();
    } else {
      option_id = it->second;
    }
    if (option_id == kRelease) {
      have_release_ = true;
    }
    if (option_id == kFactory) {
      have_factory_ = true;
    }
    option_list_.push_back(make_pair(option_id, option_string));
  }
}

void CommandHandler::ParseCommandLine(int argc, char** argv) {

  string input_string;
  string content_string;
  int option_index = 0;
  char c;
  static struct option long_options[] = {
    {"to", required_argument, 0, 't'},
    {"from", required_argument, 0, 'f'},
    {"contents", required_argument, 0, 'c'},
    {"board", required_argument, 0, 'b'}
  };
  while ((c = getopt_long(argc, argv, "t:f:c:b:",
                          long_options, &option_index)) != -1) {
    switch(c) {
     case 't':
      output_device_.assign(optarg);
      break;
     case 'f':
      input_string.assign(optarg);
      break;
     case 'c':
      content_string.assign(optarg);
      break;
     case 'b':
      board_.assign(optarg);
      break;
     case '?':
      Usage();
      break;
     default:
      cerr << "Unexpected return value " << c << " from getopt." << endl;
      exit(1);
      break;
    }
  }
  if (output_device_.empty() || input_string.empty() ||
      content_string.empty()) {
    Usage();
  }
  // TODO(chunyen): Auto detect board?
  if (board_.empty()) {
    cerr << "Auto detect board is not supported yet.";
    Usage();
  }
  // Find the running directory from argv[0]
  string program_name(argv[0]);
  string home_dir(".");

  size_t pos;
  if ((pos = program_name.find_last_of('/')) != string::npos) {
    home_dir.assign(program_name.substr(0, program_name.length() - pos));
  }
  cerr << home_dir << endl;
  ParseInputDevice(input_string, home_dir);
  ParseContent(content_string);
}

}  // namespace chromeos_memento_updater