/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
 *     Copyright 2016 Couchbase, Inc.
 *
 *   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 "scram-sha.h"

#include "core/crypto/cbcrypto.h"
#include "core/logger/logger.hxx"
#include "core/platform/base64.h"
#include "core/platform/random.h"
#include "core/platform/string_hex.h"
#include "stringutils.h"

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

namespace couchbase::core::sasl::mechanism::scram
{

using AttributeMap = std::map<char, std::string>;

/**
 * Decode the attribute list into a set. The attribute list looks like:
 * "k=value,y=value" etc
 *
 * @param list the list to parse
 * @param attributes where to store the attributes
 * @return true if success, false otherwise
 */
static bool
decodeAttributeList(const std::string& list, AttributeMap& attributes)
{
    size_t pos = 0;

    CB_LOG_TRACE("decoding attribute list [{}]", list);

    while (pos < list.length()) {
        auto equal = list.find('=', pos);
        if (equal == std::string::npos) {
            // syntax error!!
            CB_LOG_ERROR("decode attribute list [{}] failed: no '='", list);
            return false;
        }

        if ((equal - pos) != 1) {
            CB_LOG_ERROR("decode attribute list [{}] failed: no key is multichar", list);
            return false;
        }

        char key = list.at(pos);
        pos = equal + 1;

        // Make sure we haven't seen this key before..
        if (attributes.find(key) != attributes.end()) {
            CB_LOG_ERROR("decode attribute list [{}] failed: key [{}] is multichar", list, key);
            return false;
        }

        auto comma = list.find(',', pos);
        if (comma == std::string::npos) {
            attributes.try_emplace(key, list.substr(pos));
            pos = list.length();
        } else {
            attributes.try_emplace(key, list.substr(pos, comma - pos));
            pos = comma + 1;
        }
    }

    return true;
}

/********************************************************************
 * Common API
 *******************************************************************/
std::string
ScramShaBackend::getAuthMessage()
{
    if (client_first_message_bare.empty()) {
        throw std::logic_error("can't call getAuthMessage without client_first_message_bare is set");
    }
    if (server_first_message.empty()) {
        throw std::logic_error("can't call getAuthMessage without server_first_message is set");
    }
    if (client_final_message_without_proof.empty()) {
        throw std::logic_error("can't call getAuthMessage without client_final_message_without_proof is set");
    }
    return client_first_message_bare + "," + server_first_message + "," + client_final_message_without_proof;
}

void
ScramShaBackend::addAttribute(std::ostream& out, char key, const std::string& value, bool more)
{
    out << key << '=';

    switch (key) {
        case 'n': // username ..
            out << encode_username(sasl_prep(value));
            break;

        case 'r': // client nonce.. printable characters
            for (const auto& c : value) {
                if (c == ',' || (isprint(c) == 0)) {
                    throw std::invalid_argument("ScramShaBackend::addAttribute: Invalid character in client nonce");
                }
            }
            out << value;
            break;

        case 'c': // base64 encoded GS2 header and channel binding data
        case 's': // base64 encoded salt
        case 'p': // base64 encoded client proof
        case 'v': // base64 encoded server signature
            out << couchbase::core::base64::encode(value);
            break;

        case 'i': // iterator count
            // validate that it is an integer value
            try {
                (void)std::stoi(value);
            } catch (...) {
                throw std::invalid_argument("ScramShaBackend::addAttribute: Iteration count must be a numeric value");
            }
            out << value;
            break;

        case 'e':
            for (const auto& c : value) {
                if (c == ',' || (isprint(c) == 0)) {
                    throw std::invalid_argument("ScramShaBackend::addAttribute: Invalid character in error message");
                }
            }
            out << value;
            break;

        default:
            throw std::invalid_argument("ScramShaBackend::addAttribute: Invalid key");
    }

    if (more) {
        out << ',';
    }
}

void
ScramShaBackend::addAttribute(std::ostream& out, char key, int value, bool more)
{
    out << key << '=';

    std::string base64_encoded;

    switch (key) {
        case 'n': // username ..
        case 'r': // client nonce.. printable characters
        case 'c': // base64 encoded GS2 header and channel binding data
        case 's': // base64 encoded salt
        case 'p': // base64 encoded client proof
        case 'v': // base64 encoded server signature
        case 'e': // error message
            throw std::invalid_argument("ScramShaBackend::addAttribute: Invalid value (should not be int)");

        case 'i': // iterator count
            out << value;
            break;

        default:
            throw std::invalid_argument("ScramShaBackend::addAttribute: Invalid key");
    }

    if (more) {
        out << ',';
    }
}

/**
 * Generate the Server Signature. It is computed as:
 *
 * SaltedPassword  := Hi(Normalize(password), salt, i)
 * ServerKey       := HMAC(SaltedPassword, "Server Key")
 * ServerSignature := HMAC(ServerKey, AuthMessage)
 */
std::string
ScramShaBackend::getServerSignature()
{
    auto serverKey = couchbase::core::crypto::HMAC(algorithm, getSaltedPassword(), "Server Key");

    return couchbase::core::crypto::HMAC(algorithm, serverKey, getAuthMessage());
}

/**
 * Generate the Client Proof. It is computed as:
 *
 * SaltedPassword  := Hi(Normalize(password), salt, i)
 * ClientKey       := HMAC(SaltedPassword, "Client Key")
 * StoredKey       := H(ClientKey)
 * AuthMessage     := client-first-message-bare + "," +
 *                    server-first-message + "," +
 *                    client-final-message-without-proof
 * ClientSignature := HMAC(StoredKey, AuthMessage)
 * ClientProof     := ClientKey XOR ClientSignature
 */
std::string
ScramShaBackend::getClientProof()
{
    auto clientKey = couchbase::core::crypto::HMAC(algorithm, getSaltedPassword(), "Client Key");
    auto storedKey = couchbase::core::crypto::digest(algorithm, clientKey);
    std::string authMessage = getAuthMessage();
    auto clientSignature = couchbase::core::crypto::HMAC(algorithm, storedKey, authMessage);

    // Client Proof is ClientKey XOR ClientSignature
    const auto* ck = clientKey.data();
    const auto* cs = clientSignature.data();

    std::string proof;
    proof.resize(clientKey.size());

    auto total = proof.size();
    for (std::size_t ii = 0; ii < total; ++ii) {
        proof[ii] = ck[ii] ^ cs[ii];
    }

    return proof;
}

ClientBackend::ClientBackend(GetUsernameCallback& user_cb,
                             GetPasswordCallback& password_cb,
                             ClientContext& ctx,
                             Mechanism mech,
                             couchbase::core::crypto::Algorithm algo)
  : MechanismBackend(user_cb, password_cb, ctx)
  , ScramShaBackend(mech, algo)
{
    couchbase::core::RandomGenerator randomGenerator;

    std::array<char, 8> nonce{};
    if (!randomGenerator.getBytes(nonce.data(), nonce.size())) {
        CB_LOG_ERROR_RAW("failed to generate server nonce");
        throw std::bad_alloc();
    }

    clientNonce = couchbase::core::to_hex({ nonce.data(), nonce.size() });
}

std::pair<error, std::string_view>
ClientBackend::start()
{
    std::stringstream out;
    out << "n,,";
    addAttribute(out, 'n', usernameCallback(), true);
    addAttribute(out, 'r', clientNonce, false);

    client_first_message = out.str();
    client_first_message_bare = client_first_message.substr(3); // skip n,,

    return { error::OK, client_first_message };
}

std::pair<error, std::string_view>
ClientBackend::step(std::string_view input)
{
    if (input.empty()) {
        return { error::BAD_PARAM, {} };
    }

    if (server_first_message.empty()) {
        server_first_message.assign(input.data(), input.size());

        AttributeMap attributes;
        if (!decodeAttributeList(server_first_message, attributes)) {
            return { error::BAD_PARAM, {} };
        }

        for (const auto& attribute : attributes) {
            switch (attribute.first) {
                case 'r': // combined nonce
                    nonce_ = attribute.second;
                    break;
                case 's':
                    salt = couchbase::core::base64::decode_to_string(attribute.second);
                    break;
                case 'i':
                    try {
                        iterationCount = static_cast<unsigned int>(std::stoul(attribute.second));
                    } catch (...) {
                        return { error::BAD_PARAM, {} };
                    }
                    break;
                default:
                    return { error::BAD_PARAM, {} };
            }
        }

        if (attributes.find('r') == attributes.end() || attributes.find('s') == attributes.end() ||
            attributes.find('i') == attributes.end()) {
            CB_LOG_ERROR_RAW("missing r/s/i in server message");
            return { error::BAD_PARAM, {} };
        }

        // I've got the SALT, lets generate the salted password
        if (!generateSaltedPassword(passwordCallback())) {
            CB_LOG_ERROR_RAW("failed to generated salted password");
            return { error::FAIL, {} };
        }

        // Ok so we have salted hased password :D

        std::stringstream out;
        addAttribute(out, 'c', "n,,", true);
        addAttribute(out, 'r', nonce_, false);
        client_final_message_without_proof = out.str();
        out << ",";

        addAttribute(out, 'p', getClientProof(), false);

        client_final_message = out.str();

        return { error::CONTINUE, client_final_message };
    }
    server_final_message.assign(input.data(), input.size());

    AttributeMap attributes;
    if (!decodeAttributeList(server_final_message, attributes)) {
        CB_LOG_ERROR_RAW("SCRAM: failed to decode server-final-message");
        return { error::BAD_PARAM, {} };
    }

    if (attributes.find('e') != attributes.end()) {
        CB_LOG_ERROR("failed to authenticate: {}", attributes['e']);
        return { error::FAIL, {} };
    }

    if (attributes.find('v') == attributes.end()) {
        CB_LOG_ERROR_RAW("syntax error server final message is missing 'v'");
        return { error::BAD_PARAM, {} };
    }

    if (auto encoded = couchbase::core::base64::encode(getServerSignature()); encoded != attributes['v']) {
        CB_LOG_ERROR_RAW("incorrect ServerKey received");
        return { error::FAIL, {} };
    }

    return { error::OK, {} };
}

bool
ClientBackend::generateSaltedPassword(const std::string& secret)
{
    try {
        saltedPassword = couchbase::core::crypto::PBKDF2_HMAC(algorithm, secret, salt, iterationCount);
        return true;
    } catch (...) {
        return false;
    }
}

} // namespace couchbase::core::sasl::mechanism::scram