/***************************************************************************
 *
 * Copyright (c) 2000, 2001 BalaBit IT Ltd, Budapest, Hungary
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 * $Id: ctracker.c,v 1.30.2.3 2001/09/19 19:22:16 bazsi Exp $
 *
 ***************************************************************************/

#include <zorp/zorp.h>
#include <zorp/log.h>
#include <zorp/ctrack.h>
#include <zorp/policy.h>
#include <zorp/pyvars.h>
#include <zorp/poll.h>
#include <zorp/thread.h>
#include <zorp/packsock.h>
#include <zorp/pyproxy.h>
#include <zorp/pyctrack.h>
#include <zorp/pysockaddr.h>

#if ENABLE_CONNTRACK

typedef struct _ZConnectPool
{
  ZProxy *proxy;
  struct sockaddr_in from;
  struct sockaddr_in to;
  int fd;
  time_t time;
  gchar *key;
  gpointer session;
} ZConnectPool;

typedef struct _ZProxyPool
{
  ZProxy *proxy;
  guint sessions;
  gchar *key;
} ZProxyPool;

void
z_tracker_sockaddr_copy(struct sockaddr_in *to, struct sockaddr_in *from)
{
  z_enter();
  to->sin_family      = from->sin_family;
  to->sin_port        = from->sin_port;
  to->sin_addr.s_addr = from->sin_addr.s_addr;
  z_leave();
}

static guint
inline int_to_str(gchar *to, guint size, guint num)
{
  guint i;

  i = 0;

  while(num > 0 && i < size - 2)
    {
      to[i] = (num % 10) + '0';
      num = num / 10;
      i++;
    }

  to[i] = ' ';
  i++;

  return i;
}

static void
z_tracker_pool_setkey(gchar *key, guint keylen, struct sockaddr_in *from, struct sockaddr_in *to, const gchar *key2)
{
  int i = 0;
  
  z_enter();
  
  i += int_to_str(&key[i], keylen - i, from->sin_addr.s_addr);
  i += int_to_str(&key[i], keylen - i, from->sin_port);
  i += int_to_str(&key[i], keylen - i, to->sin_addr.s_addr);
  i += int_to_str(&key[i], keylen - i, to->sin_port);
  key[i] = 0;
  
  if (key2)
    strncpy(&key[i], key2, keylen - i - 1);
  
  key[keylen-1] = 0;
  z_leave();
}

void
z_tracker_ppol_setkey(ZCProxy *self, ZConnectPool *value, gchar *key)
{
  struct sockaddr_in from;
  struct sockaddr_in to;
  z_enter();
  if(self->sess)
    {
      z_cp();
      strncpy(key, value->key, MAX_POOL_KEYSIZE);
      key[MAX_POOL_KEYSIZE-1] = 0;
      z_leave();
      return;
    }
  
  memset(&from, 0, sizeof(struct sockaddr_in));
  memset(&to,   0, sizeof(struct sockaddr_in));
  
  if(self->srcip)
    from.sin_addr.s_addr = value->from.sin_addr.s_addr;
  
  if(self->srcport)
    from.sin_port = value->from.sin_port;
  
  if(self->dstip)
    to.sin_addr.s_addr = value->to.sin_addr.s_addr;
  
  if(self->dstport)
    to.sin_port = value->to.sin_port;
  
  z_tracker_pool_setkey(key, MAX_POOL_KEYSIZE, &from, &to, NULL);
  z_leave();
}

void
z_tracker_ppool_inc(ZCProxy *self, ZConnectPool *value)
{
  gchar key[MAX_POOL_KEYSIZE];
  ZProxyPool *pool;
  z_enter();
  if (!self->proxypool)
    {
      z_leave();
      return;
    }
  z_tracker_ppol_setkey(self, value, key);
  pool = g_hash_table_lookup(self->proxypool, key);
  if (pool)
    pool->sessions++;
  else
    {
      pool = g_new0(ZProxyPool, 1);
      pool->proxy = value->proxy;
      pool->key = g_new0(gchar, strlen(key)+1);
      strcpy(pool->key, key);
      pool->sessions = 1;
      g_hash_table_insert(self->proxypool, pool->key, pool);
    }
  z_leave();
}

