/*
 * Copyright 2006-2019, Haiku, Inc. All Rights Reserved.
 * Distributed under the terms of the MIT License.
 *
 * Authors:
 *		Axel Dörfler, axeld@pinc-software.de
 *		James Woodcock
 */


#include <arpa/inet.h>
#include <errno.h>
#include <getopt.h>
#include <net/if.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <SupportDefs.h>

#include <net_stat.h>
#include <syscalls.h>


extern const char* __progname;
const char* kProgramName = __progname;

static int sResolveNames = 1;

enum filter_flags {
	FILTER_FAMILY_MASK		= 0x0000ff,
	FILTER_PROTOCOL_MASK	= 0x00ff00,
	FILTER_STATE_MASK		= 0xff0000,

	// Families
	FILTER_AF_INET			= 0x000001,
	FILTER_AF_INET6			= 0x000002,
	FILTER_AF_UNIX			= 0x000004,

	// Protocols
	FILTER_IPPROTO_TCP		= 0x000100,
	FILTER_IPPROTO_UDP		= 0x000200,

	// States
	FILTER_STATE_LISTEN		= 0x010000,
};

// AF_INET/AF_INET6 families
static void inet_print_address(sockaddr* address);


static void
inet_print_address(sockaddr* address)
{
	if ((address->sa_family != AF_INET && address->sa_family != AF_INET6) ||
		address->sa_len == 0) {
		printf("%-21s", "-");
		return;
	}

	bool resolvedHostName = false;
	char hostName[128];
	if (sResolveNames) {
		// First try to resolve the host name
		if (getnameinfo(address, address->sa_len, hostName, sizeof(hostName),
			NULL, 0, NI_NAMEREQD) == 0) {
			resolvedHostName = true;
		}
	}
	// Now get the port and (optionally) the numeric host address
	char port[32];
	if (getnameinfo(address, address->sa_len,
		resolvedHostName ? NULL : hostName,
		resolvedHostName ? 0 : sizeof(hostName),
		port, sizeof(port),
		NI_NUMERICHOST | (sResolveNames ? 0 : NI_NUMERICSERV)) != 0) {
		printf("%-21s", "-");
		return;
	}

	char buffer[128];
	int length;
	if (address->sa_family != AF_INET
		|| ((sockaddr_in*)address)->sin_addr.s_addr != INADDR_ANY) {
		length = snprintf(buffer, sizeof(buffer),
			resolvedHostName ? "%s" : "[%s]", hostName);
	} else {
		// Special handling for IPv4 any address
		strcpy(buffer, "*");
		length = 1;
	}
	if (length < sizeof(buffer)) {
		length += snprintf(buffer + length, sizeof(buffer) - length,
			":%s", port);
	}

	printf("%-21s", buffer);
}


//	#pragma mark -


void
usage(int status)
{
	printf("Usage: %s [-nh]\n", kProgramName);
	printf("Options:\n");
	printf("	-n	don't resolve names\n");
	printf("	-h	this help\n");
	printf("Filter options:\n");
	printf("	-4	IPv4\n");
	printf("	-6	IPv6\n");
	printf("	-x	Unix\n");
	printf("	-t	TCP\n");
	printf("	-u	UDP\n");
	printf("	-l	listen state\n");

	exit(status);
}


//	#pragma mark -


int
main(int argc, char** argv)
{
	int optionIndex = 0;
	int opt;
	int filter = 0;

	const static struct option kLongOptions[] = {
		{"help", no_argument, 0, 'h'},
		{"numeric", no_argument, 0, 'n'},

		{"inet", no_argument, 0, '4'},
		{"inet6", no_argument, 0, '6'},
		{"unix", no_argument, 0, 'x'},

		{"tcp", no_argument, 0, 't'},
		{"udp", no_argument, 0, 'u'},

		{"listen", no_argument, 0, 'l'},

		{0, 0, 0, 0}
	};

	do {
		opt = getopt_long(argc, argv, "hn46xtul", kLongOptions,
			&optionIndex);
		switch (opt) {
			case -1:
				// end of arguments, do nothing
				break;

			case 'n':
				sResolveNames = 0;
				break;

			// Family filter
			case '4':
				filter |= FILTER_AF_INET;
				break;
			case '6':
				filter |= FILTER_AF_INET6;
				break;
			case 'x':
				filter |= FILTER_AF_UNIX;
				break;
			// Protocol filter
			case 't':
				filter |= FILTER_IPPROTO_TCP;
				break;
			case 'u':
				filter |= FILTER_IPPROTO_UDP;
				break;
			// State filter
			case 'l':
				filter |= FILTER_STATE_LISTEN;
				break;

			case 'h':
			default:
				usage(0);
				break;
		}
	} while (opt != -1);

	bool printProgram = true;
		// TODO: add some more program options... :-)

	printf("Proto  Recv-Q Send-Q Local Address         Foreign Address       "
		"State        Program\n");

	uint32 cookie = 0;
	int family = -1;
	net_stat stat;
	while (_kern_get_next_socket_stat(family, &cookie, &stat) == B_OK) {
		// Filter families
		if ((filter & FILTER_FAMILY_MASK) != 0) {
			if ((stat.family != AF_INET && stat.family != AF_INET6
					&& stat.family != AF_UNIX)
				|| ((filter & FILTER_AF_INET) == 0 && stat.family == AF_INET)
				|| ((filter & FILTER_AF_INET6) == 0 && stat.family == AF_INET6)
				|| ((filter & FILTER_AF_UNIX) == 0 && stat.family == AF_UNIX))
				continue;
		}
		// Filter protocols
		if ((filter & FILTER_PROTOCOL_MASK) != 0) {
			if ((stat.protocol != IPPROTO_TCP && stat.protocol != IPPROTO_UDP)
				|| ((filter & FILTER_IPPROTO_TCP) == 0
					&& stat.protocol == IPPROTO_TCP)
				|| ((filter & FILTER_IPPROTO_UDP) == 0
					&& stat.protocol == IPPROTO_UDP))
				continue;
		}
		if ((filter & FILTER_STATE_MASK) != 0) {
			if ((filter & FILTER_STATE_LISTEN) == 0
				|| strcmp(stat.state, "listen") != 0)
				continue;
		}

		protoent* proto = getprotobynumber(stat.protocol);
		if (proto != NULL)
			printf("%-6s ", proto->p_name);
		else
			printf("%-6d ", stat.protocol);

		printf("%6lu ", stat.receive_queue_size);
		printf("%6lu ", stat.send_queue_size);

		inet_print_address((sockaddr*)&stat.address);
		printf(" ");
		inet_print_address((sockaddr*)&stat.peer);
		printf(" %-12s ", stat.state);

		team_info info;
		if (printProgram && get_team_info(stat.owner, &info) == B_OK) {
			// remove arguments
			char* name = strchr(info.args, ' ');
			if (name != NULL)
				name[0] = '\0';

			// remove path name
			name = strrchr(info.args, '/');
			if (name != NULL)
				name++;
			else
				name = info.args;

			printf("%" B_PRId32 "/%s\n", stat.owner, name);
		} else
			printf("%" B_PRId32 "\n", stat.owner);
	}

	return 0;
}