/*

 * Copyright (c) 2023, Sam Atkins <atkinssj@serenityos.org>

 *

 * SPDX-License-Identifier: BSD-2-Clause

 */

#include <AK/Array.h>

#include <AK/Base64.h>

#include <AK/Vector.h>

#include <LibCrypto/Hash/SHA2.h>

#include <LibWeb/SRI/SRI.h>

namespace Web::SRI {

constexpr Array supported_hash_functions {

    // These are sorted by strength, low to high.

    // NOTE: We are specifically told to refuse MD5 and SHA1.

    //       https://w3c.github.io/webappsec-subresource-integrity/#hash-functions

    "sha256"sv,

    "sha384"sv,

    "sha512"sv,

};

// https://w3c.github.io/webappsec-subresource-integrity/#getprioritizedhashfunction

static StringView get_prioritized_hash_function(StringView a, StringView b)

{

    if (a == b)

        return ""sv;

    auto a_priority = supported_hash_functions.first_index_of(a).value();

    auto b_priority = supported_hash_functions.first_index_of(b).value();

    if (a_priority > b_priority)

        return a;

    return b;

}

// https://w3c.github.io/webappsec-subresource-integrity/#apply-algorithm-to-response

ErrorOr<String> apply_algorithm_to_bytes(StringView algorithm, ByteBuffer const& bytes)

{

    // NOTE: The steps are duplicated here because each hash algorithm returns a different result type.

    if (algorithm == "sha256"sv) {

        // 1. Let result be the result of applying algorithm to bytes.

        auto result = Crypto::Hash::SHA256::hash(bytes);

        // 2. Return the result of base64 encoding result.

        return encode_base64(result.bytes());

    }

    if (algorithm == "sha384"sv) {

        // 1. Let result be the result of applying algorithm to bytes.

        auto result = Crypto::Hash::SHA384::hash(bytes);

        // 2. Return the result of base64 encoding result.

        return encode_base64(result.bytes());

    }

    if (algorithm == "sha512"sv) {

        // 1. Let result be the result of applying algorithm to bytes.

        auto result = Crypto::Hash::SHA512::hash(bytes);

        // 2. Return the result of base64 encoding result.

        return encode_base64(result.bytes());

    }

    VERIFY_NOT_REACHED();

}

// https://w3c.github.io/webappsec-subresource-integrity/#parse-metadata

ErrorOr<Vector<Metadata>> parse_metadata(StringView metadata)

{

    // 1. Let result be the empty set.

    Vector<Metadata> result;

    // 2. For each item returned by splitting metadata on spaces:

    TRY(metadata.for_each_split_view(' ', SplitBehavior::Nothing, [&](StringView item) -> ErrorOr<void> {

        // 1. Let hash-with-opt-token-list be the result of splitting item on U+003F (?).

        auto hash_with_opt_token_list = item.split_view('?');

        // 2. Let hash-expression be hash-with-opt-token-list[0].

        auto hash_expression = hash_with_opt_token_list[0];

        // 3. Let base64-value be the empty string.

        StringView base64_value;

        // 4. Let hash-expr-token-list be the result of splitting hash-expression on U+002D (-).

        auto hash_expr_token_list = hash_expression.split_view('-');

        // 5. Let algorithm be hash-expr-token-list[0].

        auto algorithm = hash_expr_token_list[0];

        // 6. If hash-expr-token-list[1] exists, set base64-value to hash-expr-token-list[1].

        if (hash_expr_token_list.size() >= 1)

            base64_value = hash_expr_token_list[1];

        // 7. If algorithm is not a hash function recognized by the user agent, continue.

        if (!supported_hash_functions.contains_slow(algorithm))

            return {};

        // 8. Let metadata be the ordered map «["alg" → algorithm, "val" → base64-value]».

        //    Note: Since no options are defined (see the §3.1 Integrity metadata), a corresponding entry is not set in metadata.

        //    If options are defined in a future version, hash-with-opt-token-list[1] can be utilized as options.

        auto metadata = Metadata {

            .algorithm = TRY(String::from_utf8(algorithm)),

            .base64_value = TRY(String::from_utf8(base64_value)),

            .options = {},

        };

        // 9. Append metadata to result.

        TRY(result.try_append(move(metadata)));

        return {};

    }));

    // 3. Return result.

    return result;

}

// https://w3c.github.io/webappsec-subresource-integrity/#get-the-strongest-metadata

ErrorOr<Vector<Metadata>> get_strongest_metadata_from_set(Vector<Metadata> const& set)

{

    // 1. Let result be the empty set and strongest be the empty string.

    Vector<Metadata> result;

    Optional<Metadata> strongest;

    // 2. For each item in set:

    for (auto const& item : set) {

        // 1. If result is the empty set, add item to result and set strongest to item, skip to the next item.

        if (result.is_empty()) {

            TRY(result.try_append(item));

            strongest = item;

            continue;

        }

        // 2. Let currentAlgorithm be the alg component of strongest.

        auto& current_algorithm = strongest->algorithm;

        // 3. Let newAlgorithm be the alg component of item.

        auto& new_algorithm = item.algorithm;

        // 4. If the result of getPrioritizedHashFunction(currentAlgorithm, newAlgorithm) is the empty string, add item to result.

        auto prioritized_hash_function = get_prioritized_hash_function(current_algorithm, new_algorithm);

        if (prioritized_hash_function.is_empty()) {

            TRY(result.try_append(item));

        }

        //    If the result is newAlgorithm, set strongest to item, set result to the empty set, and add item to result.

        else if (prioritized_hash_function == new_algorithm) {

            strongest = item;

            result.clear_with_capacity();

            TRY(result.try_append(item));

        }

    }

    // 3. Return result.

    return result;

}

// https://w3c.github.io/webappsec-subresource-integrity/#does-response-match-metadatalist

ErrorOr<bool> do_bytes_match_metadata_list(ByteBuffer const& bytes, StringView metadata_list)

{

    // 1. Let parsedMetadata be the result of parsing metadataList.

    auto parsed_metadata = TRY(parse_metadata(metadata_list));

    // 2. If parsedMetadata is empty set, return true.

    if (parsed_metadata.is_empty())

        return true;

    // 3. Let metadata be the result of getting the strongest metadata from parsedMetadata.

    auto metadata = TRY(get_strongest_metadata_from_set(parsed_metadata));

    // 4. For each item in metadata:

    for (auto const& item : metadata) {

        // 1. Let algorithm be the item["alg"].

        auto& algorithm = item.algorithm;

        // 2. Let expectedValue be the item["val"].

        auto& expected_value = item.base64_value;

        // 3. Let actualValue be the result of applying algorithm to bytes.

        auto actual_value = TRY(apply_algorithm_to_bytes(algorithm, bytes));

        // 4. If actualValue is a case-sensitive match for expectedValue, return true.

        if (actual_value == expected_value)

            return true;

    }

    // 5. Return false.

    return false;

}

}