// Copyright 2025 Google LLC
//
// 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 <flasher/file/memory.hpp>
#include <flashupdate/validator/cr51.hpp>

#include <memory>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>

using ::testing::_;
using ::testing::Invoke;
using ::testing::Return;

namespace google::cr51
{

// Mock for the Cr51SignValidator dependency
class MockCr51SignValidator : public Cr51SignValidator
{
  public:
    MOCK_METHOD(std::span<const uint8_t>, hashDescriptor,
                (struct libcr51sign_ctx*, std::span<std::byte>), (override));
    MOCK_METHOD(std::optional<struct libcr51sign_validated_regions>,
                validateDescriptor, (struct libcr51sign_ctx*, bool),
                (override));
    MOCK_METHOD(bool, isBiosKeyRotationSupport, (), (override));
    MOCK_METHOD(bool, isDevSignedImageAllowed, (), (override));
};

} // namespace google::cr51

namespace flashupdate::validator::cr51
{

class Cr51VerifyTest : public ::testing::Test
{
  protected:
    Cr51VerifyTest() :
        cr51ValidatorMock(
            std::make_unique<
                ::testing::StrictMock<google::cr51::MockCr51SignValidator>>()),
        cr51(*cr51ValidatorMock)
    {
        // Provide some dummy data for the reader
        std::vector<std::byte> imageData(1024, std::byte{0});
        reader.writeAtExact(imageData, 0);
    }

    flasher::file::Memory reader;
    const std::vector<std::string> fallbackKeys = {"key1.pem", "key2.pem"};
    std::unique_ptr<google::cr51::MockCr51SignValidator> cr51ValidatorMock;
    Cr51 cr51;
};

TEST_F(Cr51VerifyTest, RoTSupported_ValidationSucceeds)
{
    EXPECT_CALL(*cr51ValidatorMock, isBiosKeyRotationSupport())
        .WillOnce(Return(true));

    // Neither of these should be called if RoT validation succeeds.
    EXPECT_CALL(*cr51ValidatorMock, isDevSignedImageAllowed()).Times(0);
    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, false)).Times(0);

    auto successfulValidation = [](libcr51sign_ctx* ctx, bool /*useRoT*/) {
        ctx->descriptor.descriptor_offset = 0;
        ctx->descriptor.descriptor_area_size = 128;
        ctx->descriptor.image_major = 1;
        ctx->descriptor.image_minor = 2;
        ctx->descriptor.image_point = 3;
        ctx->descriptor.image_subpoint = 4;
        ctx->descriptor.hash_type = HASH_SHA2_256;
        return libcr51sign_validated_regions{};
    };

    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, true))
        .WillOnce(Invoke(successfulValidation));
    EXPECT_CALL(*cr51ValidatorMock, hashDescriptor(_, _))
        .WillOnce(Return(std::span<const uint8_t>{}));
    EXPECT_TRUE(cr51.validateImage(reader, fallbackKeys));
}

TEST_F(Cr51VerifyTest, RoTSupported_RoTValidationFails_AndDevNotAllowed)
{
    EXPECT_CALL(*cr51ValidatorMock, isBiosKeyRotationSupport())
        .WillOnce(Return(true));
    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, true))
        .WillOnce(Return(std::nullopt));
    EXPECT_CALL(*cr51ValidatorMock, isDevSignedImageAllowed())
        .WillOnce(Return(false));
    // Fallback validation should not be called.
    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, false)).Times(0);
    EXPECT_FALSE(cr51.validateImage(reader, fallbackKeys));
}

TEST_F(Cr51VerifyTest,
       RoTSupported_RoTValidationFails_DevAllowed_FallbackSucceeds)
{
    EXPECT_CALL(*cr51ValidatorMock, isBiosKeyRotationSupport())
        .WillRepeatedly(Return(true));
    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, true))
        .WillOnce(Return(std::nullopt));
    EXPECT_CALL(*cr51ValidatorMock, isDevSignedImageAllowed())
        .WillOnce(Return(true));

    auto successfulValidation = [](libcr51sign_ctx* ctx, bool /*useRoT*/) {
        ctx->descriptor.descriptor_offset = 0;
        ctx->descriptor.descriptor_area_size = 128;
        ctx->descriptor.image_major = 1;
        ctx->descriptor.image_minor = 2;
        ctx->descriptor.image_point = 3;
        ctx->descriptor.image_subpoint = 4;
        ctx->descriptor.hash_type = HASH_SHA2_256;
        return libcr51sign_validated_regions{};
    };

    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, false))
        .WillOnce(Invoke(successfulValidation));
    EXPECT_CALL(*cr51ValidatorMock, hashDescriptor(_, _))
        .WillOnce(Return(std::span<const uint8_t>{}));
    EXPECT_TRUE(cr51.validateImage(reader, fallbackKeys));
}

TEST_F(Cr51VerifyTest, RoTSupported_RoTValidationFails_DevAllowed_FallbackFails)
{
    EXPECT_CALL(*cr51ValidatorMock, isBiosKeyRotationSupport())
        .WillRepeatedly(Return(true));
    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, true))
        .WillOnce(Return(std::nullopt));
    EXPECT_CALL(*cr51ValidatorMock, isDevSignedImageAllowed())
        .WillOnce(Return(true));

    // Fallback validation also fails.
    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, false))
        .Times(fallbackKeys.size())
        .WillRepeatedly(Return(std::nullopt));

    EXPECT_FALSE(cr51.validateImage(reader, fallbackKeys));
}

TEST_F(Cr51VerifyTest, RoTNotSupported_FallbackSucceeds)
{
    EXPECT_CALL(*cr51ValidatorMock, isBiosKeyRotationSupport())
        .WillRepeatedly(Return(false));
    EXPECT_CALL(*cr51ValidatorMock, isDevSignedImageAllowed()).Times(0);

    // The RoT path validateDescriptor(_, true) should not be called.
    auto successfulValidation = [](libcr51sign_ctx* ctx, bool /*useRoT*/) {
        ctx->descriptor.descriptor_offset = 0;
        ctx->descriptor.descriptor_area_size = 128;
        ctx->descriptor.image_major = 1;
        ctx->descriptor.image_minor = 2;
        ctx->descriptor.image_point = 3;
        ctx->descriptor.image_subpoint = 4;
        ctx->descriptor.hash_type = HASH_SHA2_256;
        return libcr51sign_validated_regions{};
    };

    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, false))
        .WillOnce(Invoke(successfulValidation));
    EXPECT_CALL(*cr51ValidatorMock, hashDescriptor(_, _))
        .WillOnce(Return(std::span<const uint8_t>{}));
    EXPECT_TRUE(cr51.validateImage(reader, fallbackKeys));
}

TEST_F(Cr51VerifyTest, RoTNotSupported_FallbackFails)
{
    EXPECT_CALL(*cr51ValidatorMock, isBiosKeyRotationSupport())
        .WillRepeatedly(Return(false));

    // The RoT path validateDescriptor(_, true) should not be called.
    EXPECT_CALL(*cr51ValidatorMock, validateDescriptor(_, false))
        .Times(fallbackKeys.size())
        .WillRepeatedly(Return(std::nullopt));
    EXPECT_FALSE(cr51.validateImage(reader, fallbackKeys));
}

} // namespace flashupdate::validator::cr51