/*
 * Copyright (C) 2016 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <iostream>

#include <endian.h>
#include <inttypes.h>
#include <string.h>

#include <base/files/file_util.h>
#include <base/strings/string_util.h>
#include <base/strings/stringprintf.h>

#include "bvb_unittest_util.h"
#include "bvb_refimpl.h"

class BvbToolTest : public BaseBvbToolTest {
public:
  BvbToolTest() {}
};

// This test ensure that the version is increased in both
// bvb_boot_image.h and the bvb tool.
TEST_F(BvbToolTest, BvbVersionInSync)
{
  base::FilePath path = testdir_.Append("version.txt");
  EXPECT_COMMAND(0,
                 "./bvbtool version > %s",
                 path.value().c_str());
  std::string printed_version;
  ASSERT_TRUE(base::ReadFileToString(path, &printed_version));
  base::TrimWhitespaceASCII(printed_version, base::TRIM_ALL, &printed_version);
  std::string expected_version = base::StringPrintf("%d.%d",
                                                    BVB_MAJOR_VERSION,
                                                    BVB_MINOR_VERSION);
  EXPECT_EQ(printed_version, expected_version);
}

TEST_F(BvbToolTest, ExtractPublicKey)
{
  GenerateBootImage("SHA256_RSA2048", "", 0,
                    base::FilePath("test/testkey_rsa2048.pem"));

  base::FilePath public_key_path = testdir_.Append("public_key.bin");
  EXPECT_COMMAND(0,
                 "./bvbtool extract_public_key --key test/testkey_rsa2048.pem"
                 " --output %s",
                 public_key_path.value().c_str());

  std::string key_data;
  ASSERT_TRUE(base::ReadFileToString(public_key_path, &key_data));

  BvbBootImageHeader h;
  bvb_boot_image_header_to_host_byte_order(
      reinterpret_cast<BvbBootImageHeader*>(boot_image_.data()), &h);
  uint8_t *d = reinterpret_cast<uint8_t*>(boot_image_.data());
  size_t auxilary_data_block_offset =
      sizeof(BvbBootImageHeader) + h.authentication_data_block_size;
  EXPECT_GT(h.auxilary_data_block_size, key_data.size());
  EXPECT_EQ(0, memcmp(key_data.data(),
                      d + auxilary_data_block_offset + h.public_key_offset,
                      key_data.size()));
}

TEST_F(BvbToolTest, PayloadsAreCorrect)
{
  GenerateBootImage("SHA256_RSA2048", "", 0,
                    base::FilePath("test/testkey_rsa2048.pem"));

  BvbBootImageHeader h;
  bvb_boot_image_header_to_host_byte_order(
      reinterpret_cast<BvbBootImageHeader*>(boot_image_.data()), &h);

  uint8_t *d = reinterpret_cast<uint8_t*>(boot_image_.data()) +
      sizeof(BvbBootImageHeader) +
      h.authentication_data_block_size +
      h.auxilary_data_block_size;

  // Check that the kernel and initrd images are inserted correctly.
  for (int n = 0; n < 2; n++) {
    std::string paths[2] = {"test/dummy_kernel.bin",
                            "test/dummy_initrd.bin"};
    base::FilePath path(paths[n]);

    int64_t file_size;
    std::vector<uint8_t> file_data;
    ASSERT_TRUE(base::GetFileSize(path, &file_size));
    file_data.resize(file_size);
    ASSERT_TRUE(base::ReadFile(path,
                               reinterpret_cast<char*>(file_data.data()),
                               file_data.size()));

    switch (n) {
      case 0:
        // kernel
        EXPECT_EQ(65536U, file_data.size());
        EXPECT_EQ(file_data.size(), h.kernel_size);
        EXPECT_EQ(0, memcmp(file_data.data(),
                            d + h.kernel_offset,
                            file_data.size()));
        break;
      case 1:
        // initrd
        EXPECT_EQ(131072U, file_data.size());
        EXPECT_EQ(file_data.size(), h.initrd_size);
        EXPECT_EQ(0, memcmp(file_data.data(),
                            d + h.initrd_offset,
                            file_data.size()));
        break;
      default:
        ASSERT_TRUE(false);
        break;
    }
  }
}

TEST_F(BvbToolTest, CheckCmdline)
{
  std::string cmdline("init=/sbin/init ro x y z");
  GenerateBootImage("SHA256_RSA2048", cmdline, 0,
                    base::FilePath("test/testkey_rsa2048.pem"));

  BvbBootImageHeader h;
  bvb_boot_image_header_to_host_byte_order(
      reinterpret_cast<BvbBootImageHeader*>(boot_image_.data()), &h);

  EXPECT_EQ(0, ::strcmp(cmdline.c_str(),
                        reinterpret_cast<const char*>(h.kernel_cmdline)));
}

TEST_F(BvbToolTest, CheckAddresses)
{
  GenerateBootImage("SHA256_RSA2048", "", 0,
                    base::FilePath("test/testkey_rsa2048.pem"),
                    "--kernel_addr 0x42 --initrd_addr 43");

  BvbBootImageHeader h;
  bvb_boot_image_header_to_host_byte_order(
      reinterpret_cast<BvbBootImageHeader*>(boot_image_.data()), &h);

  EXPECT_EQ(0x42U, h.kernel_addr);
  EXPECT_EQ(43U, h.initrd_addr);
}

TEST_F(BvbToolTest, CheckProperties)
{
  GenerateBootImage("SHA256_RSA2048", "", 0,
                    base::FilePath("test/testkey_rsa2048.pem"),
                    "--prop foo:brillo "
                    "--prop bar:chromeos "
                    "--prop prisoner:24601 "
                    "--prop hexnumber:0xcafe "
                    "--prop hexnumber_capital:0xCAFE "
                    "--prop large_hexnumber:0xfedcba9876543210 "
                    "--prop larger_than_uint64:0xfedcba98765432101 "
                    "--prop almost_a_number:423x "
                    "--prop_from_file blob:test/small_blob.bin "
                    );

  BvbBootImageHeader h;
  bvb_boot_image_header_to_host_byte_order(
      reinterpret_cast<BvbBootImageHeader*>(boot_image_.data()), &h);

  EXPECT_EQ(BVB_VERIFY_RESULT_OK,
            bvb_verify_boot_image(boot_image_.data(), boot_image_.size(),
                                  nullptr, nullptr));

  const char *s;
  size_t len;
  uint64_t val;

  // Basic.
  s = bvb_lookup_property(boot_image_.data(), boot_image_.size(),
                          "foo", 0, &len);
  EXPECT_EQ(0, strcmp(s, "brillo"));
  EXPECT_EQ(6U, len);
  s = bvb_lookup_property(boot_image_.data(), boot_image_.size(),
                          "bar", 0, &len);
  EXPECT_EQ(0, strcmp(s, "chromeos"));
  EXPECT_EQ(8U, len);
  s = bvb_lookup_property(boot_image_.data(), boot_image_.size(),
                          "non-existant", 0, &len);
  EXPECT_EQ(0U, len);
  EXPECT_EQ(NULL, s);

  // Numbers.
  EXPECT_NE(0, bvb_lookup_property_uint64(
      boot_image_.data(), boot_image_.size(), "prisoner", 0, &val));
  EXPECT_EQ(24601U, val);

  EXPECT_NE(0, bvb_lookup_property_uint64(
      boot_image_.data(), boot_image_.size(), "hexnumber", 0, &val));
  EXPECT_EQ(0xcafeU, val);

  EXPECT_NE(0, bvb_lookup_property_uint64(
      boot_image_.data(), boot_image_.size(), "hexnumber_capital", 0, &val));
  EXPECT_EQ(0xcafeU, val);

  EXPECT_NE(0, bvb_lookup_property_uint64(
      boot_image_.data(), boot_image_.size(), "large_hexnumber", 0, &val));
  EXPECT_EQ(0xfedcba9876543210U, val);

  // We could catch overflows and return an error ... but we currently don't.
  EXPECT_NE(0, bvb_lookup_property_uint64(
      boot_image_.data(), boot_image_.size(), "larger_than_uint64", 0, &val));
  EXPECT_EQ(0xedcba98765432101U, val);

  // Number-parsing failures.
  EXPECT_EQ(0, bvb_lookup_property_uint64(
      boot_image_.data(), boot_image_.size(), "foo", 0, &val));

  EXPECT_EQ(0, bvb_lookup_property_uint64(
      boot_image_.data(), boot_image_.size(), "almost_a_number", 0, &val));

  // Blobs.
  //
  // test/small_blob.bin is 21 byte file full of NUL-bytes except for
  // the string "brillo ftw!" at index 2 and '\n' at the last byte.
  s = bvb_lookup_property(boot_image_.data(), boot_image_.size(),
                          "blob", 0, &len);
  EXPECT_EQ(21U, len);
  EXPECT_EQ(0, memcmp(s,     "\0\0", 2));
  EXPECT_EQ(0, memcmp(s +  2, "brillo ftw!", 11));
  EXPECT_EQ(0, memcmp(s + 13, "\0\0\0\0\0\0\0", 7));
  EXPECT_EQ('\n', s[20]);
}

TEST_F(BvbToolTest, CheckRollbackIndex)
{
  uint64_t rollback_index = 42;
  GenerateBootImage("SHA256_RSA2048", "", rollback_index,
                    base::FilePath("test/testkey_rsa2048.pem"));

  BvbBootImageHeader h;
  bvb_boot_image_header_to_host_byte_order(
      reinterpret_cast<BvbBootImageHeader*>(boot_image_.data()), &h);

  EXPECT_EQ(rollback_index, h.rollback_index);
}

TEST_F(BvbToolTest, CheckPubkeyReturned)
{
  GenerateBootImage("SHA256_RSA2048", "", 0,
                    base::FilePath("test/testkey_rsa2048.pem"));

  const uint8_t* pubkey = NULL;
  size_t pubkey_length = 0;

  EXPECT_EQ(BVB_VERIFY_RESULT_OK,
            bvb_verify_boot_image(boot_image_.data(), boot_image_.size(),
                                  &pubkey, &pubkey_length));

  BvbBootImageHeader h;
  bvb_boot_image_header_to_host_byte_order(
      reinterpret_cast<BvbBootImageHeader*>(boot_image_.data()), &h);

  EXPECT_EQ(pubkey_length, h.public_key_size);

  const uint8_t* expected_pubkey = boot_image_.data() +
      sizeof(BvbBootImageHeader) +
      h.authentication_data_block_size +
      h.public_key_offset;
  EXPECT_EQ(pubkey, expected_pubkey);
}

TEST_F(BvbToolTest, Info)
{
  GenerateBootImage("SHA256_RSA2048", "foobar=cmdline test=42", 0,
                    base::FilePath("test/testkey_rsa2048.pem"),
                    "--prop foo:brillo "
                    "--prop bar:chromeos "
                    "--prop prisoner:24601 "
                    "--prop hexnumber:0xcafe "
                    "--prop hexnumber_capital:0xCAFE "
                    "--prop large_hexnumber:0xfedcba9876543210 "
                    "--prop larger_than_uint64:0xfedcba98765432101 "
                    "--prop almost_a_number:423x "
                    "--prop_from_file blob:test/small_blob.bin "
                    "--prop_from_file large_blob:test/dummy_kernel.bin");

  base::FilePath info_path = testdir_.Append("info_output.txt");
  EXPECT_COMMAND(0,
                 "./bvbtool info_boot_image --image %s --output %s",
                 boot_image_path_.value().c_str(),
                 info_path.value().c_str());

  std::string info_data;
  ASSERT_TRUE(base::ReadFileToString(info_path, &info_data));

  ASSERT_EQ(
      "Boot Image version:       1.0\n"
      "Header Block:             8192 bytes\n"
      "Authentication Block:     4096 bytes\n"
      "Auxilary Block:           70080 bytes\n"
      "Payload Block:            196608 bytes\n"
      "Algorithm:                SHA256_RSA2048\n"
      "Rollback Index:           0\n"
      "Kernel:                   65536 bytes\n"
      "Initrd:                   131072 bytes\n"
      "Kernel Load Address:      0x10008000\n"
      "Initrd Load Address:      0x11000000\n"
      "Kernel Cmdline:           foobar=cmdline test=42\n"
      "Properties:\n"
      "    foo: 'brillo'\n"
      "    bar: 'chromeos'\n"
      "    prisoner: '24601'\n"
      "    hexnumber: '0xcafe'\n"
      "    hexnumber_capital: '0xCAFE'\n"
      "    large_hexnumber: '0xfedcba9876543210'\n"
      "    larger_than_uint64: '0xfedcba98765432101'\n"
      "    almost_a_number: '423x'\n"
      "    blob: '\\x00\\x00brillo "
      "ftw!\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\n'\n"
      "    large_blob: (65536 bytes)\n",
      info_data);
}

TEST_F(BvbToolTest, ImageHashes) {
  // test/dummy_rootfs.bin is a 1,049,600 byte (1025 Kib) file.
  EXPECT_COMMAND(
      0, "cp test/dummy_rootfs.bin %s/rootfs.bin", testdir_.value().c_str());
  EXPECT_COMMAND(
      0,
      "./bvbtool add_image_hashes --salt d00df00d  --image %s/rootfs.bin",
      testdir_.value().c_str());

  // We don't want to impose the requirement of having the
  // veritysetup(1) command available on builders but leave it here so
  // it can be manually enabled when making changes.
  //
  // (The fact that we end up with the correct root hash and tree size
  // demonstrates correctness sufficiently well.)
  if (false) {
    EXPECT_COMMAND(0,
                   "veritysetup --no-superblock --format=1 --hash=sha1 "
                   "--data-block-size=4096 --hash-block-size=4096 "
                   "--salt=d00df00d "
                   "--data-blocks=257 "
                   "--hash-offset=1052672 "
                   "verify "
                   "%s/rootfs.bin %s/rootfs.bin "
                   "d0e3b9865f45fc66c1a64796dae1666647103f72",
                   testdir_.value().c_str(),
                   testdir_.value().c_str());
  }

  base::FilePath info_path = testdir_.Append("info_output.txt");
  EXPECT_COMMAND(
      0,
      "./bvbtool info_image_hashes --image %s/rootfs.bin --output %s",
      testdir_.value().c_str(),
      info_path.value().c_str());

  std::string info_data;
  ASSERT_TRUE(base::ReadFileToString(info_path, &info_data));

  // Note that image size is rounded up to block size (4096).
  ASSERT_EQ(
      "Footer version:        1.0\n"
      "Version of dm-verity:  1\n"
      "Image Size:            1052672 bytes\n"
      "Tree Offset:           1052672\n"
      "Tree Size:             16384 bytes\n"
      "Data Block Size:       4096 bytes\n"
      "Hash Block Size:       4096 bytes\n"
      "Hash Algorithm:        sha1\n"
      "Salt:                  d00df00d\n"
      "Root Hash:             d0e3b9865f45fc66c1a64796dae1666647103f72\n",
      info_data);

  // Check that bvbtool injects the directives for setting up the
  // rootfs for the given integrity-checked file system BEFORE the
  // user-supplied command-line.
  GenerateBootImage("SHA256_RSA2048",
                    "some_option=42",
                    0,
                    base::FilePath("test/testkey_rsa2048.pem"),
                    base::StringPrintf("--rootfs_with_hashes %s/rootfs.bin",
                                       testdir_.value().c_str()));
  BvbBootImageHeader h;
  bvb_boot_image_header_to_host_byte_order(
      reinterpret_cast<BvbBootImageHeader*>(boot_image_.data()), &h);

  EXPECT_EQ(
      "dm=\"1 vroot none ro 1,"
      "0 2056 verity 1 PARTUUID=$(ANDROID_SYSTEM_PARTUUID) "
      "PARTUUID=$(ANDROID_SYSTEM_PARTUUID) 4096 4096 257 257 sha1 "
      "d0e3b9865f45fc66c1a64796dae1666647103f72 d00df00d\" "
      "some_option=42",
      std::string(reinterpret_cast<const char*>(h.kernel_cmdline)));
}