void
z_tracker_ppool_dec(ZCProxy *self, ZConnectPool *value)
{
  gchar key[MAX_POOL_KEYSIZE];
  ZProxyPool *pool;
  z_enter();
  if (!self->proxypool)
    {
      z_leave();
      return;
    }
  z_tracker_ppol_setkey(self, value, key);
  pool = g_hash_table_lookup(self->proxypool, key);
  if (!pool)
    {
      z_log(NULL, CORE_ERROR, 3, "Lost proxy. Internal error;");
      return;
    }
  pool->sessions--;
  if (pool->sessions <= 0)
    {
      if (value->proxy != pool->proxy)
        {
          z_log(NULL, CORE_ERROR, 3, "Internal error; value->proxy!=pool->proxy");
          z_leave();
        }
        
      (value->proxy->free)(value->proxy);
      g_hash_table_remove(self->proxypool, key);
      g_free(pool->key);
      g_free(pool);
    }
  z_leave();
}

ZProxy *
z_tracker_ppol_search(ZCProxy *self, ZConnectPool *value)
{
  gchar key[MAX_POOL_KEYSIZE];
  ZProxyPool *pool;
  z_enter();
  if (!self->proxypool)
    {
      z_leave();
      return NULL;
    }

  z_tracker_ppol_setkey(self, value, key);
  pool = g_hash_table_lookup(self->proxypool, key);
  if (pool)
    {
      z_leave();
      return pool->proxy;
    }
  
  z_leave();
  return NULL;
}

static ZProxy *
z_tracker_start_new_proxy(ZCProxy *self, ZConnectPool *value)
{
  ZPolicyObj *res;
  ZPolicyObj *arg;
  ZProxy *proxy = NULL;
  ZSockAddr *from, *to;
  gboolean called;

  z_enter();
  
  from = z_sockaddr_new((struct sockaddr *)&value->from, sizeof(struct sockaddr_in));
  to = z_sockaddr_new((struct sockaddr *)&value->to, sizeof(struct sockaddr_in));

  z_policy_lock(self->thread);
  arg = z_policy_var_build("(iOO)", value->fd, z_py_zorp_sockaddr_new(from),
                                            z_py_zorp_sockaddr_new(to));
  res = z_policy_call(self->handler, "startProxyInstance", arg, &called);

  if (res !=NULL && res != z_policy_none)
    proxy = ((ZorpProxy *)res)->proxy;

  z_sockaddr_unref(from);
  z_sockaddr_unref(to);
  z_policy_var_unref(res);
  z_policy_unlock(self->thread);
  z_leave();
  return proxy;
}

void
z_tracker_search_proxy(ZCProxy *self, ZConnectPool *value)
{
  z_enter();  
  value->proxy = z_tracker_ppol_search(self, value);
  
  if (!value->proxy)
    {
      value->proxy = z_tracker_start_new_proxy(self, value);
      if (value->proxy)
        {
          (value->proxy->connect_init)(value->proxy, self, self->other);
          z_tracker_ppool_inc(self, value);
        }
    }

  z_leave();
  return;
}

static guint
z_tracker_add_connect(ZCProxy *self, struct sockaddr_in *from,
                                     struct sockaddr_in *to)
{
  gint fd;
  gint socksize = sizeof(struct sockaddr_in);
  
  z_enter();
  fd=socket(AF_INET, SOCK_DGRAM, 0);
  if (bind(fd, from, sizeof(struct sockaddr_in)) < 0)
    {
      close(fd);
      z_leave();
      return -1;
    }
    
  if (connect(fd, to, sizeof(struct sockaddr_in)) < 0)
    {
      close(fd);
      z_leave();
      return -1;
    }

  z_conntrack_router_newfd_with_data(self->router, fd, self);

  getsockname(fd, from, &socksize);
  
  z_leave();
  return fd;
}


static void
z_tracker_remove(ZCProxy *self, ZConnectPool *value)
{
  z_enter();
  
  z_conntrack_router_remove_fd(self->router, value->fd);
  
  close(value->fd);
  
  g_hash_table_remove(self->trackpool, value->key);
  if (value->proxy)
    {
      (value->proxy->connect_stop)(value->proxy, self, value->session);
      z_tracker_ppool_dec(self, value);
    }
  g_free(value->key);
  value->key = NULL;
  g_free(value);
  z_leave();
}

