/*
 *      tcpvsadm - TCP Virtual Server ADMinistration program
 *
 *      Version: $Id$
 *
 *      Authors: Wensong Zhang <wensong@linuxvirtualserver.org>
 *
 *      Note that most code is taken from ipvsadm.c.
 *
 *      This program 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.
 *
 *      This program 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 this program; if not, write to the Free Software
 *      Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 */

#undef __KERNEL__       /* Makefile lazyness ;) */
#include <stdio.h>
#include <stdlib.h>
#include <getopt.h>
#include <netdb.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <ctype.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/param.h>
#include <arpa/inet.h>

#include <asm/types.h>          /* For __uXX types */
#include <net/if.h>
#include <netinet/ip_icmp.h>
#include <netinet/udp.h>
#include <netinet/tcp.h>

#include "tcp_vs.h"

#define TCPVS_PROC_FILE                  "/proc/net/ktcpvs/config"

/* default scheduler */
#define DEF_SCHED	"wlc"

/* printing format flags */
#define FMT_NONE	0x0000
#define FMT_NUMERIC	0x0001

int string_to_number(const char *s, int min, int max);
int host_to_addr(const char *name, struct in_addr *addr);
char * addr_to_host(struct in_addr *addr);
char * addr_to_anyname(struct in_addr *addr);
int service_to_port(const char *name, unsigned short proto);
char * port_to_service(int port, unsigned short proto);
char * port_to_anyname(int port, unsigned short proto);

int parse_service(char *buf, u_int16_t proto,
		  u_int32_t *addr, u_int16_t *port);
int parse_netmask(char *buf, u_int32_t *addr);
int parse_timeout(char *buf, unsigned *timeout);

void list_tcpvs(unsigned int options);
int write_tcpvs(struct tcp_vs_ctl *ctl);

void usage_exit(char *program, const int exit_status);
void fail(int err, char *text);


static struct option long_options[] =
{
	{"add-service", 0, 0, 'A'},
	{"edit-service", 0, 0, 'E'},
	{"delete-service", 0, 0, 'D'},
	{"clear", 0, 0, 'C'},
	{"list", 0, 0, 'L'},
	{"add-server", 0, 0, 'a'},
	{"edit-server", 0, 0, 'e'},
	{"delete-server", 0, 0, 'd'},
	{"add-rule", 0, 0, '1'},
	{"del-rule", 0, 0, '2'},
	{"help", 0, 0, 'h'},
	{"ident", 1, 0, 'i'},
	{"port", 1, 0, 'P'},
	{"scheduler", 1, 0, 's'},
	{"real-server", 1, 0, 'r'},
	{"weight", 1, 0, 'w'},
	{"pattern", 1, 0, 'p'},
	{"numeric", 0, 0, 'n'},
	{0, 0, 0, 0}
};


