#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <sys/select.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h>

struct sock {
  int sock;
  unsigned short port;
  char *direction;
  uint64_t external, internal;
};

int divert_socket (struct sock *sock)
{
  struct sockaddr_in sin;

  if ((sock->sock = socket(PF_INET, SOCK_RAW, IPPROTO_DIVERT)) == -1) return -1;
  memset(&sin, 0, sizeof(struct sockaddr_in));
  sin.sin_family = AF_INET;
  sin.sin_port = htons(sock->port);
  return bind(sock->sock, (struct sockaddr*) &sin, sizeof(struct sockaddr_in));
}

int is_in_mynetworks (struct in_addr ia)
{
  ia = ia;
  return 1;
}

#define NIPQUAD(a) ((a)>>24)&0xff,((a)>>16)&0xff,((a)>>8)&0xff,(a)&0xff
int handle_data (fd_set *fds, struct sock *sock)
{
  char buffer[65536];
  ssize_t len;
  struct sockaddr_in sin;
  socklen_t sin_size = sizeof(struct sockaddr_in);
  struct ip *ip = (struct ip*) buffer;

  if (!FD_ISSET(sock->sock, fds)) return 0;
  if ((len = recvfrom(sock->sock, buffer, sizeof(buffer), 0, (struct sockaddr*) &sin, &sin_size)) == -1)
    return -1;
  if ((unsigned) len < sizeof(struct ip)) return 0;
  if (ip->ip_v != 4) return 0;
  if (ntohs(ip->ip_len) > len) return 0;
  if (ip->ip_hl * 4 > ntohs(ip->ip_len)) return 0;
#if 0
  printf("%u.%u.%u.%u -> %u.%u.%u.%u : %u\n", NIPQUAD(ntohl(ip->ip_src.s_addr)),
         NIPQUAD(ntohl(ip->ip_dst.s_addr)), ntohs(ip->ip_len));
#endif
  if (is_in_mynetworks(ip->ip_dst) && is_in_mynetworks(ip->ip_src))
    sock->internal += ntohs(ip->ip_len);
  else
    sock->external += ntohs(ip->ip_len);
  return 0;
}

int main ()
{
  int i, pkt = 0;
  fd_set fds;
  struct timeval timeout;
  struct sock socks[] = {
    { -1, 10001, "in", 0, 0, },
    { -1, 10002, "out", 0, 0, },
  };
  const int socks_size = sizeof(socks) / sizeof(struct sock);

  for (i = 0; i < socks_size; ++i)
    if (divert_socket(&socks[i]) == -1) {
      perror("divert_socket()");
      return 1;
    }
setgid(1003);setuid(1003);
  for (;;) {
    FD_ZERO(&fds);
    for (i = 0; i < socks_size; ++i)
      FD_SET(socks[i].sock, &fds);
    timeout.tv_sec = 2;
    timeout.tv_usec = 0;
    i = select(socks[socks_size - 1].sock + 1, &fds, NULL, NULL, &timeout);
    if (i == -1) {
      perror("select()");
      return 1;
    }
    if (i > 0) {
      for (i = 0; i < socks_size; ++i)
        if (handle_data(&fds, &socks[i]) == -1) {
          perror("handle_data()");
          return 1;
        }
      if ((pkt = ++pkt % 100) == 0)
        for (i = 0; i < socks_size; ++i)
          printf("%s external:%llu internal:%llu\n", socks[i].direction,
                 socks[i].external, socks[i].internal);
    }
  }

  return 0;
}