/*
 * librdkafka - Apache Kafka C library
 *
 * Copyright (c) 2019, Magnus Edenhill
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

 /**
  * Example of utilizing the Windows Certificate store with SSL.
  */

#include <iostream>
#include <string>
#include <cstdlib>
#include <cstdio>
#include <csignal>
#include <cstring>
#include <sstream>

#include "../win32/wingetopt.h"
#include <windows.h>
#include <wincrypt.h>

  /*
   * Typically include path in a real application would be
   * #include <librdkafka/rdkafkacpp.h>
   */
#include "rdkafkacpp.h"



class ExampleStoreRetriever {
public:
        ExampleStoreRetriever (std::string const &subject, std::string const &pass)
                : m_cert_subject(subject), m_password(pass),
                m_cert_store(NULL), m_cert_ctx(NULL) {
                load_certificate();
        }

        ~ExampleStoreRetriever() {
                if (m_cert_ctx)
                        CertFreeCertificateContext(m_cert_ctx);

                if (m_cert_store)
                        CertCloseStore(m_cert_store, 0);
        }

        /* @returns the public key in DER format */
        const std::vector<unsigned char> get_public_key () {
                std::vector<unsigned char> buf((size_t)m_cert_ctx->cbCertEncoded);
                buf.assign((const char *)m_cert_ctx->pbCertEncoded,
                        (const char *)m_cert_ctx->pbCertEncoded +
                           (size_t)m_cert_ctx->cbCertEncoded);
                return buf;
        }

        /* @returns the private key in PCKS#12 format */
        const std::vector<unsigned char> get_private_key () {
                ssize_t ret = 0;
                /*
                 * In order to export the private key the certificate
                 * must first be marked as exportable.
                 *
                 * Steps to export the certificate
                 * 1) Create an in-memory cert store
                 * 2) Add the certificate to the store
                 * 3) Export the private key from the in-memory store
                 */

                 /* Create an in-memory cert store */
                HCERTSTORE hMemStore = CertOpenStore(CERT_STORE_PROV_MEMORY,
                                                     0, NULL, 0, NULL);
                if (!hMemStore)
                        throw "Failed to create in-memory cert store: " +
                                GetErrorMsg(GetLastError());

                /* Add certificate to store */
                if (!CertAddCertificateContextToStore(hMemStore,
                                                      m_cert_ctx,
                                                      CERT_STORE_ADD_USE_EXISTING,
                                                      NULL))
                        throw "Failed to add certificate to store: " +
                                GetErrorMsg(GetLastError());

                /*
                 * Export private key from cert
                 */
                CRYPT_DATA_BLOB db = { NULL };

                std::wstring w_password(m_password.begin(), m_password.end());

                /* Acquire output size */
                if (!PFXExportCertStoreEx(hMemStore,
                                          &db,
                                          w_password.c_str(),
                                          NULL,
                                          EXPORT_PRIVATE_KEYS |
                                          REPORT_NO_PRIVATE_KEY |
                                          REPORT_NOT_ABLE_TO_EXPORT_PRIVATE_KEY))
                        throw "Failed to export private key: " + GetErrorMsg(GetLastError());

                std::vector<unsigned char> buf;

                buf.resize(db.cbData);
                db.pbData = &buf[0];

                /* Extract key */
                if (!PFXExportCertStoreEx(hMemStore,
                                          &db,
                                          w_password.c_str(),
                                          NULL,
                                          EXPORT_PRIVATE_KEYS |
                                          REPORT_NO_PRIVATE_KEY |
                                          REPORT_NOT_ABLE_TO_EXPORT_PRIVATE_KEY))
                        throw "Failed to export private key (PFX): " + GetErrorMsg(GetLastError());

                CertCloseStore(hMemStore, 0);

                buf.resize(db.cbData);

                return buf;
        }

