/*
 * s_client command
 *
 * Copyright (C) 2019  Patrick McDermott
 *
 * This file is part of wolfssl-util.
 *
 * wolfssl-util is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * wolfssl-util is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with wolfssl-util.  If not, see <http://www.gnu.org/licenses/>.
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <errno.h>
#include <netdb.h>
#include <poll.h>
#include <stdbool.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#include <wolfssl/options.h>
#include <wolfssl/ssl.h>
#include <wolfssl/wolfcrypt/settings.h>

#include "commands.h"

#define ARRAY_SIZE(a)  (sizeof(a) / sizeof((a)[0]))
#undef MAX
#define MAX(a, b)      (((a) > (b)) ? (a) : (b))
#undef MIN
#define MIN(a, b)      (((a) < (b)) ? (a) : (b))

static _Bool
parse_host_port(char *hostport, char **host, char **port)
{
	*host = hostport;
	/* XXX: Port is required.  Otherwise, this will mangle IPv6 addresses
	 * without some more intelligent (and larger) code. */
	*port = strrchr(*host, ':');
	if (!*port) {
		fputs("Port is required\n", stderr);
		return false;
	}
	**port = '\0';
	++*port;

	return true;
}

#ifdef HAVE_SUPPORTED_CURVES
static _Bool
use_curves(WOLFSSL_CTX *ctx)
{
	static word16 curves[] = {
#ifdef HAVE_CURVE25519
		WOLFSSL_ECC_X25519,
#endif
#ifdef HAVE_ECC
#  if defined(HAVE_ECC160) || defined(HAVE_ALL_CURVES)
#    ifdef HAVE_ECC_KOBLITZ
		WOLFSSL_ECC_SECP160K1,
#    endif
#    ifndef NO_ECC_SECP  /* Ugh double negative */
		WOLFSSL_ECC_SECP160R1,
#    endif
#    ifdef HAVE_ECC_SECPR2
		WOLFSSL_ECC_SECP160R2,
#    endif
#  endif
#  if defined(HAVE_ECC192) || defined(HAVE_ALL_CURVES)
#    ifdef HAVE_ECC_KOBLITZ
		WOLFSSL_ECC_SECP192K1,
#    endif
#    ifndef NO_ECC_SECP
		WOLFSSL_ECC_SECP192R1,
#    endif
#  endif
#  if defined(HAVE_ECC224) || defined(HAVE_ALL_CURVES)
#    ifdef HAVE_ECC_KOBLITZ
		WOLFSSL_ECC_SECP224K1,
#    endif
#    ifndef NO_ECC_SECP
		WOLFSSL_ECC_SECP224R1,
#    endif
#  endif
#  if !defined(NO_ECC256) || defined(HAVE_ALL_CURVES)
#    ifdef HAVE_ECC_KOBLITZ
		WOLFSSL_ECC_SECP256K1,
#    endif
#    ifndef NO_ECC_SECP
		WOLFSSL_ECC_SECP256R1,
#    endif
#    ifdef HAVE_ECC_BRAINPOOL
		WOLFSSL_ECC_BRAINPOOLP256R1,
#    endif
#  endif
#  if defined(HAVE_ECC384) || defined(HAVE_ALL_CURVES)
#    ifndef NO_ECC_SECP
		WOLFSSL_ECC_SECP384R1,
#    endif
#    ifdef HAVE_ECC_BRAINPOOL
		WOLFSSL_ECC_BRAINPOOLP384R1,
#    endif
#  endif
#  if defined(HAVE_ECC521) || defined(HAVE_ALL_CURVES)
#    ifndef NO_ECC_SECP
		WOLFSSL_ECC_SECP521R1,
#    endif
#    ifdef HAVE_ECC_BRAINPOOL
		WOLFSSL_ECC_BRAINPOOLP512R1,
#    endif
#  endif
#endif  /* HAVE_ECC */
#ifdef HAVE_FFDHE_2048
		WOLFSSL_FFDHE_2048,
#endif
#ifdef HAVE_FFDHE_3072
		WOLFSSL_FFDHE_3072,
#endif
#ifdef HAVE_FFDHE_4096
		WOLFSSL_FFDHE_4096,
#endif
#ifdef HAVE_FFDHE_6144
		WOLFSSL_FFDHE_6144,
#endif
#ifdef HAVE_FFDHE_8192
		WOLFSSL_FFDHE_8192,
#endif
	};
	size_t i;

	for (i = 0; i < ARRAY_SIZE(curves); ++i) {
		if (wolfSSL_CTX_UseSupportedCurve(ctx, curves[i]) !=
				WOLFSSL_SUCCESS) {
			return false;
		}
	}

	return true;
}
#endif  /* HAVE_SUPPORTED_CURVES */