void
z_tracker_stop_session(ZCProxy *self, gchar *key)
{
  ZConnectPool *value;
  GList *tmplist;
  
  z_enter();
  
  if (self == NULL || key == NULL)
    {
      z_leave();
      return;
    }
    
  value = g_hash_table_lookup(self->trackpool, key);
  if (!value)
    {
      z_leave();
      return;
    }
  z_conntrack_router_remove_fd(self->router, value->fd);
  
  tmplist = g_list_find(self->timer, value);
  if (tmplist == NULL)
    z_log(NULL, CORE_ERROR, 4, "Error! Lost connection!");
  else
    self->timer = g_list_remove_link(self->timer, tmplist);
  
  close(value->fd);
  
  g_hash_table_remove(self->trackpool, value->key);

  z_tracker_ppool_dec(self, value);

  g_free(value->key);
  value->key = NULL;
  g_free(value);
  z_leave();
}
                          

static gboolean
z_tracker_pool_checktime(ZCProxy *self, ZConnectPool *value, guint remove)
{
  time_t now;
  GList *tmplist;
  
  z_enter();
  
  if (!value)
    {
      tmplist = g_list_last(self->timer);
      if (g_list_previous(tmplist) != NULL)
        tmplist = g_list_previous(tmplist);
      value = tmplist->data;
    }
  else
    {
      tmplist = g_list_find(self->timer, value);
      if (tmplist == NULL)
        {
          z_log(NULL, "plug.tracking", 1, "Error! Lost connection!");
          z_leave();
          return TRUE;
        }
    }

  if (!value)
    {
      z_leave();
      return FALSE;
    }

  self->timer = g_list_remove_link(self->timer, tmplist);
  
  now = time(NULL);
  if (remove)
    {
      if((value->time + self->timeout) < now)
        {
          z_tracker_remove(self, value);
          z_leave();
          return TRUE;
        }
    }
  
  value->time = now;
  self->timer = g_list_prepend(self->timer, value);

  z_leave();
  return FALSE;
}

static gboolean
z_tracker_pool_insert(ZCProxy *self, ZConnectPool *value)
{
  time_t now;
  z_enter();

  z_tracker_pool_checktime(self, NULL, 1);
  
  now = time(NULL);
  value->time = now;
  g_hash_table_insert(self->trackpool, value->key, value);

  self->timer = g_list_prepend(self->timer, value);

  z_leave();
  return TRUE;
}

static ZConnectPool *
z_tracker_load_new_proxy(ZCProxy *self, gchar *oldkey, int fd, ZPacket *packet)
{
  ZConnectPool *new = g_new0(ZConnectPool, 1);

  z_enter();

  new->key = g_new0(gchar, strlen(oldkey)+1);
  strcpy(new->key, oldkey);

  if (!z_tracker_pool_insert(self, new))
    {
      g_free(new);
      z_leave();
      return NULL;
    }
  
  new->fd = fd;
  
  z_tracker_sockaddr_copy(&new->to, &packet->packheader.from);
  z_tracker_sockaddr_copy(&new->from, &packet->packheader.to);
  
  z_tracker_search_proxy(self, new);

  if (new->proxy)
    {
      new->session = (new->proxy->connect_setup)(new->proxy, new->key);
      z_tracker_ppool_inc(self, new);
    }

  z_leave();
  return new;
}

static ZConnectPool *
z_tracker_pool_get(ZCProxy *self, gchar *key, guint remove)
{
  ZConnectPool *value;
  
  z_enter();
  
  value = g_hash_table_lookup(self->trackpool, key);
  
  if (value)
    {
      if (z_tracker_pool_checktime(self, value, remove))
        {
          value = NULL;
        }
    }
  
  z_leave();
  return value;
}

