// diffie_hellman.hpp
#pragma once
#include <openssl/dh.h>
#include <memory>
#include <vector>

namespace crypto {
	class diffie_hellman {
	public:
		diffie_hellman();
		~diffie_hellman() = default;

		bool create_p_g();
		bool set_p_g(const std::vector<std::uint8_t>& p, const std::vector<std::uint8_t>& g);

		std::vector<std::uint8_t> get_shared_p() const;
		std::vector<std::uint8_t> get_shared_g() const;
		std::vector<std::uint8_t> get_pub_key() const;
		std::vector<std::uint8_t> get_pri_key(const std::vector<std::uint8_t>& pub_key) const;

	private:
		std::unique_ptr<DH, decltype(&DH_free)> dh_;
	};
}

// diffie_hellman.cpp
#include "diffie_hellman.hpp"
#include <openssl/bn.h>

namespace crypto {
	namespace {
		std::vector<std::uint8_t> bignum_to_bytes(const BIGNUM* bn) {
			std::vector<std::uint8_t> result(BN_num_bytes(bn));
			BN_bn2bin(bn, result.data());

			return result;
		}

		std::unique_ptr<BIGNUM, decltype(&BN_free)> bytes_to_big_num(const std::vector<std::uint8_t>& uint8_ts) {
			return std::unique_ptr<BIGNUM, decltype(&BN_free)>{ BN_bin2bn(uint8_ts.data(), uint8_ts.size(), nullptr), &BN_free };
		}
	}


	diffie_hellman::diffie_hellman() : dh_ { DH_new(), & DH_free } {
	}

	bool diffie_hellman::create_p_g() {
		constexpr int key_bits{ 512 };

		if (DH_generate_parameters_ex(dh_.get(), key_bits, DH_GENERATOR_2, nullptr) != 1)  {
			return false;
		}
		return DH_generate_key(dh_.get()) == 1;
	}

	bool diffie_hellman::set_p_g(const std::vector<std::uint8_t>& p, const std::vector<std::uint8_t>& g) {
		auto p_bn = BN_bin2bn(p.data(), p.size(), nullptr);
		auto g_bn = BN_bin2bn(g.data(), g.size(), nullptr);

		if (DH_set0_pqg(dh_.get(), p_bn, nullptr, g_bn) != 1) {
			return false;
		}
		return DH_generate_key(dh_.get()) == 1;
	}

	std::vector<std::uint8_t> diffie_hellman::get_shared_p() const {
		return bignum_to_bytes(DH_get0_p(dh_.get()));
	}

	std::vector<std::uint8_t> diffie_hellman::get_shared_g() const {
		return bignum_to_bytes(DH_get0_g(dh_.get()));
	}

	std::vector<std::uint8_t> diffie_hellman::get_pub_key() const {
		return bignum_to_bytes(DH_get0_pub_key(dh_.get()));
	}


	std::vector<std::uint8_t> diffie_hellman::get_pri_key(const std::vector<std::uint8_t>& pub_key) const {
		auto pub_key_bn{ bytes_to_big_num(pub_key) };

		std::vector<std::uint8_t> result(DH_size(dh_.get()));
		const int ret = DH_compute_key(result.data(), pub_key_bn.get(), dh_.get());

		return ret < 0 ? std::vector<std::uint8_t>{} : result;
	}
}

// main.cpp
#include "diffie_hellman.hpp"
#include <cassert>

int main()
{
	crypto::diffie_hellman request;
	assert(request.create_p_g());
	auto p = request.get_shared_p();
	auto g = request.get_shared_g();
	auto&& pubkey_of_request = request.get_pub_key();

	crypto::diffie_hellman response;
	assert(response.set_p_g(p, g));
	auto&& pubkey_of_response = response.get_pub_key();

	auto&& key1 = request.get_pri_key(pubkey_of_response);
	auto&& key2 = response.get_pri_key(pubkey_of_request);

	assert(key1 == key2);

	return 0;
}