/*
* (C) 2014,2015 Jack Lloyd
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include "tests.h"

#if defined(BOTAN_HAS_BCRYPT)
   #include <botan/bcrypt.h>
#endif

#if defined(BOTAN_HAS_PASSHASH9)
   #include <botan/passhash9.h>
#endif

#if defined(BOTAN_HAS_ARGON2_FMT)
   #include "test_rng.h"
   #include <botan/argon2fmt.h>
#endif

#if defined(BOTAN_HAS_PKCS12_KDF)
   #include <botan/pwdhash.h>
#endif

namespace Botan_Tests {

namespace {

#if defined(BOTAN_HAS_BCRYPT)
class Bcrypt_Tests final : public Text_Based_Test {
   public:
      Bcrypt_Tests() : Text_Based_Test("passhash/bcrypt.vec", "Password,Passhash") {}

      Test::Result run_one_test(const std::string& /*header*/, const VarMap& vars) override {
         // Encoded as binary so we can test binary inputs
         const std::vector<uint8_t> password_vec = vars.get_req_bin("Password");
         const std::string password(reinterpret_cast<const char*>(password_vec.data()), password_vec.size());

         const std::string passhash = vars.get_req_str("Passhash");

         Test::Result result("bcrypt");
         result.test_is_true("correct hash accepted", Botan::check_bcrypt(password, passhash));

         // self-test low levels for each test password
         for(uint16_t level = 4; level <= 6; ++level) {
            const std::string gen_hash = Botan::generate_bcrypt(password, this->rng(), level);
            result.test_is_true("generated hash accepted", Botan::check_bcrypt(password, gen_hash));
         }

         return result;
      }

      std::vector<Test::Result> run_final_tests() override {
         Test::Result result("bcrypt");

         const uint64_t start = Test::timestamp();

         const std::string password = "ag00d1_2BE5ur3";

         const uint16_t max_level = (Test::run_long_tests() ? 15 : 10);

         auto& rng = this->rng();

         for(uint16_t level = 4; level <= max_level; ++level) {
            const std::string gen_hash = Botan::generate_bcrypt(password, rng, level);
            result.test_is_true("generated hash accepted", Botan::check_bcrypt(password, gen_hash));
         }

         result.test_throws("Invalid bcrypt version rejected", "Unknown bcrypt version 'q'", [&rng]() {
            Botan::generate_bcrypt("pass", rng, 4, 'q');
         });

         result.set_ns_consumed(Test::timestamp() - start);

         return {result};
      }
};

BOTAN_REGISTER_TEST("pbkdf", "bcrypt", Bcrypt_Tests);

#endif

#if defined(BOTAN_HAS_ARGON2_FMT)
class Argon2_Tests final : public Text_Based_Test {
   public:
      Argon2_Tests() : Text_Based_Test("passhash/argon2.vec", "Password,Passhash", "Mode,M,T,P,Salt,OutLen") {}

      Test::Result run_one_test(const std::string& header, const VarMap& vars) override {
         const std::string password = vars.get_req_str("Password");
         const std::string passhash = vars.get_req_str("Passhash");

         Test::Result result("Argon2 password hash");

         if(header == "Verify") {
            const bool accepted = Botan::argon2_check_pwhash(password.data(), password.size(), passhash);
            result.test_is_true("correct hash accepted", accepted);
         } else if(header == "Generate") {
            const std::vector<uint8_t> salt = vars.get_req_bin("Salt");
            const uint8_t y = vars.get_req_u8("Mode");
            const size_t M = vars.get_req_sz("M");
            const size_t t = vars.get_req_sz("T");
            const size_t p = vars.get_req_sz("P");
            const size_t out_len = vars.get_req_sz("OutLen");

            Fixed_Output_RNG rng(salt);

            const std::string generated =
               Botan::argon2_generate_pwhash(password.data(), password.size(), rng, p, M, t, y, salt.size(), out_len);

            result.test_str_eq("expected hash generated", generated, passhash);
            const bool accepted = Botan::argon2_check_pwhash(password.data(), password.size(), generated);
            result.test_is_true("generated hash accepted", accepted);
         } else {
            throw Test_Error("Unexpected header in Argon2 password hash test file");
         }

         return result;
      }
};