static int
connect_socket(const char *host, const char *port)
{
	struct addrinfo  hints = {0};
	struct addrinfo *result;
	int              s;
	struct addrinfo *rp;
	int              sfd;

	hints.ai_family   = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_flags    = AI_PASSIVE;
	hints.ai_protocol = 0;

	s = getaddrinfo(host, port, &hints, &result);
	if (s != 0) {
		fprintf(stderr, "Failed to resolve host and port: %s\n",
				gai_strerror(s));
		return -1;
	}

	for (rp = result; rp != NULL; rp = rp->ai_next) {
		sfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
		if (sfd == -1) {
			continue;
		}
		if (connect(sfd, rp->ai_addr, rp->ai_addrlen) == -1) {
			close(sfd);
			continue;
		}
		break;
	}
	if (rp == NULL) {
		fputs("Failed to connect\n", stderr);
		sfd = -1;
	}

	freeaddrinfo(result);

	return sfd;
}

static _Bool
write_all(int fd, const void *buf, size_t count)
{
	ssize_t ret;

	while (count > 0) {
		while ((ret = write(fd, buf, count)) < 0 && errno == EINTR) {
			continue;
		}
		if (ret < 0) {
			return false;
		}
		buf = ((const char *) buf) + ret;
		count -= ret;
	}

	return true;
}

static _Bool
poll_fds(int sfd, WOLFSSL *ssl)
{
	struct pollfd fds[2] = {
		{ .fd = -1, .events = POLLIN|POLLERR, .revents = 0 },
		{ .fd = -1, .events = POLLIN|POLLERR, .revents = 0 },
	};
	char          buf[MAX(8192, WOLFSSL_MAX_ERROR_SZ)];
	ssize_t       len;
	int           ret;

	fds[0].fd = STDIN_FILENO;
	fds[1].fd = sfd;

	for (;;) {
		while (poll(fds, ARRAY_SIZE(fds), -1) < 0 && (errno == EINTR ||
					errno == EAGAIN)) {
			continue;
		}
		if (fds[0].revents > 0) {  /* stdin */
			len = read(STDIN_FILENO, buf, sizeof(buf));
			if (len < 0) {
				fputs("Input read error\n", stderr);
				return false;
			} else if (len == 0) {
				fds[0].fd = -1;  /* Stop polling. */
			} else if ((ret = wolfSSL_write(ssl, buf, len)) <= 0) {
				wolfSSL_ERR_error_string(wolfSSL_get_error(ssl,
							ret), buf);
				fprintf(stderr, "Socket write error: %s\n",
						buf);
				return false;
			}
		}
		if (fds[1].revents > 0) {  /* socket */
			ret = wolfSSL_read(ssl, buf, MIN(sizeof(buf), 1024));
			if (ret < 0) {
				wolfSSL_ERR_error_string(wolfSSL_get_error(ssl,
							ret), buf);
				fprintf(stderr, "Socket read error: %s\n",
						buf);
				return false;
			} else if (ret == 0) {
				fds[1].fd = -1;  /* Stop polling. */
				close(STDOUT_FILENO);  /* Signal socket EOF. */
			} else if (write_all(STDOUT_FILENO, buf, (size_t) ret)
					== false) {
				return false;
			}
		}
		if (fds[0].fd == fds[1].fd) {  /* Both -1 (no longer polled) */
			return true;
		}
	}

	/* Unreached */
}

