/**
 * The MIT License (MIT)
 *
 * Copyright (c) 2013-2020 Winlin
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#include <srs_app_rtc_dtls.hpp>

using namespace std;

#include <string.h>

#include <srs_kernel_log.hpp>
#include <srs_kernel_error.hpp>
#include <srs_app_config.hpp>
#include <srs_core_autofree.hpp>
#include <srs_rtmp_stack.hpp>

#include <srtp2/srtp.h>
#include <openssl/ssl.h>

// The return value of verify_callback controls the strategy of the further verification process. If verify_callback
// returns 0, the verification process is immediately stopped with "verification failed" state. If SSL_VERIFY_PEER is
// set, a verification failure alert is sent to the peer and the TLS/SSL handshake is terminated. If verify_callback
// returns 1, the verification process is continued. If verify_callback always returns 1, the TLS/SSL handshake will
// not be terminated with respect to verification failures and the connection will be established. The calling process
// can however retrieve the error code of the last verification error using SSL_get_verify_result(3) or by maintaining
// its own error storage managed by verify_callback.
// @see https://www.openssl.org/docs/man1.0.2/man3/SSL_CTX_set_verify.html
static int verify_callback(int preverify_ok, X509_STORE_CTX *ctx)
{
    // Always OK, we don't check the certificate of client,
    // because we allow client self-sign certificate.
    return 1;
}

SrsDtlsCertificate::SrsDtlsCertificate()
{
    dtls_cert = NULL;
    dtls_pkey = NULL;
    eckey = NULL;
}

SrsDtlsCertificate::~SrsDtlsCertificate()
{
    if (eckey) {
        EC_KEY_free(eckey);
    }

    if (dtls_pkey) {
        EVP_PKEY_free(dtls_pkey);
    }

    if (dtls_cert) {
        X509_free(dtls_cert);
    }
}

srs_error_t SrsDtlsCertificate::initialize()
{
    srs_error_t err = srs_success;

    // Initialize once.
    if (dtls_cert) {
        return err;
    }

#if OPENSSL_VERSION_NUMBER < 0x10100000L // v1.1.x
    // Initialize SSL library by registering algorithms
    // The SSL_library_init() and OpenSSL_add_ssl_algorithms() functions were deprecated in OpenSSL 1.1.0 by OPENSSL_init_ssl().
    // @see https://www.openssl.org/docs/man1.1.0/man3/OpenSSL_add_ssl_algorithms.html
    // @see https://web.archive.org/web/20150806185102/http://sctp.fh-muenster.de:80/dtls/dtls_udp_echo.c
    OpenSSL_add_ssl_algorithms();
#endif

    // Initialize SRTP first.
    srs_assert(srtp_init() == 0);

    // Whether use ECDSA certificate.
    ecdsa_mode = _srs_config->get_rtc_server_ecdsa();

    // Create keys by RSA or ECDSA.
    dtls_pkey = EVP_PKEY_new();
    srs_assert(dtls_pkey);
    if (!ecdsa_mode) { // By RSA
        RSA* rsa = RSA_new();
        srs_assert(rsa);

        // Initialize the big-number for private key.
        BIGNUM* exponent = BN_new();
        srs_assert(exponent);
        BN_set_word(exponent, RSA_F4);

        // Generates a key pair and stores it in the RSA structure provided in rsa.
        // @see https://www.openssl.org/docs/man1.0.2/man3/RSA_generate_key_ex.html
        int key_bits = 1024;
        RSA_generate_key_ex(rsa, key_bits, exponent, NULL);

        // @see https://www.openssl.org/docs/man1.1.0/man3/EVP_PKEY_type.html
        srs_assert(EVP_PKEY_set1_RSA(dtls_pkey, rsa) == 1);

        RSA_free(rsa);
        BN_free(exponent);
    }
    if (ecdsa_mode) { // By ECDSA, https://stackoverflow.com/a/6006898
        eckey = EC_KEY_new();
        srs_assert(eckey);

        // Should use the curves in ClientHello.supported_groups
        // For example:
        //      Supported Group: x25519 (0x001d)
        //      Supported Group: secp256r1 (0x0017)
        //      Supported Group: secp384r1 (0x0018)
        // @remark The curve NID_secp256k1 is not secp256r1, k1 != r1.
        // TODO: FIXME: Parse ClientHello and choose the curve.
        // Note that secp256r1 in openssl is called NID_X9_62_prime256v1, not NID_secp256k1
        // @see https://stackoverflow.com/questions/41950056/openssl1-1-0-b-is-not-support-secp256r1openssl-ecparam-list-curves
        EC_GROUP* ecgroup = EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1);
        //EC_GROUP* ecgroup = EC_GROUP_new_by_curve_name(NID_secp384r1);
        srs_assert(ecgroup);
#if OPENSSL_VERSION_NUMBER < 0x10100000L // v1.1.x
        // For openssl 1.0, we must set the group parameters, so that cert is ok.
        // @see https://github.com/monero-project/monero/blob/master/contrib/epee/src/net_ssl.cpp#L225
        EC_GROUP_set_asn1_flag(ecgroup, OPENSSL_EC_NAMED_CURVE);
#endif

        srs_assert(EC_KEY_set_group(eckey, ecgroup) == 1);
        srs_assert(EC_KEY_generate_key(eckey) == 1);

        // @see https://www.openssl.org/docs/man1.1.0/man3/EVP_PKEY_type.html
        srs_assert(EVP_PKEY_set1_EC_KEY(dtls_pkey, eckey) == 1);

        EC_GROUP_free(ecgroup);
    }

    // Create certificate, from previous generated pkey.
    // TODO: Support ECDSA certificate.
    dtls_cert = X509_new();
    srs_assert(dtls_cert);
    if (true) {
        X509_NAME* subject = X509_NAME_new();
        srs_assert(subject);

        int serial = rand();
        ASN1_INTEGER_set(X509_get_serialNumber(dtls_cert), serial);

        const std::string& aor = "ossrs.net";
        X509_NAME_add_entry_by_txt(subject, "CN", MBSTRING_ASC, (unsigned char *) aor.data(), aor.size(), -1, 0);

        X509_set_issuer_name(dtls_cert, subject);
        X509_set_subject_name(dtls_cert, subject);

        int expire_day = 365;
        const long cert_duration = 60*60*24*expire_day;

        X509_gmtime_adj(X509_get_notBefore(dtls_cert), 0);
        X509_gmtime_adj(X509_get_notAfter(dtls_cert), cert_duration);

        X509_set_version(dtls_cert, 2);
        srs_assert(X509_set_pubkey(dtls_cert, dtls_pkey) == 1);
        srs_assert(X509_sign(dtls_cert, dtls_pkey, EVP_sha1()) != 0);

        X509_NAME_free(subject);
    }

    // Show DTLS fingerprint
    if (true) {
        char fp[100] = {0};
        char *p = fp;
        unsigned char md[EVP_MAX_MD_SIZE];
        unsigned int n = 0;

        // TODO: FIXME: Unused variable.
        /*int r = */X509_digest(dtls_cert, EVP_sha256(), md, &n);

        for (unsigned int i = 0; i < n; i++, ++p) {
            sprintf(p, "%02X", md[i]);
            p += 2;

            if(i < (n-1)) {
                *p = ':';
            } else {
                *p = '\0';
            }
        }

        fingerprint.assign(fp, strlen(fp));
        srs_trace("fingerprint=%s", fingerprint.c_str());
    }

    return err;
}

