/*-
 * Copyright (c) 2004 Free (Olivier Beyssac)
 * 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 AND CONTRIBUTORS ``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 OR CONTRIBUTORS 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 <stdio.h>
#include <stdlib.h>
#include <syslog.h>
#include <time.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include "ipinfo.h"
#include "iptree.h"
#include "options.h"


#define get_pos(x, y) (x + y * 256)


extern struct options opt;


#define IPTREE_M_IPLIST    (0)
#define IPTREE_M_BLACKLIST (1)


/* dlist is a common double linked list made of dl_nodes */
typedef struct dlist *dlist;


struct dl_node
{
  struct ipinfo *p;
  struct dl_node *prev;
  struct dl_node *next;
};


struct dlist
{
  struct dl_node *first;
  struct dl_node *last;
  size_t count;
};


/* iptree stores IP addresses
   IP x.y.z.t is in the tree at position iptree->[x + 256 * y]->c[z]->d[t] */
struct ipt_leaf
{
  size_t count;
  struct ipinfo *d[256];
};


struct ipt_node
{
  size_t count;
  struct ipt_leaf *c[256];
};


struct iptree
{
  struct ipt_node *ab[256 * 256];
  dlist ipl;                   /* Whole IPs list */
  dlist bl;                    /* Blacklisted IPs */
};


static void iptree_free(const iptree ipt, struct ipinfo *val);


static void *xcalloc(size_t nmemb, size_t size)
{
  void *tmp;

  if ((tmp = calloc(nmemb, size)) == NULL)
    syslog(LOG_ERR, "unable to calloc()");

  return tmp;
}


/* Dump dlist to file */
static void dlist_dump2file(const dlist dl, FILE *f)
{
  struct dl_node *p;

  for (p = dl->first; p; p = p->next)
    fwrite(p->p, sizeof(struct ipinfo), 1, f);
}


/* Remove the first element of dlist and return its address */
static struct dl_node *dlist_clean(const iptree ipt, const int mode)
{
  dlist l;
  struct dl_node *p;

  if (mode == IPTREE_M_IPLIST) {
    l = ipt->ipl;
    l->first->p->iplp = NULL;
  } else {
    l = ipt->bl;
    l->first->p->blp = NULL;
  }
  
  if (l->first == NULL)
    return NULL;
  
  p = l->first;
  
  /* Free only if
     - we are in the IP list and entry is not blacklisted,
     - we want to explicitly delete from blacklist and entry is
       not in IP list */
  if ((mode == IPTREE_M_IPLIST && !p->p->blp)
      || (mode == IPTREE_M_BLACKLIST && !p->p->iplp)) {
    if (opt.log_level >= 3)
      syslog(LOG_INFO, "Going to free %p (last=%lu - mode=%d)",
             (void *)p->p, (unsigned long)p->p->last, mode);
    iptree_free(ipt, p->p);
  }

  l->first = p->next;
  if (l->first)
    l->first->prev = NULL;
  l->count--;
  p->prev = NULL;
  p->next = NULL;

  return p;
}


/* Return a newly allocated dlist */
static dlist dlist_init(void)
{
  dlist tmp;

  tmp = xcalloc(sizeof(struct dlist), 1);
  
  return tmp;
}


/* Add an ipinfo record at the end of the wanted dlist (mode is IPTREE_M_IPLIST
   or M_BLACKLIST)
   Return 1 if success, 0 if allocation failed */