static ZConnectPool *
z_tracker_new_sender(ZCProxy *self, ZPacket *packet, ZProxy *proxy, gpointer session)
{
  ZConnectPool *new = g_new0(ZConnectPool, 1);

  gchar oldkey[MAX_POOL_KEYSIZE];
  gchar prot_key[MAX_POOL_KEYSIZE];
    
  z_enter();
  
  memset(prot_key, 0, MAX_POOL_KEYSIZE);
  memset(oldkey, 0, MAX_POOL_KEYSIZE);
  if (self->keygen)
    {
      (self->keygen)(self, packet, prot_key, MAX_POOL_KEYSIZE);
    }
  
  new->fd = z_tracker_add_connect(self, &packet->packheader.from, &packet->packheader.to);  
  
  z_tracker_pool_setkey(oldkey, MAX_POOL_KEYSIZE, &packet->packheader.to,
                           &packet->packheader.from, prot_key);

  new->key = g_new0(gchar, strlen(oldkey)+1);
  strcpy(new->key, oldkey);

  z_tracker_sockaddr_copy(&new->from, &packet->packheader.from);
  z_tracker_sockaddr_copy(&new->to, &packet->packheader.to);

  if (!z_tracker_pool_insert(self, new))
    {
      g_free(new);
      z_leave();
      return NULL;
    }

  new->proxy = proxy;
  new->session = session;

  z_tracker_ppool_inc(self, new);

  z_leave();
  return new;
}

static ZConnectPool *
z_tracker_search_bykey(ZCProxy *self, gchar *key)
{
  ZConnectPool *value;

  z_enter();
  
  value = z_tracker_pool_get(self, key, 0);
  
  z_leave();
  return value;
}

static ZProxy *
z_tracker_search(ZCProxy *self, ZPacket *packet, int fd, gpointer *session)
{
  ZConnectPool *value;
  gchar key[MAX_POOL_KEYSIZE];
  gchar prot_key[MAX_POOL_KEYSIZE];

  z_enter();
  
  memset(prot_key, 0, MAX_POOL_KEYSIZE);
  memset(key, 0, MAX_POOL_KEYSIZE);  
  if (self->keygen)
    {
      (self->keygen)(self, packet, prot_key, MAX_POOL_KEYSIZE);
    }
  z_tracker_pool_setkey(key, MAX_POOL_KEYSIZE, &packet->packheader.from,
                                        &packet->packheader.to, prot_key);
  
  value = z_tracker_pool_get(self, key, 1);
  
  if (!value)
    {
      if (self->way == 0)
        {
          value = z_tracker_load_new_proxy(self, key, -1, packet);
        }
      else
        {
          z_leave();
          return NULL;
        }
    }
  
  *session = value->session;
  
  z_leave();
  return value->proxy;
}

// Innentol vannak az exportalt fuggvenyek.

void
z_tracker_init(ZCProxy *self, ZPolicyObj *handler, ZCTRouter *router)
{
  z_enter();
  
  self->handler = handler;
  z_policy_var_ref(handler);
  
  self->router = router;
  strcpy(self->session_id, "Fake session");
  self->setattr = z_pyvars_setattr;
  self->getattr = z_pyvars_getattr;
  self->vars = z_pyvars_new();
  self->thread = PyThreadState_New(z_python_global_state->interp);
  
  self->trackpool = g_hash_table_new(g_str_hash, g_str_equal);
  
  self->proxypool = g_hash_table_new(g_str_hash, g_str_equal);
  
  self->timeout = 30000;
  self->max_packet = -1;
  self->timer = g_list_alloc();
  self->way = 0;
  
  self->srcip =   FALSE;
  self->srcport = FALSE;
  self->dstip =   FALSE;
  self->dstport = FALSE;
  self->sess =    FALSE;

  z_proxy_var_new((ZProxy *)self, "timeout",          Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->timeout);
  z_proxy_var_new((ZProxy *)self, "maxpacket",        Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->max_packet);
  z_proxy_var_new((ZProxy *)self, "source_ip",        Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->srcip);
  z_proxy_var_new((ZProxy *)self, "source_port",      Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->srcport);
  z_proxy_var_new((ZProxy *)self, "destination_ip",   Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->dstip);
  z_proxy_var_new((ZProxy *)self, "destination_port", Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->dstport);
  z_proxy_var_new((ZProxy *)self, "all_session",      Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->sess);
  
  self->sender_fd = socket(AF_INET, SOCK_DGRAM, 0);

  z_conntrack_router_sethandler(self->router, self);
    
  z_leave();
}