X509* SrsDtlsCertificate::get_cert()
{
    return dtls_cert;
}

EVP_PKEY* SrsDtlsCertificate::get_public_key()
{
    return dtls_pkey;
}
    
EC_KEY* SrsDtlsCertificate::get_ecdsa_key() 
{
    return eckey;
}
    
std::string SrsDtlsCertificate::get_fingerprint()
{
    return fingerprint;
}

bool SrsDtlsCertificate::is_ecdsa()
{
    return ecdsa_mode;
}

ISrsDtlsCallback::ISrsDtlsCallback()
{
}

ISrsDtlsCallback::~ISrsDtlsCallback()
{
}

SrsDtls::SrsDtls(ISrsDtlsCallback* cb)
{
    callback = cb;
    handshake_done = false;

    role_ = SrsDtlsRoleServer;
    version_ = SrsDtlsVersionAuto;
}

SrsDtls::~SrsDtls()
{
    if (dtls_ctx) {
        SSL_CTX_free(dtls_ctx);
        dtls_ctx = NULL;
    }

    if (dtls) {
        // this function will free bio_in and bio_out
        SSL_free(dtls);
        dtls = NULL;
    }
}

SSL_CTX* SrsDtls::build_dtls_ctx()
{
    SSL_CTX* dtls_ctx;
#if OPENSSL_VERSION_NUMBER < 0x10002000L // v1.0.2
    dtls_ctx = SSL_CTX_new(DTLSv1_method());
#else
    if (version_ == SrsDtlsVersion1_0) {
        dtls_ctx = SSL_CTX_new(DTLSv1_method());
    } else if (version_ == SrsDtlsVersion1_2) {
        dtls_ctx = SSL_CTX_new(DTLSv1_2_method());
    } else {
        // SrsDtlsVersionAuto, use version-flexible DTLS methods
        dtls_ctx = SSL_CTX_new(DTLS_method());
    }
#endif

    if (_srs_rtc_dtls_certificate->is_ecdsa()) { // By ECDSA, https://stackoverflow.com/a/6006898
#if OPENSSL_VERSION_NUMBER >= 0x10002000L // v1.0.2
        // For ECDSA, we could set the curves list.
        // @see https://www.openssl.org/docs/man1.0.2/man3/SSL_CTX_set1_curves_list.html
        SSL_CTX_set1_curves_list(dtls_ctx, "P-521:P-384:P-256");
#endif

        // For openssl <1.1, we must set the ECDH manually.
        // @see https://stackoverrun.com/cn/q/10791887
#if OPENSSL_VERSION_NUMBER < 0x10100000L // v1.1.x
    #if OPENSSL_VERSION_NUMBER < 0x10002000L // v1.0.2
        SSL_CTX_set_tmp_ecdh(dtls_ctx, _srs_rtc_dtls_certificate->get_ecdsa_key());
    #else
        SSL_CTX_set_ecdh_auto(dtls_ctx, 1);
    #endif
#endif
    }

    // Setup DTLS context.
    if (true) {
        // We use "ALL", while you can use "DEFAULT" means "ALL:!EXPORT:!LOW:!aNULL:!eNULL:!SSLv2"
        // @see https://www.openssl.org/docs/man1.0.2/man1/ciphers.html
        srs_assert(SSL_CTX_set_cipher_list(dtls_ctx, "ALL") == 1);

        // Setup the certificate.
        srs_assert(SSL_CTX_use_certificate(dtls_ctx, _srs_rtc_dtls_certificate->get_cert()) == 1);
        srs_assert(SSL_CTX_use_PrivateKey(dtls_ctx, _srs_rtc_dtls_certificate->get_public_key()) == 1);

        // Server will send Certificate Request.
        // @see https://www.openssl.org/docs/man1.0.2/man3/SSL_CTX_set_verify.html
        // TODO: FIXME: Config it, default to off to make the packet smaller.
        SSL_CTX_set_verify(dtls_ctx, SSL_VERIFY_PEER | SSL_VERIFY_CLIENT_ONCE, verify_callback);
        // The depth count is "level 0:peer certificate", "level 1: CA certificate",
        // "level 2: higher level CA certificate", and so on.
        // @see https://www.openssl.org/docs/man1.0.2/man3/SSL_CTX_set_verify.html
        SSL_CTX_set_verify_depth(dtls_ctx, 4);

        // Whether we should read as many input bytes as possible (for non-blocking reads) or not.
        // @see https://www.openssl.org/docs/man1.0.2/man3/SSL_CTX_set_read_ahead.html
        SSL_CTX_set_read_ahead(dtls_ctx, 1);

        // TODO: Maybe we can use SRTP-GCM in future.
        // @see https://bugs.chromium.org/p/chromium/issues/detail?id=713701
        // @see https://groups.google.com/forum/#!topic/discuss-webrtc/PvCbWSetVAQ
        // @remark Only support SRTP_AES128_CM_SHA1_80, please read ssl/d1_srtp.c
        srs_assert(SSL_CTX_set_tlsext_use_srtp(dtls_ctx, "SRTP_AES128_CM_SHA1_80") == 0);
    }

    return dtls_ctx;
}

