/*-
 * Copyright (c) 2005 Fredrik Lindberg. <fli at shapeshifter dot se>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <sys/param.h>
#include <sys/socket.h>
#include <sys/sysctl.h>
#include <sys/time.h>
#include <sys/linker.h>

#include <net/if.h>
#include <net/if_mib.h>
#include <net/if_var.h>
#include <netinet/in.h>
#include <netinet/in_var.h>
#include <arpa/inet.h>

#include <errno.h>
#include <fcntl.h>
#include <kvm.h>
#include <stdio.h>
#include <stdlib.h>
#if __FreeBSD_version >= 500000
# include <stdint.h>
#else
# include <inttypes.h>
#endif
#include <string.h>
#include <unistd.h>

#include "bw.h"

#ifndef IFNAMSIZ
#define IFNAMSIZ	16
#endif

static size_t _kvm_read(kvm_t *, unsigned long, void *, size_t);
uint32_t _kvm_addrlookup(kvm_t *, const void *, uint8_t, uint32_t *); 

/*
 * Quick and dirty way to access stats on a
 * interface basis. 
 * Doesn't require uid 0 nor access to /dev/mem
 */
int
bw_sysctl(const char *iface, uint64_t *bw, unsigned int sleep_time)
{
	int     name[6], i;
	size_t  len;
	struct  ifmibdata ifmb[2];
	u_char  flg_found = 0;
	int     ifcount;

	/* Grab total number of interfaces */
	len = sizeof(ifcount);
	if (sysctlbyname("net.link.generic.system.ifcount", &ifcount, &len, (void *)0, 0) == -1) {
		perror("sysctl");
		return -1;
	}

	len = sizeof(struct ifmibdata);
	name[0] = CTL_NET;
	name[1] = AF_LINK;
	name[2] = NETLINK_GENERIC;
	name[3] = IFMIB_IFDATA;
	name[5] = IFDATA_GENERAL;

	/* Find the right interface index */
	i = 1;
	while (flg_found == 0 && i <= ifcount)
	{
		name[4] = i;
		sysctl(name, 6, &ifmb[0], &len, 0, 0);
		if (strncmp(ifmb[0].ifmd_name, iface, strlen(iface)) == 0)
			flg_found = 1;
		i++;
	}

	if (flg_found == 0)
		return -1;

	name[4] = i - 1;

	sysctl(name, 6, &ifmb[0], &len, 0, 0);
	sleep(sleep_time); 
	sysctl(name, 6, &ifmb[1], &len, 0, 0);

	bw[0] = ifmb[1].ifmd_data.ifi_ibytes - ifmb[0].ifmd_data.ifi_ibytes;
	bw[1] = ifmb[1].ifmd_data.ifi_obytes - ifmb[0].ifmd_data.ifi_obytes;

	return 0;
}

/*
 * Function used to lookup in-kernel address to
 * correct ifnet structure. 
 */
uint32_t
_kvm_addrlookup(kvm_t *kvm, const void *src, uint8_t type, uint32_t *addr_ifnet) 
{

	uint32_t addr, addr_ifaddr, src_addr = 0;
	struct 	kld_sym_lookup ksl;
	struct	ifnethead ifnethead;
	struct	ifnet ifn;
	struct	ifaddr ifaddr;

#if __FreeBSD_version < 500000
	char 	ifn_buffer[IFNAMSIZ];
#endif

	memset(&ksl, 0, sizeof(ksl));
	ksl.version = sizeof(struct kld_sym_lookup); 
	ksl.symname = "ifnet";

	/* Lookup the address to ``ifnet'' symbol */ 
	if (kldsym(0, KLDSYM_LOOKUP, &ksl) != 0) {
		perror("kldsym");
		return -1;
	}

	addr = ksl.symvalue;


	/* Grab the address of the list head */
	if (_kvm_read(kvm, addr, &ifnethead, sizeof(struct ifnethead)) < 0) { 
		return -1;
	}

	/*
	 * Interfaces and assigned addresses are stored in a linked, chained list.
	 * We search it linear to find the memory address offset we are 
	 * looking for (interface or address)
	 * If ``src_addr'' is 0 after this section the searched item couln't be found.
 	 */
	for (addr = (u_long)TAILQ_FIRST(&ifnethead); addr > 0; addr = (u_long)TAILQ_NEXT(&ifn, if_link)) {
		_kvm_read(kvm, addr, &ifn, sizeof(struct ifnet));

		/* 
		 * This stuff below requires a hack, because of changes to the 
		 * ifnet structure between 4.x and 5.x.
		 */

		if (type == TYPE_IF) {
#if __FreeBSD_version >= 500000  
			if (strncmp(ifn.if_xname, (char *)src, strlen((char *)src)) == 0)
#else
			_kvm_read(kvm, (u_long)ifn.if_name, ifn_buffer, IFNAMSIZ);
			if ((strncmp(ifn_buffer, (char *)src, strlen((char *)src) - 1) == 0) && 
				((int)*((char *)src + (strlen((char *)src) - 1) - 48) == ifn.if_unit)) 
#endif
				src_addr = addr;
				if (addr_ifnet != NULL)
					*addr_ifnet = addr;
		}

		else if (type == TYPE_ADR) {
			for (addr_ifaddr = (u_long)TAILQ_FIRST(&ifn.if_addrhead); addr_ifaddr > 0;
				 addr_ifaddr = (u_long)TAILQ_NEXT(&ifaddr, ifa_link)) {

				struct sockaddr a;
				struct sockaddr_in *sin;
	
				_kvm_read(kvm, addr_ifaddr, &ifaddr, sizeof(struct ifaddr));
				_kvm_read(kvm, (u_long)ifaddr.ifa_addr, &a, sizeof(struct sockaddr));

				sin = (struct sockaddr_in *)&a;

				if (a.sa_family == AF_INET) {
					if (sin->sin_addr.s_addr == (*((in_addr_t *)(&src))))
						src_addr = addr_ifaddr; 
						if (addr_ifnet != NULL)
							*addr_ifnet = addr;
				}
				/* XXX: AF_INET6 should be implemented */
			}
		}
	}

	return src_addr;
}