int
s_client(int argc, char **argv)
{
	char           *host       = NULL;
	char           *port       = NULL;
	const char     *servername = NULL;
	int             ret        = EXIT_FAILURE;
	WOLFSSL_METHOD *method;
	WOLFSSL_CTX    *ctx        = NULL;
	WOLFSSL        *ssl        = NULL;
	int             sfd        = -1;
	int             err;
	char            buf[WOLFSSL_MAX_ERROR_SZ];

	for (; argc > 0; --argc, ++argv) {
		if (strcmp(*argv, "-quiet") == 0) {
			/* No-op */
		} else if (strcmp(*argv, "-connect") == 0) {
			--argc, ++argv;
			if (parse_host_port(*argv, &host, &port) == false) {
				return EXIT_FAILURE;
			}
		} else if (strcmp(*argv, "-servername") == 0) {
			--argc, ++argv;
			servername = *argv;
		} else if (**argv == '-') {
			fprintf(stderr, "Unsupported option \"%s\"\n", *argv);
		}
	}

	wolfSSL_Init();

	if ((method = wolfTLSv1_2_client_method()) == NULL ||
			(ctx = wolfSSL_CTX_new(method)) == NULL) {
		fputs("Out of memory\n", stderr);
		goto error;
	}

#if defined(HAVE_CA_CERTS) && HAVE_CA_CERTS
	if (wolfSSL_CTX_load_verify_locations_ex(ctx, NULL, CA_CERTS,
				WOLFSSL_LOAD_FLAG_IGNORE_ERR) !=
			WOLFSSL_SUCCESS) {
		fputs("Failed to load CA certificates\n", stderr);
		goto error;
	}
#endif

	if (
#ifdef HAVE_CERTIFICATE_STATUS_REQUEST_V2
			wolfSSL_CTX_EnableOCSPStapling(ctx) != WOLFSSL_SUCCESS
			|| wolfSSL_CTX_UseOCSPStaplingV2(ctx,
				WOLFSSL_CSR2_OCSP_MULTI, 0) !=
				WOLFSSL_SUCCESS ||
#endif
#ifdef HAVE_CERTIFICATE_STATUS_REQUEST
			wolfSSL_CTX_EnableOCSPStapling(ctx) != WOLFSSL_SUCCESS
			|| wolfSSL_CTX_UseOCSPStapling(ctx, WOLFSSL_CSR_OCSP, 0)
				!= WOLFSSL_SUCCESS ||
#endif
#ifdef HAVE_OCSP
			wolfSSL_CTX_EnableOCSP(ctx, WOLFSSL_OCSP_CHECKALL) !=
				WOLFSSL_SUCCESS ||
#endif
#if defined(HAVE_CRL) && defined(HAVE_CRL_IO)
			wolfSSL_CTX_EnableCRL(ctx, WOLFSSL_CRL_CHECKALL) !=
				WOLFSSL_SUCCESS ||
#endif
#ifdef HAVE_SNI
			(servername != NULL && wolfSSL_CTX_UseSNI(ctx,
				WOLFSSL_SNI_HOST_NAME, servername,
				strlen(servername)) != WOLFSSL_SUCCESS) ||
#endif
#ifdef HAVE_SUPPORTED_CURVES
			use_curves(ctx) == false ||
#endif
			(ssl = wolfSSL_new(ctx)) == NULL ||
			wolfSSL_check_domain_name(ssl, servername) !=
				WOLFSSL_SUCCESS
			) {
		fputs("Out of memory\n", stderr);
		goto error;
	}

	sfd = connect_socket(host, port);
	if (sfd == -1) {
		goto error;
	}
	wolfSSL_set_fd(ssl, sfd);

	if ((err = wolfSSL_connect(ssl)) != WOLFSSL_SUCCESS) {
		err = wolfSSL_get_error(ssl, err);
		wolfSSL_ERR_error_string(err, buf);
		fprintf(stderr, "Handshake error: %s\n", buf);
		goto error;
	}

	if (poll_fds(sfd, ssl) == false) {
		goto error;
	}

	ret = EXIT_SUCCESS;
error:
	close(sfd);
	wolfSSL_free(ssl);
	wolfSSL_CTX_free(ctx);
	wolfSSL_Cleanup();

	return ret;
}