srs_error_t SrsDtls::initialize(std::string role, std::string version)
{
    srs_error_t err = srs_success;

    role_ = SrsDtlsRoleServer;
    if (role == "active") {
        role_ = SrsDtlsRoleClient;
    }

    if (version == "dtls1.0") {
        version_ = SrsDtlsVersion1_0;
    } else if (version == "dtls1.2") {
        version_ = SrsDtlsVersion1_2;
    } else {
        version_ = SrsDtlsVersionAuto;
    }

    dtls_ctx = build_dtls_ctx();

    // TODO: FIXME: Leak for SSL_CTX* return by build_dtls_ctx.
    if ((dtls = SSL_new(dtls_ctx)) == NULL) {
        return srs_error_new(ERROR_OpenSslCreateSSL, "SSL_new dtls");
    }

    if (role == "active") {
        // Dtls setup active, as client role.
        SSL_set_connect_state(dtls);
        SSL_set_max_send_fragment(dtls, 1500);
    } else {
        // Dtls setup passive, as server role.
        SSL_set_accept_state(dtls);
    }

    if ((bio_in = BIO_new(BIO_s_mem())) == NULL) {
        return srs_error_new(ERROR_OpenSslBIONew, "BIO_new in");
    }

    if ((bio_out = BIO_new(BIO_s_mem())) == NULL) {
        BIO_free(bio_in);
        return srs_error_new(ERROR_OpenSslBIONew, "BIO_new out");
    }

    SSL_set_bio(dtls, bio_in, bio_out);

    return err;
}