/*
 * Read max speed from "device"
 */
uint32_t
bw_maxspeed(const void *src, uint8_t type) 
{
	kvm_t	*kvm;
	uint32_t src_addr = 0;
	struct ifnet _ifnet;

	/* Open a filedescriptor to the kernel memory interface */
	if ((kvm = kvm_open(0, 0, 0, O_RDONLY, NULL)) == NULL) {
		return -1;
	}

	_kvm_addrlookup(kvm, src, type, &src_addr);

	_ifnet.if_baudrate = 0;
	kvm_read(kvm, src_addr, &_ifnet, sizeof(struct ifnet));
	kvm_close(kvm);

	return _ifnet.if_baudrate;
}

/*
 * Does a "device" exists ?
 */
int8_t
bw_exists(const void *src, uint8_t type)
{
	kvm_t	 *kvm;
	uint32_t src_addr = 0;
	int8_t	 exists = 0;
	int      name[6], i, ifcount;
	size_t   len;
	struct   ifmibdata ifmb[2];

	/* Use the sysctl-method to find interface, no uid 0 required */
	if (type == TYPE_IF) {
		char *iface = (char *)src;
		/* Grab total number of interfaces */
		len = sizeof(ifcount);
		if (sysctlbyname("net.link.generic.system.ifcount", &ifcount, &len, (void *)0, 0) == -1) {
			return 0;
		}

		len = sizeof(struct ifmibdata);
		name[0] = CTL_NET;
		name[1] = AF_LINK;
		name[2] = NETLINK_GENERIC;
		name[3] = IFMIB_IFDATA;
		name[5] = IFDATA_GENERAL;

		/* Find the right interface index */
		for (i = 1; i <= ifcount; i++) {
			name[4] = i;
			sysctl(name, 6, &ifmb[0], &len, 0, 0);
			if (strncmp(ifmb[0].ifmd_name, iface, strlen(iface)) == 0) {
				exists = 1;
				break;
			}
		}
	}
	/* To find an ip-address we must use kvm, thus requiring uid 0*/
	else if (type == TYPE_ADR) {

		/* Open a filedescriptor to the kernel memory interface */
		if ((kvm = kvm_open(0, 0, 0, O_RDONLY, NULL)) == NULL) {
			return 0;
		}

		src_addr = _kvm_addrlookup(kvm, src, type, NULL);
		kvm_close(kvm);

		exists = (src_addr > 0) ? 1 : 0;
	}

	/* Set errno to device not configured if we can't fint it */
	if (!exists)
		errno = ENXIO; 

	return exists;
}


/*
 * Grabs statistics directly from kernel memory using
 * kvm(3), requires access to /dev/mem
 * 
 * Can produce statistics on both interface and
 * assigned address basis.
 */
int
bw_kvm(const void *src, uint8_t type, uint64_t *bw, unsigned int sleep_time) 
{
	kvm_t	*kvm;
	uint32_t src_addr = 0;

	/* Open a filedescriptor to the kernel memory interface */
	if ((kvm = kvm_open(0, 0, 0, O_RDONLY, NULL)) == NULL) {
		perror("kvm");
		return -1;
	}

	src_addr = _kvm_addrlookup(kvm, src, type, NULL);

	/*
	 * Perform the actual measure
	 */
	if (src_addr > 0) {
		if (type == TYPE_IF) {
			struct ifnet _ifnet[2];

			kvm_read(kvm, src_addr, &_ifnet[0], sizeof(struct ifnet));
			sleep(sleep_time);
			kvm_read(kvm, src_addr, &_ifnet[1], sizeof(struct ifnet));

			bw[0] = _ifnet[1].if_ibytes - _ifnet[0].if_ibytes;
			bw[1] = _ifnet[1].if_obytes - _ifnet[0].if_obytes;

		}
		else if (type == TYPE_ADR) {
			struct ifaddr _ifaddr[2];

			kvm_read(kvm, src_addr, &_ifaddr[0], sizeof(struct ifaddr));
			sleep(sleep_time);
			kvm_read(kvm, src_addr, &_ifaddr[1], sizeof(struct ifaddr));

			bw[0] = _ifaddr[1].if_ibytes - _ifaddr[0].if_ibytes;
			bw[1] = _ifaddr[1].if_obytes - _ifaddr[0].if_obytes;
		} 	
	}
	else {
		kvm_close(kvm);
		return -1;
	}

	kvm_close(kvm);	

	return 0;
}

/*
 * Wrapper function to kvm_read
 */
static size_t
_kvm_read(kvm_t *kvm, unsigned long addr, void *buf, size_t nbytes)
{
	size_t tmp; 
	if ((tmp = kvm_read(kvm, addr, buf, nbytes)) < 0) {
		perror("kvm_read");
	}

	return tmp;
}