/*
* TLS Policy tests
*
* (C) 2016 Juraj Somorovsky
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include "tests.h"

#if defined(BOTAN_HAS_TLS)
   #include <botan/assert.h>
   #include <botan/tls_ciphersuite.h>
   #include <botan/tls_exceptn.h>
   #include <botan/tls_policy.h>
#endif

#if defined(BOTAN_HAS_RSA)
   #include <botan/rsa.h>
#endif

#if defined(BOTAN_HAS_ECC_GROUP)
   #include <botan/ec_group.h>
#endif

#if defined(BOTAN_HAS_ECDH)
   #include <botan/ecdh.h>
#endif

#if defined(BOTAN_HAS_ECDSA)
   #include <botan/ecdsa.h>
#endif

#if defined(BOTAN_HAS_DIFFIE_HELLMAN)
   #include <botan/dh.h>
   #include <botan/dl_group.h>
#endif

namespace Botan_Tests {

namespace {

#if defined(BOTAN_HAS_TLS)
class TLS_Policy_Unit_Tests final : public Test {
   public:
      std::vector<Test::Result> run() override {
         std::vector<Test::Result> results;

         results.push_back(test_peer_key_acceptable_rsa(this->rng()));
         results.push_back(test_peer_key_acceptable_ecdh(this->rng()));
         results.push_back(test_peer_key_acceptable_ecdsa(this->rng()));
         results.push_back(test_peer_key_acceptable_dh());
         results.push_back(test_key_exchange_groups_to_offer());
         results.push_back(test_require_extended_master_secret());
         results.push_back(test_dtls_refuses_weak_ciphersuites());

         return results;
      }

      // RFC 9147 4.5.3 forbids _CCM_8 in DTLS absent forgery safeguards Botan does
      // not implement, and CBC+HMAC keeps a residual timing channel that DTLS
      // makes observable many times per association. Both are refused for DTLS
      // while remaining available to stream TLS.
      static Test::Result test_dtls_refuses_weak_ciphersuites() {
         Test::Result result("TLS Policy DTLS ciphersuite refusals");

         class Permissive_Policy final : public Botan::TLS::Policy {
            public:
               std::vector<std::string> allowed_ciphers() const override {
                  return {"AES-256/CCM(8)", "AES-128/CCM(8)", "AES-256/CCM", "AES-128/CCM", "AES-128/GCM", "AES-128"};
               }

               std::vector<std::string> allowed_macs() const override { return {"AEAD", "SHA-256", "SHA-1"}; }

               bool allow_tls12() const override { return true; }

               bool allow_dtls12() const override { return true; }
         };

         const Permissive_Policy policy;

         auto classify = [&](Botan::TLS::Protocol_Version version) {
            size_t ccm_8 = 0;
            size_t cbc = 0;
            for(const auto id : policy.ciphersuite_list(version)) {
               const auto suite = Botan::TLS::Ciphersuite::by_id(id);
               if(!suite.has_value()) {
                  continue;
               }
               if(suite->uses_short_authentication_tag()) {
                  ccm_8 += 1;
               }
               if(suite->cbc_ciphersuite()) {
                  cbc += 1;
               }
            }
            return std::make_pair(ccm_8, cbc);
         };

         const auto [tls_ccm_8, tls_cbc] = classify(Botan::TLS::Protocol_Version::TLS_V12);
         const auto [dtls_ccm_8, dtls_cbc] = classify(Botan::TLS::Protocol_Version::DTLS_V12);

         // Guard against the test passing because the policy offered none at all.
   #if defined(BOTAN_HAS_AEAD_CCM)
         result.test_is_true("TLS 1.2 still offers CCM_8 suites", tls_ccm_8 > 0);
   #endif
   #if defined(BOTAN_HAS_TLS_CBC)
         result.test_is_true("TLS 1.2 still offers CBC suites", tls_cbc > 0);
   #endif
         BOTAN_UNUSED(tls_ccm_8, tls_cbc);

         result.test_sz_eq("DTLS 1.2 offers no CCM_8 suite", dtls_ccm_8, 0);
         result.test_sz_eq("DTLS 1.2 offers no CBC suite", dtls_cbc, 0);
         result.test_sz_gte("DTLS 1.2 still has usable suites",
                            policy.ciphersuite_list(Botan::TLS::Protocol_Version::DTLS_V12).size(),
                            1);

         return result;
      }

      static Test::Result test_require_extended_master_secret() {
         Test::Result result("TLS Policy require_extended_master_secret");

         const Botan::TLS::Policy default_policy;
         result.test_is_true("default Policy requires EMS", default_policy.require_extended_master_secret());

         using TP = Botan::TLS::Text_Policy;
         result.test_is_true("default text policy requires EMS", TP("").require_extended_master_secret());
         result.test_is_true("text policy override disables EMS requirement",
                             !TP("require_extended_master_secret = false").require_extended_master_secret());
         result.test_is_true("text policy override re-enables EMS requirement",
                             TP("require_extended_master_secret = true").require_extended_master_secret());
         return result;
      }

   private:
      static Test::Result test_peer_key_acceptable_rsa([[maybe_unused]] Botan::RandomNumberGenerator& rng) {
         Test::Result result("TLS Policy RSA key verification");
   #if defined(BOTAN_HAS_RSA)
         auto rsa_key_1024 = std::make_unique<Botan::RSA_PrivateKey>(rng, 1024);
         const Botan::TLS::Policy policy;

         try {
            policy.check_peer_key_acceptable(*rsa_key_1024);
            result.test_failure("Incorrectly accepting 1024 bit RSA keys");
         } catch(Botan::TLS::TLS_Exception&) {
            result.test_success("Correctly rejecting 1024 bit RSA keys");
         }

         auto rsa_key_2048 = std::make_unique<Botan::RSA_PrivateKey>(rng, 2048);
         policy.check_peer_key_acceptable(*rsa_key_2048);
         result.test_success("Correctly accepting 2048 bit RSA keys");
   #endif
         return result;
      }

      static Test::Result test_peer_key_acceptable_ecdh([[maybe_unused]] Botan::RandomNumberGenerator& rng) {
         Test::Result result("TLS Policy ECDH key verification");
   #if defined(BOTAN_HAS_ECDH)

         const Botan::TLS::Policy policy;

         if(Botan::EC_Group::supports_named_group("secp192r1")) {
            const auto group_192 = Botan::EC_Group::from_name("secp192r1");
            auto ecdh_192 = std::make_unique<Botan::ECDH_PrivateKey>(rng, group_192);

            try {
               policy.check_peer_key_acceptable(*ecdh_192);
               result.test_failure("Incorrectly accepting 192 bit EC keys");
            } catch(Botan::TLS::TLS_Exception&) {
               result.test_success("Correctly rejecting 192 bit EC keys");
            }
         }

         if(Botan::EC_Group::supports_named_group("secp256r1")) {
            const auto group_256 = Botan::EC_Group::from_name("secp256r1");
            auto ecdh_256 = std::make_unique<Botan::ECDH_PrivateKey>(rng, group_256);
            policy.check_peer_key_acceptable(*ecdh_256);
            result.test_success("Correctly accepting 256 bit EC keys");
         }
   #endif
         return result;
      }

      static Test::Result test_peer_key_acceptable_ecdsa([[maybe_unused]] Botan::RandomNumberGenerator& rng) {
         Test::Result result("TLS Policy ECDSA key verification");
   #if defined(BOTAN_HAS_ECDSA)
         const Botan::TLS::Policy policy;

         if(Botan::EC_Group::supports_named_group("secp192r1")) {
            const auto group_192 = Botan::EC_Group::from_name("secp192r1");
            auto ecdsa_192 = std::make_unique<Botan::ECDSA_PrivateKey>(rng, group_192);

            try {
               policy.check_peer_key_acceptable(*ecdsa_192);
               result.test_failure("Incorrectly accepting 192 bit EC keys");
            } catch(Botan::TLS::TLS_Exception&) {
               result.test_success("Correctly rejecting 192 bit EC keys");
            }
         }

         if(Botan::EC_Group::supports_named_group("secp256r1")) {
            const auto group_256 = Botan::EC_Group::from_name("secp256r1");
            auto ecdsa_256 = std::make_unique<Botan::ECDSA_PrivateKey>(rng, group_256);
            policy.check_peer_key_acceptable(*ecdsa_256);
            result.test_success("Correctly accepting 256 bit EC keys");
         }
   #endif
         return result;
      }

      static Test::Result test_peer_key_acceptable_dh() {
         Test::Result result("TLS Policy DH key verification");
   #if defined(BOTAN_HAS_DIFFIE_HELLMAN)
         const BigInt g("2");
         const BigInt p("58458002095536094658683755258523362961421200751439456159756164191494576279467");
         const Botan::DL_Group grp(p, g);
         const Botan::BigInt x("46205663093589612668746163860870963912226379131190812163519349848291472898748");
         auto dhkey = std::make_unique<Botan::DH_PrivateKey>(grp, x);

         const Botan::TLS::Policy policy;
         try {
            policy.check_peer_key_acceptable(*dhkey);
            result.test_failure("Incorrectly accepting short bit DH keys");
         } catch(Botan::TLS::TLS_Exception&) {
            result.test_success("Correctly rejecting short bit DH keys");
         }
   #endif
         return result;
      }

      static Test::Result test_key_exchange_groups_to_offer() {
         Test::Result result("TLS Policy key share offering");

         const Botan::TLS::Policy default_policy;
         result.test_sz_eq(
            "default TLS Policy offers exactly one", default_policy.key_exchange_groups_to_offer().size(), 1);
         result.test_is_true(
            "default TLS Policy offers preferred group",
            default_policy.key_exchange_groups().front() == default_policy.key_exchange_groups_to_offer().front());

         using TP = Botan::TLS::Text_Policy;

         result.test_sz_eq(
            "default behaviour from text policy (size)", TP("").key_exchange_groups_to_offer().size(), 1);
         result.test_is_true("default behaviour from text policy (preferred)",
                             TP("").key_exchange_groups().front() == TP("").key_exchange_groups_to_offer().front());

         result.test_is_true("no offerings",
                             TP("key_exchange_groups_to_offer = none").key_exchange_groups_to_offer().empty());

         const std::string two_groups = "key_exchange_groups_to_offer = secp256r1 ffdhe/ietf/4096";
         result.test_sz_eq("list of offerings (size)", TP(two_groups).key_exchange_groups_to_offer().size(), 2);
         result.test_is_true("list of offerings (0)",
                             TP(two_groups).key_exchange_groups_to_offer()[0] == Botan::TLS::Group_Params::SECP256R1);
         result.test_is_true("list of offerings (1)",
                             TP(two_groups).key_exchange_groups_to_offer()[1] == Botan::TLS::Group_Params::FFDHE_4096);

         return result;
      }
};

BOTAN_REGISTER_TEST("tls", "tls_policy", TLS_Policy_Unit_Tests);

#endif

}  // namespace

}  // namespace Botan_Tests