srs_error_t SrsDtls::do_handshake()
{
    srs_error_t err = srs_success;

    int ret = SSL_do_handshake(dtls);

    unsigned char *out_bio_data;
    int out_bio_len = BIO_get_mem_data(bio_out, &out_bio_data);

    int ssl_err = SSL_get_error(dtls, ret); 
    switch(ssl_err) {   
        case SSL_ERROR_NONE: {
            handshake_done = true;
            if ((callback == NULL) || ((err = callback->on_dtls_handshake_done()) != srs_success)) {
                return srs_error_wrap(err, "dtls handshake done handle");
            }
            break;
        }  

        case SSL_ERROR_WANT_READ: {   
            break;
        }   

        case SSL_ERROR_WANT_WRITE: {   
            break;
        }

        default: {   
            break;
        }   
    }   

    if (out_bio_len && callback) {
        if ((err = callback->write_dtls_data(out_bio_data, out_bio_len)) != srs_success) {
            return srs_error_wrap(err, "send dtls packet");
        }
    }

    return err;
}

srs_error_t SrsDtls::on_dtls(char* data, int nb_data)
{
    srs_error_t err = srs_success;
    if (BIO_reset(bio_in) != 1) {
        return srs_error_new(ERROR_OpenSslBIOReset, "BIO_reset");
    }
    if (BIO_reset(bio_out) != 1) {
        return srs_error_new(ERROR_OpenSslBIOReset, "BIO_reset");
    }

    if (BIO_write(bio_in, data, nb_data) <= 0) {
        // TODO: 0 or -1 maybe block, use BIO_should_retry to check.
        return srs_error_new(ERROR_OpenSslBIOWrite, "BIO_write");
    }

    if (!handshake_done) {
        err = do_handshake();
    } else {
        while (BIO_ctrl_pending(bio_in) > 0) {
            char dtls_read_buf[8092];
            int nb = SSL_read(dtls, dtls_read_buf, sizeof(dtls_read_buf));

            if (nb > 0 && callback) {
                if ((err = callback->on_dtls_application_data(dtls_read_buf, nb)) != srs_success) {
                    return srs_error_wrap(err, "dtls application data process");
                }
            }
        }
    }

    return err;
}

srs_error_t SrsDtls::start_active_handshake()
{
    if (role_ == SrsDtlsRoleClient) {
        return do_handshake();
    }

    return srs_success;
}

const int SRTP_MASTER_KEY_KEY_LEN = 16;
const int SRTP_MASTER_KEY_SALT_LEN = 14;
srs_error_t SrsDtls::get_srtp_key(std::string& recv_key, std::string& send_key)
{
    srs_error_t err = srs_success;

    unsigned char material[SRTP_MASTER_KEY_LEN * 2] = {0};  // client(SRTP_MASTER_KEY_KEY_LEN + SRTP_MASTER_KEY_SALT_LEN) + server
    static const string dtls_srtp_lable = "EXTRACTOR-dtls_srtp";
    if (!SSL_export_keying_material(dtls, material, sizeof(material), dtls_srtp_lable.c_str(), dtls_srtp_lable.size(), NULL, 0, 0)) {
        return srs_error_new(ERROR_RTC_SRTP_INIT, "SSL_export_keying_material failed");
    }

    size_t offset = 0;

    std::string client_master_key(reinterpret_cast<char*>(material), SRTP_MASTER_KEY_KEY_LEN);
    offset += SRTP_MASTER_KEY_KEY_LEN;
    std::string server_master_key(reinterpret_cast<char*>(material + offset), SRTP_MASTER_KEY_KEY_LEN);
    offset += SRTP_MASTER_KEY_KEY_LEN;
    std::string client_master_salt(reinterpret_cast<char*>(material + offset), SRTP_MASTER_KEY_SALT_LEN);
    offset += SRTP_MASTER_KEY_SALT_LEN;
    std::string server_master_salt(reinterpret_cast<char*>(material + offset), SRTP_MASTER_KEY_SALT_LEN);

    if (role_ == SrsDtlsRoleClient) {
        recv_key = server_master_key + server_master_salt;
        send_key = client_master_key + client_master_salt;
    } else {
        recv_key = client_master_key + client_master_salt;
        send_key = server_master_key + server_master_salt;
    }

    return err;
}

