// Copyright (C) FIT VUT
// Petr Lampa <lampa@fit.vutbr.cz>
// $Id$
// vi:set ts=8 sts=8 sw=8:
//

#ifndef __NDWATCH_PCAP_H__
#define __NDWATCH_PCAP_H__

extern "C" {
#include <pcap.h>
#include <errno.h>
}

class PCAP
{
private:
	pcap_t *_handle[32];
	int _fds[32];
	string _devs[32];
	int _num;
	int _maxfd;
	bool _fd0nonbl;
	fd_set _fdset;

	PCAP(const PCAP&);
	PCAP& operator=(const PCAP&);
public:
	PCAP(): _num(0), _maxfd(0), _fd0nonbl(false), _fdset() { FD_ZERO(&_fdset); }
	bool open(const string &ifc="");
	bool set_filter(const string &str);
	bool read(Buffer &pkt, string &dev, unsigned long timeout = 0);
	bool write(const string &dev, Buffer &pkt);
	void close();
	~PCAP() { close(); }
};

bool PCAP::open(const string &ifc) 
{ 
	char errbuf[PCAP_ERRBUF_SIZE];

	if (ifc == "") {
		char *dev = pcap_lookupdev(errbuf);
		if (dev == NULL) {
			if (debug > 0) cerr << "pcap_lookup: " << errbuf << endl;
			return false;
		}
		_devs[_num] = dev;
	} else {
		_devs[_num] = ifc;
	}
	if ((_handle[_num] = pcap_open_live(_devs[_num].c_str(), BUFSIZ, 1, 1000, errbuf)) == NULL) {
		if (debug > 0) cerr << "pcap_open(" << _devs[_num] << ")" << errbuf << endl;
		return false;
	}
	_fds[_num] = pcap_get_selectable_fd(_handle[_num]);
	FD_SET(_fds[_num], &_fdset);
	if (_fds[_num] > _maxfd) _maxfd = _fds[_num];
	if (_num == 1) {   // when more then 1, set also the first one
		if (pcap_setnonblock(_handle[0], 1, errbuf) < 0) {
			if (debug > 0) cerr << "pcap_setnonblock: " << errbuf << endl;
			return false;
		}
	}
	if (_num > 0) {
		if (pcap_setnonblock(_handle[_num], 1, errbuf) < 0) {
			if (debug > 0) cerr << "pcap_setnonblock: " << errbuf << endl;
			return false;
		}
	}
	_num++;
	return true;
}

bool PCAP::set_filter(const string &str)
{
	struct bpf_program bp;
	if (_num == 0) return false;
	// linux const problem XXX
	if (pcap_compile(_handle[0], &bp, const_cast<char *>(str.c_str()), 1, 0) == -1) return false;
	for (int i = 0; i < _num; i++) {
		if (pcap_setfilter(_handle[i], &bp) == -1) return false;
	}
	pcap_freecode(&bp);
	return true;
}

// timeout in ms
bool PCAP::read(Buffer &pkt, string &dev, unsigned long timeout)
{
	struct pcap_pkthdr header;
	char errbuf[PCAP_ERRBUF_SIZE];
	const Octet *p;
	if (_num == 0) return false;
	if (_num == 1 && timeout == 0) {
		if (_fd0nonbl) {
			if (pcap_setnonblock(_handle[0], 0, errbuf) < 0) {
			}
			_fd0nonbl = false;
		}
		p = reinterpret_cast<const Octet *>(pcap_next(_handle[0], &header));
		if (!p) return false;
		pkt.buf(p, header.len);
		dev = _devs[0];
		return true;
	}
	if (_num == 1 && timeout && !_fd0nonbl) {
		if (pcap_setnonblock(_handle[0], 1, errbuf) < 0) {
		}
		_fd0nonbl = true;
	}
	fd_set readfds;
	for (int i = 0; i < _num; i++) {
		p = reinterpret_cast<const Octet *>(pcap_next(_handle[i], &header));
		if (p) {
			pkt.buf(p, header.len);
			dev = _devs[i];
			return true;
		}
	}
	int n;
	do {
		memcpy(&readfds, &_fdset, sizeof(fd_set));
		struct timeval *to, toval;
		if (timeout) {
			to = &toval;
			toval.tv_usec = (timeout%1000)*1000;
			toval.tv_sec = timeout/1000;
		} else {
			to = NULL;
		}
		n = select(_maxfd, &readfds, NULL, NULL, to);
		if (n > 0) {
			for (int i = 0; i < _num; i++) {
				if (!FD_ISSET(_fds[i], &readfds)) continue;
				p = reinterpret_cast<const Octet *>(pcap_next(_handle[i], &header));
				if (p) {
					pkt.buf(p, header.len);
					dev = _devs[i];
					return true;
				}
			}
		}
	} while (n < 0 && errno == EINTR);
	return false;
}

bool PCAP::write(const string &dev, Buffer &pkt)
{
	if (_num == 0) return false;
	for (int i = 0; i < _num; i++) {
		if (_devs[i] != dev) continue;
		if (static_cast<unsigned>(pcap_inject(_handle[0], pkt.buf(), pkt.length())) != pkt.length()) return false;
		return true;
	}
	return false;
}

void PCAP::close() {
	for (int i = 0; i < _num; i++) pcap_close(_handle[i]);
}
#endif
