#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>

#define BUF_SIZE 500

static int local_bind(int sock, int family, const char *host,
		      unsigned short port)
{
	int s;

	if (family == AF_INET) {
		struct sockaddr_in sin = {
			.sin_family = AF_INET,
			.sin_port = htons(port),
		};

		s = inet_pton(AF_INET, host, &sin.sin_addr);
		if (s <= 0) {
			if (s == 0)
				fprintf(stderr,
					"%s: not in valid ipv4 format\n", host);
			else
				perror("inet_pton");
			exit(EXIT_FAILURE);
		}

		return bind(sock, (struct sockaddr *)&sin, sizeof(sin));
	} else if (family == AF_INET6) {
		struct sockaddr_in6 sin6 = {
			.sin6_family = AF_INET6,
			.sin6_port = htons(port),
		};

		s = inet_pton(AF_INET6, host, &sin6.sin6_addr);
		if (s <= 0) {
			if (s == 0)
				fprintf(stderr,
					"%s: not in valid ipv6 format\n", host);
			else
				perror("inet_pton");
			exit(EXIT_FAILURE);
		}

		return bind(sock, (struct sockaddr *)&sin6, sizeof(sin6));
	} else {
		fprintf(stderr, "Unknown family: %d\n", family);
	}
	return 0;
}

static void md5set(int sock, struct sockaddr *sa, socklen_t salen,
		   const char *password)
{
	struct tcp_md5sig md5sig;

	memset(&md5sig, 0, sizeof(md5sig));
	memcpy(&md5sig.tcpm_addr, sa, salen);
	memcpy(md5sig.tcpm_key, password, md5sig.tcpm_keylen =
	       strlen(password));

	if (setsockopt(sock, IPPROTO_TCP, TCP_MD5SIG, &md5sig, sizeof md5sig) <
	    0)
		perror("setsockopt (TCP_MD5SIG)");
}

static void usage(void)
{
	fprintf(stderr,
		"Usage: client [-6|-4] host port [lhost lport passwd]\n");
	exit(EXIT_FAILURE);
}


int main(int argc, char *argv[])
{
	struct addrinfo *result, *rp;
	int sfd, s;
	ssize_t nread;
	char buf[BUF_SIZE];
	struct addrinfo hints = {
		.ai_socktype = SOCK_STREAM,
	};

	if (argc > 1 && argv[1][0] == '-') {
		if (strcmp(argv[1], "-4") == 0) {
			hints.ai_family = AF_INET;
		} else if (strcmp(argv[1], "-6") == 0) {
			hints.ai_family = AF_INET6;
		} else
			usage();
		--argc, ++argv;
	}

	if (argc < 3)
		usage();

	/* Obtain address(es) matching host/port */
	s = getaddrinfo(argv[1], argv[2], &hints, &result);
	if (s != 0) {
		fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(s));
		exit(EXIT_FAILURE);
	}

	/* getaddrinfo() returns a list of address structures.
	   Try each address until we successfully connect(2).
	   If socket(2) (or connect(2)) fails, we (close the socket
	   and) try the next address. */

	for (rp = result; rp != NULL; rp = rp->ai_next) {
		sfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
		if (sfd == -1) {
			perror("socket");
			continue;
		}

		if (argc > 3) {
			local_bind(sfd, rp->ai_family, argv[3], atoi(argv[4]));
			md5set(sfd, rp->ai_addr, rp->ai_addrlen, argv[5]);
		}

		printf("connecting to %s..",
		       inet_ntop(rp->ai_family, rp->ai_addr, buf, sizeof(buf)));
		fflush(stdout);

		if (connect(sfd, rp->ai_addr, rp->ai_addrlen) != -1)
			break;	/* Success */

		perror("connect");
		close(sfd);
	}

	if (rp == NULL) {	/* No address succeeded */
		fprintf(stderr, "Could not connect\n");
		exit(EXIT_FAILURE);
	}

	freeaddrinfo(result);	/* No longer needed */

	printf("connected\n");
	while ((nread = read(0, buf, sizeof(buf))) > 0) {
		if (write(sfd, buf, nread) < 0)
			break;
		nread = read(sfd, buf, sizeof(buf));
		if (nread <= 0)
			break;
		if (write(1, buf, nread) != nread)
			break;
	}

	exit(EXIT_SUCCESS);
}