SrsSRTP::SrsSRTP()
{
    recv_ctx_ = NULL;
    send_ctx_ = NULL;
}

SrsSRTP::~SrsSRTP()
{
    if (recv_ctx_) {
        srtp_dealloc(recv_ctx_);
    }

    if (send_ctx_) {
        srtp_dealloc(send_ctx_);
    }
}

srs_error_t SrsSRTP::initialize(string recv_key, std::string send_key)
{
    srs_error_t err = srs_success;

    srtp_policy_t policy;
    bzero(&policy, sizeof(policy));

    // TODO: Maybe we can use SRTP-GCM in future.
    // @see https://bugs.chromium.org/p/chromium/issues/detail?id=713701
    // @see https://groups.google.com/forum/#!topic/discuss-webrtc/PvCbWSetVAQ
    srtp_crypto_policy_set_aes_cm_128_hmac_sha1_80(&policy.rtp);
    srtp_crypto_policy_set_aes_cm_128_hmac_sha1_80(&policy.rtcp);

    policy.ssrc.value = 0;
    // TODO: adjust window_size
    policy.window_size = 8192;
    policy.allow_repeat_tx = 1;
    policy.next = NULL;

    // init recv context
    policy.ssrc.type = ssrc_any_inbound;
    uint8_t *rkey = new uint8_t[recv_key.size()];
    SrsAutoFreeA(uint8_t, rkey);
    memcpy(rkey, recv_key.data(), recv_key.size());
    policy.key = rkey;

    if (srtp_create(&recv_ctx_, &policy) != srtp_err_status_ok) {
        return srs_error_new(ERROR_RTC_SRTP_INIT, "srtp_create recv failed");
    }

    policy.ssrc.type = ssrc_any_outbound;
    uint8_t *skey = new uint8_t[send_key.size()];
    SrsAutoFreeA(uint8_t, skey);
    memcpy(skey, send_key.data(), send_key.size());
    policy.key = skey;

    if (srtp_create(&send_ctx_, &policy) != srtp_err_status_ok) {
        return srs_error_new(ERROR_RTC_SRTP_INIT, "srtp_create recv failed");
    }

    return err;
}

srs_error_t SrsSRTP::protect_rtp(const char* plaintext, char* cipher, int& nb_cipher)
{
    srs_error_t err = srs_success;

    memcpy(cipher, plaintext, nb_cipher);
    // TODO: FIXME: Wrap error code.
    if (srtp_protect(send_ctx_, cipher, &nb_cipher) != 0) {
        return srs_error_new(ERROR_RTC_SRTP_PROTECT, "rtp protect failed");
    }

    return err;
}

srs_error_t SrsSRTP::protect_rtcp(const char* plaintext, char* cipher, int& nb_cipher)
{
    srs_error_t err = srs_success;

    memcpy(cipher, plaintext, nb_cipher);
    // TODO: FIXME: Wrap error code.
    if (srtp_protect_rtcp(send_ctx_, cipher, &nb_cipher) != 0) {
        return srs_error_new(ERROR_RTC_SRTP_PROTECT, "rtcp protect failed");
    }

    return err;
}

srs_error_t SrsSRTP::protect_rtp2(void* rtp_hdr, int* len_ptr)
{
    srs_error_t err = srs_success;

    // TODO: FIXME: Wrap error code.
    if (srtp_protect(send_ctx_, rtp_hdr, len_ptr) != 0) {
        return srs_error_new(ERROR_RTC_SRTP_PROTECT, "rtp protect");
    }

    return err;
}

srs_error_t SrsSRTP::unprotect_rtp(const char* cipher, char* plaintext, int& nb_plaintext)
{
    srs_error_t err = srs_success;

    memcpy(plaintext, cipher, nb_plaintext);
    srtp_err_status_t r0 = srtp_unprotect(recv_ctx_, plaintext, &nb_plaintext);
    if (r0 != srtp_err_status_ok) {
        return srs_error_new(ERROR_RTC_SRTP_UNPROTECT, "unprotect r0=%u", r0);
    }

    return err;
}

srs_error_t SrsSRTP::unprotect_rtcp(const char* cipher, char* plaintext, int& nb_plaintext)
{
    srs_error_t err = srs_success;

    memcpy(plaintext, cipher, nb_plaintext);
    // TODO: FIXME: Wrap error code.
    if (srtp_unprotect_rtcp(recv_ctx_, plaintext, &nb_plaintext) != srtp_err_status_ok) {
        return srs_error_new(ERROR_RTC_SRTP_UNPROTECT, "rtcp unprotect failed");
    }

    return err;
}