BOTAN_REGISTER_TEST("pbkdf", "argon2_pass", Argon2_Tests);

#endif

#if defined(BOTAN_HAS_PASSHASH9)
class Passhash9_Tests final : public Text_Based_Test {
   public:
      Passhash9_Tests() : Text_Based_Test("passhash/passhash9.vec", "Password,Passhash,PRF") {}

      Test::Result run_one_test(const std::string& /*header*/, const VarMap& vars) override {
         // Encoded as binary so we can test binary inputs
         const std::vector<uint8_t> password_vec = vars.get_req_bin("Password");
         const std::string password(reinterpret_cast<const char*>(password_vec.data()), password_vec.size());

         const std::string passhash = vars.get_req_str("Passhash");
         const std::size_t prf = vars.get_req_sz("PRF");

         Test::Result result("passhash9");

         if(Botan::is_passhash9_alg_supported(uint8_t(prf))) {
            result.test_is_true("correct hash accepted", Botan::check_passhash9(password, passhash));
         }

         for(uint8_t alg_id = 0; alg_id <= 4; ++alg_id) {
            if(Botan::is_passhash9_alg_supported(alg_id)) {
               const std::string gen_hash = Botan::generate_passhash9(password, this->rng(), 2, alg_id);

               if(!result.test_is_true("generated hash accepted", Botan::check_passhash9(password, gen_hash))) {
                  result.test_note("rejected hash", gen_hash);
               }
            }
         }

         const uint16_t max_level = (Test::run_long_tests() ? 14 : 8);

         for(uint16_t level = 1; level <= max_level; ++level) {
            const uint8_t alg_id = 1;  // default used by generate_passhash9()
            if(Botan::is_passhash9_alg_supported(alg_id)) {
               const std::string gen_hash = Botan::generate_passhash9(password, this->rng(), level, alg_id);
               if(!result.test_is_true("generated hash accepted", Botan::check_passhash9(password, gen_hash))) {
                  result.test_note("rejected hash", gen_hash);
               }
            }
         }

         return result;
      }

      std::vector<Test::Result> run_final_tests() override {
         Test::Result result("passhash9");

         result.test_is_false("Unknown algorithm is unknown", Botan::is_passhash9_alg_supported(255));

         auto& rng = this->rng();

         result.test_throws("Throws if algorithm not supported",
                            "Passhash9: Algorithm id 255 is not defined",
                            [&rng]() { Botan::generate_passhash9("pass", rng, 3, 255); });

         result.test_throws(
            "Throws if iterations is too high", "Requested passhash9 work factor 513 is too large", []() {
               Botan::check_passhash9("floof", "$9$AgIB3c5J3kvAuML84sZ5hWT9WzJtiYRPLCEARaujS7I6IKbNCwp0");
            });
         return {result};
      }
};

BOTAN_REGISTER_TEST("pbkdf", "passhash9", Passhash9_Tests);

#endif

#if defined(BOTAN_HAS_PKCS12_KDF)

class PKCS12_KDF_Tests final : public Text_Based_Test {
   public:
      PKCS12_KDF_Tests() : Text_Based_Test("pbkdf/pkcs12_kdf.vec", "Passphrase,Salt,Iterations,Output") {}

      Test::Result run_one_test(const std::string& algo_spec, const VarMap& vars) override {
         const std::string passphrase = vars.get_req_str("Passphrase");
         const std::vector<uint8_t> salt = vars.get_req_bin("Salt");
         const size_t iterations = vars.get_req_sz("Iterations");
         const std::vector<uint8_t> expected = vars.get_req_bin("Output");

         Test::Result result(algo_spec);

         auto pwdhash_fam = Botan::PasswordHashFamily::create(algo_spec);
         if(!pwdhash_fam) {
            result.note_missing(algo_spec);
            return result;
         }

         auto pwdhash = pwdhash_fam->from_iterations(iterations);

         std::vector<uint8_t> derived(expected.size());
         pwdhash->hash(derived, passphrase, salt);
         result.test_bin_eq("derived key", derived, expected);

         return result;
      }
};

BOTAN_REGISTER_TEST("pbkdf", "pkcs12_kdf", PKCS12_KDF_Tests);

#endif

}  // namespace

}  // namespace Botan_Tests