static int dlist_add(const iptree ipt, struct ipinfo *val,
		     const int mode, const time_t req_time)
{
  dlist l;
  unsigned int list_size;
  struct dl_node *p;
  time_t delta;

  if (val->blp) {
    if (opt.log_level >= 3)
      syslog(LOG_INFO, "%u.%u.%u.%u already in bl",
	     val->iptbl[3], val->iptbl[2], val->iptbl[1], val->iptbl[0]);
    return 1;
  }
  
  if (mode == IPTREE_M_IPLIST) {
    l = ipt->ipl;
    p = val->iplp;
    list_size = opt.list_size;
  } else {
    l = ipt->bl;
    p = val->blp;
    list_size = opt.blacklist_size;
  }

  if (p == NULL) {
    /* Add entry to list */
    l->count++;
    
    /* See if we can get an expirable, already allocated node, otherwise alloc
       a brand new one */
    if (l->count >= list_size + 1)
      p = dlist_clean(ipt, mode);
    else if ((p = xcalloc(sizeof(struct dl_node), 1)) == NULL) {
      syslog(LOG_ERR, "dlist_add failed to allocate %d bytes",
	     sizeof(struct dl_node));
      return 0;
    }
    
    if (mode == IPTREE_M_IPLIST)
      val->iplp = p;
    else
      val->blp = p;
    
    p->p = val;
    
    if (opt.log_level >= 3)
      syslog(LOG_INFO, "%s list size: %u/%u",
	     mode == IPTREE_M_IPLIST?"IP":"black", l->count, list_size);
    
    if (l->first == NULL) {
      l->first = p;
      l->last = p;
      return 1;
    }
  } else {
    /* Already existing entry, see if we have to blacklist */
    delta = val->last - val->start;
    if (mode == IPTREE_M_IPLIST && !val->blp
	&& delta >= opt.interval
	&& val->count * opt.interval / delta >= opt.max_req) {
      /* Blacklist
	 We will not remove the entry from the IP list and just wait
	 until the entry expires (no more updates once it's
	 blacklisted) */
      dlist_add(ipt, val, IPTREE_M_BLACKLIST, req_time);

      if (opt.log_level >= 1)
	syslog(LOG_INFO,
	       "%u.%u.%u.%u put in bl: %u reqs / %lu secs",
	       val->iptbl[3], val->iptbl[2], val->iptbl[1], val->iptbl[0],
	       val->count, (unsigned long)delta);
    }

    if (l->last == p)
      return 1;

    if (l->first == p)
      l->first = p->next;
  }

  /* Do the generic moving stuff */
  l->last->next = p;

  if (p->prev)
    p->prev->next = p->next;
  if (p->next)
    p->next->prev = p->prev;

  p->next = NULL;
  p->prev = l->last;
  l->last = p;
  l->first->prev = NULL;

  return 1;
}


#if 0
/* Call for debugging purposes only: brute force dump */
static void iptree_dump(const iptree ipt)
{
  size_t i, j, k;

  for (i = 0; i < 256 * 256; i++)
    if (ipt->ab[i])
      for (j = 0; j < 256; j++)
	if (ipt->ab[i]->c[j])
	  for (k = 0; k < 256; k++)
	    if (ipt->ab[i]->c[j]->d[k])
	      syslog(LOG_INFO, "iptree: %u.%u.%u.%u: %u %lu %lu",
		     i%256, i/256, j, k,
		     ipt->ab[i]->c[j]->d[k]->count,
		     ipt->ab[i]->c[j]->d[k]->start,
		     ipt->ab[i]->c[j]->d[k]->last);
}
#endif

 
/* Read dlist from file */
static void iptree_getfromfile(const iptree ipt, FILE *f)
{
  struct ipinfo ipi;
  int svlog_level;
  unsigned int val;
  
  svlog_level = opt.log_level;
  opt.log_level = 0;
  
  while (fread(&ipi, sizeof(struct ipinfo), 1, f)) {
    val = (ipi.iptbl[3]<<24) + (ipi.iptbl[2]<<16) + (ipi.iptbl[1]<<8)
      + ipi.iptbl[0];
    iptree_add(ipt, val, ipi.start, ipi.last, ipi.count, 0);
  }
  
  opt.log_level = svlog_level;
}
 

/* Return a newly allocated iptree */
extern iptree iptree_init(void)
{
  iptree tmp;
  struct stat sb;
  FILE *f;

  /* FIXME: should add a test */
  tmp = malloc(sizeof(struct iptree));
  memset(tmp->ab, 0, 256 * 256);
  tmp->ipl = dlist_init();
  tmp->bl = dlist_init();
  
  if (stat(opt.wl_filename, &sb) == 0 && S_ISREG(sb.st_mode)) {
    double records = (double)sb.st_size / (double)sizeof(struct ipinfo);
    int check = sb.st_size / sizeof(struct ipinfo);

    /* Sanity check */
    if ((double)check != records)
      return tmp;
    
    if ((f = fopen(opt.wl_filename, "r")) != NULL) {
      syslog(LOG_INFO, "loading information from previous dump");
      iptree_getfromfile(tmp, f);
      fclose(f);
    }
  }

  return tmp;
}