 private:
         void load_certificate () {
                 if (m_cert_ctx)
                         return;

                 m_cert_store = CertOpenStore(CERT_STORE_PROV_SYSTEM,
                                              0,
                                              NULL,
                                              CERT_SYSTEM_STORE_CURRENT_USER,
                                              L"My");
                 if (!m_cert_store)
                         throw "Failed to open cert store: " + GetErrorMsg(GetLastError());

                 m_cert_ctx = CertFindCertificateInStore(m_cert_store,
                                                         X509_ASN_ENCODING,
                                                         0,
                                                         CERT_FIND_SUBJECT_STR,
                                                         /* should probally do a better std::string to std::wstring conversion */
                                                         std::wstring(m_cert_subject.begin(),
                                                                      m_cert_subject.end()).c_str(),
                                                         NULL);
                 if (!m_cert_ctx) {
                         CertCloseStore(m_cert_store, 0);
                         m_cert_store = NULL;
                         throw "Certificate " + m_cert_subject + " not found in cert store: " + GetErrorMsg(GetLastError());
                 }
         }

         std::string GetErrorMsg (unsigned long error) {
                 char *message = NULL;
                 size_t ret = FormatMessageA(
                         FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
                         nullptr,
                         error,
                         0,
                         (char*)&message,
                         0,
                         nullptr);
                 if (ret == 0) {
                         std::stringstream ss;

                         ss << std::string("could not format message for ") << error;
                         return ss.str();
                 } else {
                         std::string result(message, ret);
                         LocalFree(message);
                         return result;
                 }
         }

 private:
         std::string m_cert_subject;
         std::string m_password;
         PCCERT_CONTEXT m_cert_ctx;
         HCERTSTORE m_cert_store;
};


class PrintingSSLVerifyCb : public RdKafka::SslCertificateVerifyCb {
        /* This SSL cert verification callback simply prints the certificates
         * in the certificate chain.
         * It provides no validation, everything is ok. */
public:
        bool ssl_cert_verify_cb (const std::string &broker_name,
                                 int32_t broker_id,
                                 int *x509_error,
                                 int depth,
                                 const char *buf, size_t size,
                                 std::string &errstr) {
                PCCERT_CONTEXT ctx = CertCreateCertificateContext(
                        X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
                        (const uint8_t*)buf, static_cast<unsigned long>(size));

                if (!ctx)
                        std::cerr << "Failed to parse certificate" << std::endl;

                char subject[256] = "n/a";
                char issuer[256] = "n/a";

                CertGetNameStringA(ctx, CERT_NAME_FRIENDLY_DISPLAY_TYPE,
                                   0, NULL,
                                   subject, sizeof(subject));

                CertGetNameStringA(ctx, CERT_NAME_FRIENDLY_DISPLAY_TYPE,
                                   CERT_NAME_ISSUER_FLAG, NULL,
                                   issuer, sizeof(issuer));

                std::cerr << "Broker " << broker_name <<
                        " (" << broker_id << "): " <<
                        "certificate depth " << depth <<
                        ", X509 error " << *x509_error <<
                        ", subject " << subject <<
                        ", issuer " << issuer << std::endl;

                if (ctx)
                        CertFreeCertificateContext(ctx);

                return true;
        }
};


/**
* @brief Print the brokers in the cluster.
*/
static void print_brokers (RdKafka::Handle *handle,
                           const RdKafka::Metadata *md) {
        std::cout << md->brokers()->size() << " broker(s) in cluster " <<
                handle->clusterid(0) << std::endl;

        /* Iterate brokers */
        RdKafka::Metadata::BrokerMetadataIterator ib;
        for (ib = md->brokers()->begin(); ib != md->brokers()->end(); ++ib)
                std::cout << "  broker " << (*ib)->id() << " at "
                << (*ib)->host() << ":" << (*ib)->port() << std::endl;

}