int main(int argc, char **argv)
{
	int cmd, c;
	int parse;
	char *optstr = "";
	int result;
	unsigned int options = FMT_NONE;
	struct tcp_vs_ctl ctl;

	/*
	 *   If no other arguement, list TCPVS_PROC_FILE
	 */
	if (argc == 1){
		list_tcpvs(options);
		exit(0);
	}

	memset(&ctl, 0, sizeof(struct tcp_vs_ctl));

	/*
	 *	Want user virtual server control
	 */
	if ((cmd = getopt_long(argc, argv, "AEDCaedLlh",
			       long_options, NULL)) == EOF)
		usage_exit(argv[0], -1);

	switch (cmd) {
	case 'A':
		ctl.cmd = TCP_VS_CMD_ADD;
		optstr = "i:s:P:";
		break;
	case 'E':
		ctl.cmd = TCP_VS_CMD_SET;
		optstr = "i:s:P:";
		break;
	case 'D':
		ctl.cmd = TCP_VS_CMD_DEL;
		optstr = "i:";
		break;
	case 'a':
		ctl.cmd = TCP_VS_CMD_ADD_DEST;
		optstr = "i:w:r:";
		break;
	case 'e':
		ctl.cmd = TCP_VS_CMD_SET_DEST;
		optstr = "i:w:r:";
		break;
	case 'd':
		ctl.cmd = TCP_VS_CMD_DEL_DEST;
		optstr = "i:w:r:";
		break;
	case 'C':
		ctl.cmd = TCP_VS_CMD_FLUSH;
		optstr = "";
		break;
	case '1':
		ctl.cmd = TCP_VS_CMD_ADD_RULE;
		optstr = "i:p:r:";
		break;
	case '2':
		ctl.cmd = TCP_VS_CMD_DEL_RULE;
		optstr = "i:p:r:";
		break;
	case 'L':
	case 'l':
		ctl.cmd = TCP_VS_CMD_LIST;
		optstr = "n";
		break;
	case 'h':
		usage_exit(argv[0], 0);
		break;
	default:
		usage_exit(argv[0], -1);
	}

	/* weight=0 is allowed, which means that server is quiesced */
	ctl.weight = -1;

	/* the default serverport, will implement to configure it later */
	ctl.serverport = 8080;

	while ((c=getopt_long(argc, argv, optstr,
			      long_options, NULL)) != EOF) {
		switch (c) {
		case 'i':
			if (strlen(ctl.name) != 0)
				fail(2, "multiple tcpvs specified");
			strncpy(ctl.name, optarg, KTCPVS_VSNAME_MAXLEN);
			break;
		case 's':
			if (strlen(ctl.sched_name) != 0)
				fail(2, "multiple scheduler specified");
			strncpy(ctl.sched_name,
				optarg, KTCPVS_SCHEDNAME_MAXLEN);
			break;
		case 'P':
			ctl.serverport = string_to_number(optarg, 1, 65534);
			if (ctl.serverport == -1)
				fail(2, "illegal server port specified");
			break;
		case 'r':
			parse = parse_service(optarg, IPPROTO_TCP,
					      &ctl.daddr, &ctl.dport);
			if (parse == 0)
				fail(2, "illegal dest address:port specified");
			/* copy vport to dport if none specified */
			if (parse == 1)
				ctl.dport = ctl.serverport;
			break;
		case 'w':
			if (ctl.weight != -1)
				fail(2, "multiple server weights specified");
			if ((ctl.weight=
			     string_to_number(optarg,0,65535)) == -1)
				fail(2, "illegal weight specified");
			break;
		case 'p':
			if (strlen(ctl.pattern) != 0)
				fail(2, "multiple scheduler specified");
			strncpy(ctl.pattern,
				optarg, KTCPVS_PATTERN_MAXLEN);
			ctl.len = strlen(optarg);
			break;
		case 'n':
			options |= FMT_NUMERIC;
			break;
		default:
			fail(2, "invalid option");
		}
	}

	if (optind < argc)
		fail(2, "unknown arguments found in command line");

	if (ctl.cmd == TCP_VS_CMD_LIST) {
		list_tcpvs(options);
		exit(0);
	}

	if (ctl.cmd == TCP_VS_CMD_ADD
	    || ctl.cmd == TCP_VS_CMD_SET) {
		/*
		 * Set the default scheduling algorithm if not specified
		 */
		if (strlen(ctl.sched_name) == 0)
			strcpy(ctl.sched_name, DEF_SCHED);
	}

	if (ctl.cmd == TCP_VS_CMD_ADD_DEST
	    || ctl.cmd == TCP_VS_CMD_SET_DEST) {
		/*
		 * Set the default weight 1 if not specified
		 */
		if (ctl.weight == -1)
			ctl.weight = 1;
	}

	result = write_tcpvs(&ctl);

	return 0;
}


void list_tcpvs(unsigned int format)
{
	static char buffer[1024];
	FILE *handle;

	handle = fopen(TCPVS_PROC_FILE, "r");
	if (!handle) {
		fprintf(stderr, "Could not open the %s file\n"
			"Are you sure that the KTCPVS module is inserted "
			"in the kernel?\n", TCPVS_PROC_FILE);
		exit(1);
	}

	/*
	 * Print the KTCPVS information
	 */
	while (!feof(handle)) {
		if (fgets(buffer, sizeof(buffer), handle))
			printf("%s", buffer);
	}

	fclose(handle);
}


int write_tcpvs(struct tcp_vs_ctl *ctl)
{
	FILE *handle;
	int ret;

	handle = fopen(TCPVS_PROC_FILE, "w");
	if (!handle) {
		fprintf(stderr, "Could not open the %s file\n"
			"Are you sure that the KTCPVS module is inserted "
			"in the kernel?\n", TCPVS_PROC_FILE);
		exit(1);
	}

	ret = fwrite(ctl, sizeof(*ctl), 1, handle);
	if (ret < 0) {
		/*
		 * We will add some code to parse the error no here later
		 */
		fprintf(stderr, "write the %s file error.\n", TCPVS_PROC_FILE);
		exit (1);
	}

	fclose(handle);
	return ret;
}


int string_to_number(const char *s, int min, int max)
{
	int number;
	char *end;

	number = (int)strtol(s, &end, 10);
	if (*end == '\0' && end != s) {
		/* we parsed a number, let's see if we want this */
		if (min <= number && number <= max)
			return number;
		else
			return -1;
	} else
		return -1;
}


/*
 * Get netmask.
 * Return 0 if failed,
 *	  1 if addr read
 */
int parse_netmask(char *buf, u_int32_t *addr)
{
	struct in_addr inaddr;

	if (inet_aton(buf, &inaddr) != 0)
		*addr = inaddr.s_addr;
	else if (host_to_addr(buf, &inaddr) != -1)
		*addr = inaddr.s_addr;
	else
		return 0;

	return 1;
}

/*
 * Get IP address and port from the argument.
 * Return 0 if failed,
 *	  1 if addr read
 *        2 if addr and port read
 */