/* Add an IP address to a tree */
extern int iptree_add(const iptree ipt, const unsigned int val,
		      const time_t cr_time, const time_t req_time,
		      const int count, const int options)
{
  union {
    long l;
    unsigned char c[sizeof(val)];
  } u;
  unsigned char *iptbl;
  unsigned long pos;
  size_t i;

  u.l = val;
  iptbl = u.c;

  pos = get_pos(iptbl[3], iptbl[2]);

  if (!ipt->ab[pos]) {
    ipt->ab[pos] = xcalloc(sizeof(struct ipt_node), 1);
    if (ipt->ab[pos] == NULL) {
      syslog(LOG_ERR, "iptree_add failed to alloc %d bytes for %u.%u.%u.%u 1",
	     sizeof(struct ipt_node), iptbl[3], iptbl[2], iptbl[1], iptbl[0]);
      return 0;
    }
  }
  
  if (!ipt->ab[pos]->c[iptbl[1]]) {
    ipt->ab[pos]->c[iptbl[1]] = xcalloc(sizeof(struct ipt_leaf), 1);
    if (ipt->ab[pos]->c[iptbl[1]] == NULL) {
      syslog(LOG_ERR, "iptree_add failed to alloc %d bytes for %u.%u.%u.%u 2",
	     sizeof(struct ipt_node), iptbl[3], iptbl[2], iptbl[1], iptbl[0]);
      if (!ipt->ab[pos]->count) {
	syslog(LOG_ERR, "iptree_add freed block %u.%u.", iptbl[3], iptbl[2]);
	free(ipt->ab[pos]);
	ipt->ab[pos] = NULL;
      }
      return 0;
    }
  }
  
  if (!ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]) {
    ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]] = xcalloc(sizeof(struct ipinfo), 1);
    if (ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]] == NULL) {
      syslog(LOG_ERR, "iptree_add failed to alloc %d bytes for %u.%u.%u.%u 3",
	     sizeof(struct ipt_node), iptbl[3], iptbl[2], iptbl[1], iptbl[0]);
      if (!ipt->ab[pos]->c[iptbl[1]]->count) {
	syslog(LOG_ERR, "freed block %u.%u.%u", iptbl[3], iptbl[2], iptbl[1]);
	free(ipt->ab[pos]->c[iptbl[1]]);
	ipt->ab[pos]->c[iptbl[1]] = NULL;
      }
      if (!ipt->ab[pos]->count) {
	syslog(LOG_ERR, "freed block %u.", iptbl[3]);
	free(ipt->ab[pos]);
	ipt->ab[pos] = NULL;
      }
      return 0;
    }
    ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->start = cr_time;
    ipt->ab[pos]->c[iptbl[1]]->count++;
    ipt->ab[pos]->count++;
  }

  for (i = 0; i < sizeof(long); i++)
    ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->iptbl[i] = iptbl[i];

  ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->options |= options;
  
  /* Decrement only if IP count is greater than zero */
  if (count > 0 || ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->count > 0)
    ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->count += count;

  ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->last = req_time;

  if ((options & IPINFO_OPT_FORCED)) {
    if (dlist_add(ipt, ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]],
                  IPTREE_M_BLACKLIST, req_time) == 0) {
      syslog(LOG_INFO, "iptree_add cleaning last added bl entry %u.%u.%u.%u (%p)",
             iptbl[3], iptbl[2], iptbl[1], iptbl[0],
	     (void *)ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]);
      iptree_free(ipt, ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]);
      
      return 0;
    } else
      return 1;
  }
  
  if (dlist_add(ipt, ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]],
                IPTREE_M_IPLIST, req_time) == 0
      && ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->count == 1
      && ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->blp == NULL) {
    syslog(LOG_INFO, "iptree_add cleaning last added entry %u.%u.%u.%u (%p)",
	   iptbl[3], iptbl[2], iptbl[1], iptbl[0],
	   (void *)ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]);
    iptree_free(ipt, ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]);
    
    return 0;
  }
 
  return 1;
}