int main (int argc, char **argv) {
        std::string brokers;
        std::string errstr;
        std::string cert_subject;
        std::string priv_key_pass;

        /*
         * Create configuration objects
         */
        RdKafka::Conf *conf = RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL);
        RdKafka::Conf *tconf = RdKafka::Conf::create(RdKafka::Conf::CONF_TOPIC);

        int opt;
        while ((opt = getopt(argc, argv, "b:d:X:s:p:")) != -1) {
                switch (opt) {
                        case 'b':
                                brokers = optarg;
                                break;
                        case 'd':
                                if (conf->set("debug", optarg, errstr) != RdKafka::Conf::CONF_OK) {
                                        std::cerr << errstr << std::endl;
                                        exit(1);
                                }
                                break;
                        case 'X':
                        {
                                char *name, *val;

                                name = optarg;
                                if (!(val = strchr(name, '='))) {
                                        std::cerr << "%% Expected -X property=value, not " <<
                                                name << std::endl;
                                        exit(1);
                                }

                                *val = '\0';
                                val++;

                                if (conf->set(name, val, errstr) != RdKafka::Conf::CONF_OK) {
                                        std::cerr << errstr << std::endl;
                                        exit(1);
                                }
                        }
                        break;

                        case 's':
                                cert_subject = optarg;
                                break;

                        case 'p':
                                priv_key_pass = optarg;
                                if (conf->set("ssl.key.password", optarg, errstr) !=
                                    RdKafka::Conf::CONF_OK) {
                                        std::cerr << errstr << std::endl;
                                        exit(1);
                                }

                                break;

                        default:
                                goto usage;
                }
        }

        if (brokers.empty() || optind != argc) {
usage:
                std::string features;
                conf->get("builtin.features", features);
                fprintf(stderr,
                        "Usage: %s [options] -b <brokers> -s <cert-subject> -p <priv-key-password>\n"
                        "\n"
                        "Windows Certificate Store integration example.\n"
                        "Use certlm.msc or mmc to view your certificates.\n"
                        "\n"
                        "librdkafka version %s (0x%08x, builtin.features \"%s\")\n"
                        "\n"
                        " Options:\n"
                        "  -b <brokers>    Broker address\n"
                        "  -s <cert>       The subject name of the client's SSL certificate to use\n"
                        "  -p <pass>       The private key password\n"
                        "  -d [facs..]     Enable debugging contexts: %s\n"
                        "  -X <prop=name>  Set arbitrary librdkafka "
                        "configuration property\n"
                        "\n",
                        argv[0],
                        RdKafka::version_str().c_str(), RdKafka::version(),
                        features.c_str(),
                        RdKafka::get_debug_contexts().c_str());
                exit(1);
        }

        if (!cert_subject.empty()) {

                try {
                        /* Load certificates from the Windows store */
                        ExampleStoreRetriever certStore(cert_subject, priv_key_pass);

                        std::vector<unsigned char> pubkey, privkey;

                        pubkey = certStore.get_public_key();
                        privkey = certStore.get_private_key();

                        if (conf->set_ssl_cert(RdKafka::CERT_PUBLIC_KEY,
                                               RdKafka::CERT_ENC_DER,
                                               &pubkey[0], pubkey.size(),
                                               errstr) !=
                            RdKafka::Conf::CONF_OK)
                                throw "Failed to set public key: " + errstr;

                        if (conf->set_ssl_cert(RdKafka::CERT_PRIVATE_KEY,
                                               RdKafka::CERT_ENC_PKCS12,
                                               &privkey[0], privkey.size(),
                                               errstr) !=
                            RdKafka::Conf::CONF_OK)
                                throw "Failed to set private key: " + errstr;

                } catch (const std::string &ex) {
                        std::cerr << ex << std::endl;
                        exit(1);
                }
        }


        /*
         * Set configuration properties
         */
        conf->set("bootstrap.servers", brokers, errstr);

        /* We use the Certificiate verification callback to print the
         * certificate chains being used. */
        PrintingSSLVerifyCb ssl_verify_cb;

        if (conf->set("ssl_cert_verify_cb", &ssl_verify_cb, errstr) != RdKafka::Conf::CONF_OK) {
                std::cerr << errstr << std::endl;
                exit(1);
        }

        /* Create any type of client, producering being the cheapest. */
        RdKafka::Producer *producer = RdKafka::Producer::create(conf, errstr);
        if (!producer) {
                std::cerr << "Failed to create producer: " << errstr << std::endl;
                exit(1);
        }

        RdKafka::Metadata *metadata;

        /* Fetch metadata */
        RdKafka::ErrorCode err = producer->metadata(false, NULL, &metadata, 5000);
        if (err != RdKafka::ERR_NO_ERROR) {
                std::cerr << "%% Failed to acquire metadata: "
                        << RdKafka::err2str(err) << std::endl;
                exit(1);
        }

        print_brokers(producer, metadata);

        delete metadata;
        delete producer;

        return 0;
}