int parse_service(char *buf, u_int16_t proto, u_int32_t *addr, u_int16_t *port)
{
	char *pp;
	long prt;
	struct in_addr inaddr;

	pp = strchr(buf,':');
	if (pp) *pp = '\0';

	if (inet_aton(buf, &inaddr) != 0)
		*addr = inaddr.s_addr;
	else if (host_to_addr(buf, &inaddr) != -1)
		*addr = inaddr.s_addr;
	else
		return 0;

	if (pp == NULL)
		return 1;

	if ((prt=string_to_number(pp+1, 0, 65535)) != -1)
		*port = htons(prt);
	else if ((prt=service_to_port(pp+1, proto)) != -1)
		*port = htons(prt);
	else
		return 0;

	return 2;
}


/*
 * Get the timeout of persistent service.
 * Return 0 if failed(the timeout value is less or equal than zero).
 *	  1 if succeed.
 */
int parse_timeout(char *buf, unsigned *timeout)
{
	int i;

	if (buf == NULL) {
		*timeout = TCP_VS_TEMPLATE_TIMEOUT;
		return 1;
	}

	if ((i=string_to_number(buf, 1, 86400*31)) == -1)
		return 0;

	*timeout = i * HZ;
	return 1;
}


void usage_exit(char *program, const int exit_status)
{
	FILE *stream;

	if (exit_status != 0)
		stream = stderr;
	else
		stream = stdout;

	fprintf(stream,
		"tcpvsadm 0.0.2, 17-May-2001\n"
		"Usage:\n"
		"  %s -A|E -i ident [-s scheduler] [-P port]\n"
		"  %s -D -i ident\n"
		"  %s -C\n"
		"  %s -a|e -i ident -r server-address [-w weight]\n"
		"  %s -d -i ident -r server-address\n"
		"  %s --add-rule -i ident -p pattern -r server-address\n"
		"  %s --del-rule -i ident -p pattern -r server-address\n"
		"  %s -L|l [-n]\n"
		"  %s -h\n\n",
		program, program, program, program, program,
		program, program, program, program);

	fprintf(stream,
		"Commands:\n"
		"Either long or short options are allowed.\n"
		"  --add-service     -A        add virtual service with options\n"
		"  --edit-service    -E        edit virtual service with options\n"
		"  --delete-service  -D        delete virtual service\n"
		"  --clear           -C        clear the whole table\n"
		"  --add-server      -a        add real server with options\n"
		"  --edit-server     -e        edit real server with options\n"
		"  --delete-server   -d        delete real server\n"
		"  --add-rule                  add rule into virtual service\n"
		"  --del-rule                  del rule into virtual service\n"
		"  --list            -L|-l     list the table\n"
		"  --help            -h        display this help message\n\n"
		);

	fprintf(stream,
		"Options:\n"
		"  --ident        -i identity          service identity\n"
		"  --scheduler    -s scheduler         one of wlc|http\n"
		"                                      the default scheduler is %s.\n"
		"  --port         -p port              service port number\n"
		"  --real-server  -r server-address    server-address is host (and port)\n"
		"  --weight       -w weight            capacity of real server\n"
		"  --numeric      -n                   numeric output of addresses and ports\n",
		DEF_SCHED);

	exit(exit_status);
}


void fail(int err, char *text)
{
	printf("%s\n",text);
	exit(err);
}


int host_to_addr(const char *name, struct in_addr *addr)
{
	struct hostent *host;

	if ((host = gethostbyname(name)) != NULL) {
		if (host->h_addrtype != AF_INET ||
		    host->h_length != sizeof(struct in_addr))
			return -1;
		/* warning: we just handle h_addr_list[0] here */
		memcpy(addr, host->h_addr_list[0], sizeof(struct in_addr));
		return 0;
	}
	return -1;
}


char * addr_to_host(struct in_addr *addr)
{
	struct hostent *host;

	if ((host = gethostbyaddr((char *) addr,
				  sizeof(struct in_addr), AF_INET)) != NULL)
		return (char *) host->h_name;
	else
		return (char *) NULL;
}


char * addr_to_anyname(struct in_addr *addr)
{
	char *name;

	if ((name = addr_to_host(addr)) != NULL)
		return name;
	else
		return inet_ntoa(*addr);
}


int service_to_port(const char *name, unsigned short proto)
{
	struct servent *service;

	if (proto == IPPROTO_TCP
	    && (service = getservbyname(name, "tcp")) != NULL)
		return ntohs((unsigned short) service->s_port);
	else if (proto == IPPROTO_UDP
		 && (service = getservbyname(name, "udp")) != NULL)
		return ntohs((unsigned short) service->s_port);
	else
		return -1;
}


char * port_to_service(int port, unsigned short proto)
{
	struct servent *service;

	if (proto == IPPROTO_TCP &&
	    (service = getservbyport(htons(port), "tcp")) != NULL)
		return service->s_name;
	else if (proto == IPPROTO_UDP &&
		 (service = getservbyport(htons(port), "udp")) != NULL)
		return service->s_name;
	else
		return (char *) NULL;
}


char * port_to_anyname(int port, unsigned short proto)
{
	char *name;
	static char buf[10];

	if ((name = port_to_service(port, proto)) != NULL)
		return name;
	else {
		sprintf(buf, "%d", port);
		return buf;
	}
}