/* Remove an ipinfo record from the tree and any unused node above */
static void iptree_free(const iptree ipt, struct ipinfo *val)
{
  unsigned long pos;
  unsigned char *iptbl;

  /* Shortcut */
  iptbl = val->iptbl;

  pos = get_pos(iptbl[3], iptbl[2]);

  if (opt.log_level >= 3)
    syslog(LOG_INFO, "free IP %u.%u.%u.%u (%p)",
	   iptbl[3], iptbl[2], iptbl[1], iptbl[0], (void *)val);

  free(val);
  
  ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]] = NULL;
  ipt->ab[pos]->c[iptbl[1]]->count--;
  if (!ipt->ab[pos]->c[iptbl[1]]->count) {
    free(ipt->ab[pos]->c[iptbl[1]]);
    ipt->ab[pos]->c[iptbl[1]] = NULL;
    
    if (opt.log_level >= 3)
      syslog(LOG_INFO, "freed block %u.%u.%u.", iptbl[3], iptbl[2], iptbl[1]);
  }
  
  ipt->ab[pos]->count--;
  if (!ipt->ab[pos]->count) {
    free(ipt->ab[pos]);
    ipt->ab[pos] = NULL;
    
    if (opt.log_level >= 3)
      syslog(LOG_INFO, "freed block %u.%u.", iptbl[3], iptbl[2]);
  }
  
  return;
}


/* Dump all iptree lists to files */
extern void iptree_dump2files(const iptree ipt)
{
  FILE *f;

  if ((f = fopen(opt.wl_filename, "w+")) == NULL) {
    syslog(LOG_ERR, "fopen: %s", strerror(errno));
    return;
  }
  dlist_dump2file(ipt->ipl, f);
  if (opt.log_level >= 3)
    syslog(LOG_INFO, "IP list dumped to %s", opt.wl_filename);
  fclose(f);
  

  if ((f = fopen(opt.bl_filename, "w+")) == NULL) {
    syslog(LOG_ERR, "fopen: %s", strerror(errno));
    return;
  }
  dlist_dump2file(ipt->bl, f);
  if (opt.log_level >= 3)
    syslog(LOG_INFO, "black list dumped to %s", opt.bl_filename);
  fclose(f);
}


/* Return 1 if entry is blacklisted, 0 otherwise */
extern int iptree_get_bl(const iptree ipt, const unsigned long val,
			 const time_t req_time)
{
  unsigned char *iptbl;
  unsigned long pos;
  time_t t, delta;

  iptbl = (unsigned char *)&val;

  pos = get_pos(iptbl[3], iptbl[2]);

  if (ipt->ab[pos] && ipt->ab[pos]->c[iptbl[1]]
      && ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]
      && ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->blp) {
    t = req_time - ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->last;
    if (t <= opt.blacklist_expiration) {
      delta = ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->last
	- ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->start;
      if (ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->options & IPINFO_OPT_FORCED
	  || (delta >= opt.interval
	      && (ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->count * opt.interval
		  / delta >= opt.max_req))) {
	if (opt.log_level >= 3)
	  syslog(LOG_INFO, "request for a blacklisted IP: %u.%u.%u.%u",
		 iptbl[3], iptbl[2], iptbl[1], iptbl[0]);
	return 1;
      }
    } else {
      /* Redemption */
      ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->count = 0;
      ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->options = 0;
      ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->start = req_time;
      ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->last = req_time;

      /* if entry is still in IP list, update it */
      if (ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]->iplp) {
        if (opt.log_level >= 3)
          syslog(LOG_INFO, "re-adding %u.%u.%u.%u (%p)",
                 iptbl[3], iptbl[2], iptbl[1], iptbl[0],
                 (void *)ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]]);
	dlist_add(ipt, ipt->ab[pos]->c[iptbl[1]]->d[iptbl[0]],
		  IPTREE_M_IPLIST, req_time);
      }
    }
  }

  return 0;
}


/* Show some stats via syslog */
extern void iptree_show_stats(const iptree ipt, const time_t t)
{
  if (ipt->ipl->first == NULL)
    syslog(LOG_INFO, "IP list: %u/%u", ipt->ipl->count, opt.list_size);
  else
    syslog(LOG_INFO, "IP list: %u/%u (oldest: %lu secs)",
	   ipt->ipl->count, opt.list_size,
           (unsigned long)t - ipt->ipl->first->p->last);

  if (ipt->bl->first == NULL)
    syslog(LOG_INFO, "Blacklist: %u/%u", ipt->bl->count, opt.blacklist_size);
  else
    syslog(LOG_INFO, "Blacklist: %u/%u (oldest: %lu secs)",
	   ipt->bl->count, opt.blacklist_size,
           (unsigned long)t - ipt->bl->first->p->last);
}