void
z_tracker_sender_init(ZCProxy *self, ZCProxy *old)
{
  z_enter();
  
  self->handler = old->handler;
  z_policy_var_ref(old->handler);
  self->router = old->router;
  strcpy(self->session_id, "Fake sender session");
  self->setattr = z_pyvars_setattr;
  self->getattr = z_pyvars_getattr;
  self->vars = old->vars;
  self->thread = old->thread;
  
  self->trackpool = g_hash_table_new(g_str_hash, g_str_equal);
  
  self->proxypool = NULL;

  self->timeout = old->timeout;
  self->max_packet = old->max_packet;
  self->timer = g_list_alloc();
  self->way = 1;

  self->sender_fd = socket(AF_INET, SOCK_DGRAM, 0);
  
  z_leave();
}

void
z_tracker_destroy(ZCProxy *self)
{
  ZPolicyObj *handler;
  ZPolicyThread *thread;
  
  z_enter();
  if (self->way == 0)
    {
      z_policy_lock(self->thread);
      thread = self->thread;
      self->thread = NULL;

      z_pyvars_destroy(self->vars);
      self->vars = NULL;
        
      handler = self->handler;
      self->handler = NULL;
      z_policy_var_unref(handler);

      z_policy_unlock(thread);
  
      z_python_lock();
      PyThreadState_Clear(thread);
      z_python_unlock();
    }
  z_leave();
}

void
z_tracker_tune(ZCProxy *self, ZCProxy *other, ZCKeygen keygen)
{
  self->other = other;
  self->keygen = keygen;
}

gboolean
z_tracker_config_event(ZCProxy *self)
{
  ZPolicyObj *res;
  gboolean back = TRUE;
  gboolean called;
    
  z_enter();
  res = z_policy_call(self->handler, "config", NULL, &called);
  if (res == NULL)
    {
      if (called)
        {
          back = FALSE;
        }
    }
  z_policy_var_unref(res);
  z_leave();
  return back;
}

void
z_tracker_read(ZCProxy *self, ZPacket *packet, int fd)
{
  ZProxy *proxy;
  gpointer session;
  z_enter();
  
  proxy = z_tracker_search(self, packet, fd, &session);

  if (proxy)
    (proxy->packet_read)(proxy, packet, self->way, session);
  else
    z_log(NULL, CORE_DEBUG, 4, "UDP packet dropped. Reason: Policy deny or cannot load proxy.");
  
  z_leave();
}

gchar *
z_tracker_write(ZCProxy *self, ZPacket *packet, ZProxy *proxy, gchar *key, gpointer session)
{
  ZConnectPool *value;
  z_enter();
  
  if (key)
    value = z_tracker_search_bykey(self, key);
  else
    value = z_tracker_new_sender(self, packet, proxy, session);
  
  if (value)
    z_packsock_send(packet, self->sender_fd, &value->from, &value->to);
  
  g_free(packet);
  z_leave();
  return value->key;
}

void
z_tracker_drop(ZCProxy *self, ZPacket *packet)
{
  g_free(packet->data);
  g_free(packet);
}

gboolean
z_tracker_connect_server(ZCProxy *self, ZorpCHeader *header)
{
  ZSockAddr *to;
  ZSockAddr *newto;
  ZPolicyObj *res, *args, *ret1, *ret2;
  gboolean ret = FALSE;
  z_enter();
  
  to = z_sockaddr_new((struct sockaddr *)&header->to, sizeof(struct sockaddr_in));
  z_policy_lock(self->thread);
  args = z_policy_var_build("(O)", z_py_zorp_sockaddr_new(to));
  res = z_policy_call(self->handler, "connectServer", args, NULL);
  if (res)
    {
      if(z_policy_var_parse(res, "(OO)", &ret1, &ret2))
        {
          if (z_py_zorp_sockaddr_check(ret1))
            {
              newto = ((ZorpSockAddr *)ret1)->sa;
              z_tracker_sockaddr_copy(&header->from, (struct sockaddr_in *)&newto->sa);
            }
          else
            {
              header->from.sin_addr.s_addr = 0;
              header->from.sin_port = 0;
            }
          if (z_py_zorp_sockaddr_check(ret2))
            {
              newto = ((ZorpSockAddr *)ret2)->sa;
              z_tracker_sockaddr_copy(&header->to, (struct sockaddr_in *)&newto->sa);
              ret = TRUE;
            }
        }
    }
  z_sockaddr_unref(to);
  z_policy_var_unref(res);
  z_policy_unlock(self->thread);
  z_leave();
  return(ret);
}

#endif
