/***************************************************************************
 *
 * Copyright (c) 2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009,
 * 2010 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 version 2 as published
 * by the Free Software Foundation.
 *
 * Note that this permission is granted for only version 2 of the GPL.
 *
 * As an additional exemption you are allowed to compile & link against the
 * OpenSSL libraries as published by the OpenSSL project. See the file
 * COPYING for details.
 *
 * 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: pssl.c,v 1.51 2004/07/29 10:55:48 bazsi Exp $
 *
 * Author: Attila SZALAY <sasa@balabit.hu>
 * Auditor:
 * Last audited version:
 * Notes:
 *
 ***************************************************************************/

#include <zorp/zorp.h>
#include <zorp/stream.h>
#include <zorp/proxy.h>
#include <zorp/poll.h>
#include <zorp/policy.h>
#include <zorp/thread.h>
#include <zorp/log.h>
#include <zorp/registry.h>
#include <zorp/sockaddr.h>
#include <zorp/io.h>
#include <zorp/streamssl.h>
#include <zorp/ssl.h>
#include <zorp/plugsession.h>

#include <sys/types.h>
#include <sys/time.h>
#include <unistd.h>
#include <errno.h>
#include <openssl/err.h>


#define PROXY_BUFSIZE    1500 /* 64k */

#define PSSL_DEBUG "pssl.debug"
#define PSSL_ERROR "pssl.error"
#define PSSL_POLICY "pssl.policy"

typedef struct _PsslProxy
{
  ZProxy super;
  ZPoll *poll;
  ZPlugSessionData session_data;

  GString *client_keyfile, *client_certfile;
  GString *server_keyfile, *server_certfile;
  GString *client_ca_dir;
  GString *client_crl_dir;
  GString *server_ca_dir;
  GString *server_crl_dir;
  gboolean need_ssl[EP_MAX];
  gint handshake_timeout;
  gint verify_type[EP_MAX];
  gint verify_depth[EP_MAX];
} PsslProxy;

extern ZClass PsslProxy__class;

#define ENABLE_PKT_STATS 1

static gboolean
pssl_packet_stat_event(ZPlugSession *session, 
                       guint64 client_bytes, guint64 client_pkts, 
                       guint64 server_bytes, guint64 server_pkts,
                       gpointer user_data);

static void pssl_finish(ZPlugSession *session, gpointer user_data);


static void
pssl_config_set_defaults(PsslProxy *self)
{
  gint i;
  
  z_proxy_enter(self);
  self->handshake_timeout = 30000;
  self->session_data.timeout = 600000;
  self->session_data.copy_to_server = TRUE;
  self->session_data.copy_to_client = TRUE;
  self->session_data.buffer_size = 1500;
  self->session_data.shutdown_soft = TRUE;
  self->session_data.packet_stats = pssl_packet_stat_event;
  self->session_data.finish = pssl_finish;
          
  self->poll = z_poll_new();
  
  self->client_keyfile = g_string_sized_new(32);
  self->client_certfile = g_string_sized_new(32);

  self->client_ca_dir = g_string_sized_new(0);
  self->client_crl_dir = g_string_sized_new(0);

  self->server_ca_dir = g_string_sized_new(32);
  self->server_crl_dir = g_string_sized_new(32);

  self->server_keyfile = g_string_sized_new(0);
  self->server_certfile = g_string_sized_new(0);
  
  self->verify_type[EP_SERVER] = Z_SSL_VERIFY_REQUIRED_TRUSTED;
  self->verify_type[EP_CLIENT] = Z_SSL_VERIFY_REQUIRED_TRUSTED;
  self->verify_depth[EP_SERVER] = 1;
  self->verify_depth[EP_CLIENT] = 1;
  for (i = 0; i < EP_MAX; i++)
    self->need_ssl[i] = TRUE;
  z_proxy_return(self);
}

static void
pssl_register_vars(PsslProxy *self)
{
  z_proxy_enter(self);
  z_proxy_var_new(&self->super,
                  "timeout",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->session_data.timeout);

  z_proxy_var_new(&self->super,
                  "handshake_timeout",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->handshake_timeout);

  z_proxy_var_new(&self->super,
                  "copy_to_client",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->session_data.copy_to_client);

  z_proxy_var_new(&self->super,
                  "copy_to_server",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->session_data.copy_to_server);

  z_proxy_var_new(&self->super,
                  "shutdown_soft",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->session_data.shutdown_soft);

  z_proxy_var_new(&self->super,
                  "packet_stats_interval_packet",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->session_data.packet_stats_interval_packet);
                  
  z_proxy_var_new(&self->super,
                  "packet_stats_interval",
                  Z_VAR_TYPE_ALIAS | Z_VAR_GET | Z_VAR_GET_CONFIG,
                  "packet_stats_interval_packet");

  z_proxy_var_new(&self->super,
                  "packet_stats_interval_time",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->session_data.packet_stats_interval_time);

  z_proxy_var_new(&self->super,
                  "buffer_size",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->session_data.buffer_size);

		  
  z_proxy_var_new(&self->super,
                  "client_need_ssl",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
		  &self->need_ssl[EP_CLIENT]);
		  
  z_proxy_var_new(&self->super,
                  "client_key_file",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_STRING,
		  self->client_keyfile);
		  
  z_proxy_var_new(&self->super,
                  "client_cert_file",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_STRING,
		  self->client_certfile);
		  
  z_proxy_var_new(&self->super,
                  "client_ca_directory",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_STRING,
                  self->client_ca_dir);
                  
  z_proxy_var_new(&self->super,
                  "client_crl_directory",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_STRING,
                  self->client_crl_dir);
                  
  z_proxy_var_new(&self->super,
                  "client_verify_type",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->verify_type[EP_CLIENT]);
                  
  z_proxy_var_new(&self->super,
                  "client_verify_depth",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->verify_depth[EP_CLIENT]);


  z_proxy_var_new(&self->super,
                  "server_need_ssl",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
		  &self->need_ssl[EP_SERVER]);
		  
  z_proxy_var_new(&self->super,
                  "server_key_file",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_STRING,
		  self->server_keyfile);
		  
  z_proxy_var_new(&self->super,
                  "server_cert_file",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_STRING,
		  self->server_certfile);
		  
  z_proxy_var_new(&self->super,
                  "server_ca_directory",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_STRING,
                  self->server_ca_dir);
                  
  z_proxy_var_new(&self->super,
                  "server_crl_directory",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_STRING,
                  self->server_crl_dir);
                  
  z_proxy_var_new(&self->super,
                  "server_verify_type",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->verify_type[EP_SERVER]);
                  
  z_proxy_var_new(&self->super,
                  "server_verify_depth",
                  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT,
                  &self->verify_depth[EP_SERVER]);

  /* compatibility aliases with Zorp 0.8.x */
  z_proxy_var_new(&self->super,
                  "client_key",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_ALIAS,
		  "client_key_file");
		  
  z_proxy_var_new(&self->super,
                  "client_cert",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_ALIAS,
		  "client_cert_file");
		  
  z_proxy_var_new(&self->super,
                  "server_key",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_ALIAS,
		  "server_key_file");
		  
  z_proxy_var_new(&self->super,
                  "server_cert",
		  Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_ALIAS,
		  "server_cert_file");
  z_proxy_return(self);
}

static gboolean
pssl_packet_stat_event(ZPlugSession *session G_GNUC_UNUSED, 
                       guint64 client_bytes, guint64 client_pkts, 
                       guint64 server_bytes, guint64 server_pkts,
                       gpointer user_data)
{
  PsslProxy *self = (PsslProxy *)user_data;
  ZPolicyObj *res;
  gboolean called;
  guint resc;

  z_policy_lock(self->super.thread);
  res = z_policy_call(self->super.handler, "packetStats",
                      z_policy_var_build("iiii",
                                         (guint32) client_bytes,
                                         (guint32) client_pkts,
                                         (guint32) server_bytes,
                                         (guint32) server_pkts),
                      &called,
                      self->super.session_id);
  if (called)
    {
      resc = Z_REJECT;
      if (res)
        {
          if (!z_policy_var_parse(res, "i", &resc))
            z_proxy_log(self, PSSL_POLICY, 1, "Return value of packetStats() is not int;");
          else if (resc == Z_REJECT)
            z_proxy_log(self, PSSL_DEBUG, 6, "packetStats() return Z_REJECT, aborting session;");
        }
      else
        {
          z_proxy_log(self, PSSL_POLICY, 1, "Error during packetStats() event;");
        }
    }
  else
    {
      resc = Z_ACCEPT;
    }
  z_policy_var_unref(res);
  z_policy_unlock(self->super.thread);
  return resc == Z_ACCEPT;
}
    
static void
pssl_finish(ZPlugSession *session G_GNUC_UNUSED, gpointer user_data)
{
  PsslProxy *self = (PsslProxy *)user_data;
  
  z_poll_quit(self->poll);
}

guint
pssl_perform_handshake(PsslProxy  *self, gint side)
{
  ZSSLSession *ssl;
  X509 *peercert;
  int ret;
  ZStream *tmpstream;
  
  z_proxy_enter(self);
  ssl = z_ssl_session_new(self->super.session_id, 
                          side == EP_CLIENT ? Z_SSL_MODE_SERVER : Z_SSL_MODE_CLIENT,
                          self->client_keyfile->str,
                          self->client_certfile->str,
                          self->client_ca_dir->str,
                          self->client_crl_dir->str,
                          self->verify_depth[EP_CLIENT],
                          self->verify_type[EP_CLIENT]);

  if (!ssl)
    {
      z_proxy_log(self, PSSL_ERROR, 1, "Error initializing SSL session; side='%s'", EP_STR(side));
      z_proxy_return(self, FALSE);
    }
  SSL_set_options(ssl->ssl, SSL_MODE_ENABLE_PARTIAL_WRITE);

  z_stream_set_timeout(self->super.endpoints[side], self->handshake_timeout);
  tmpstream = self->super.endpoints[side];
  self->super.endpoints[side] = z_stream_ssl_new(tmpstream, ssl);
  z_stream_unref(tmpstream);

  if (side == EP_CLIENT)
    ret = SSL_accept(ssl->ssl);
  else
    ret = SSL_connect(ssl->ssl);

  if (ret <= 0)
    {
      char buf[1024];
      
      z_ssl_session_unref(ssl);
      z_proxy_log(self, PSSL_ERROR, 1, "SSL handshake failed; side='%s', error='%s'", EP_STR(side), z_ssl_get_error_str(buf, sizeof(buf)));
      z_proxy_return(self, FALSE);
    }
  peercert = SSL_get_peer_certificate(ssl->ssl);
  if (peercert)
    {
      gchar tmp[1024];
      
      X509_NAME_oneline(X509_get_subject_name(peercert), tmp, sizeof(tmp));
      X509_free(peercert);
      z_proxy_log(self, PSSL_DEBUG, 4, "Identified peer; side='%s', peer='%s'", EP_STR(side), tmp);
    }
    
  z_ssl_session_unref(ssl);  
  z_stream_set_timeout(self->super.endpoints[side], -2);
  z_proxy_log(self, PSSL_DEBUG, 6, "SSL handshake successful; side='%s'", EP_STR(side));
  z_proxy_return(self, TRUE);

}

static gboolean
pssl_init_streams(PsslProxy *self)
{
  ZStream *tmpstream;
  gint ep;

  z_proxy_enter(self);
  if (!self->super.endpoints[EP_CLIENT] || 
      !self->super.endpoints[EP_SERVER] ||
      !self->poll)
    z_proxy_return(self, FALSE);

  for (ep = EP_CLIENT; ep < EP_MAX; ep++)
    {
      tmpstream = self->super.endpoints[ep];
      z_stream_ref(tmpstream);
      if (self->need_ssl[ep])
        {
          z_proxy_log(self, PSSL_DEBUG, 6, "Performing SSL handshake; side='%s'", EP_STR(ep));
          if (!pssl_perform_handshake(self, ep))
            {
              z_stream_unref(tmpstream);
              z_proxy_return(self, FALSE);
            }
        }
      z_stream_set_nonblock(tmpstream, TRUE);
      z_stream_unref(tmpstream);
    }
  z_proxy_return(self, TRUE);
}

static gboolean
pssl_request_stack_event(PsslProxy *self, ZStackedProxy **stacked)
{
  ZPolicyObj *res;
  gboolean called;
  gboolean rc = TRUE;

  z_proxy_enter(self);
  z_policy_lock(self->super.thread);
  *stacked = NULL;
  res = z_policy_call(self->super.handler,
                      "requestStack",
                      NULL,
                      &called,
                      self->super.session_id);
  if (res)
    {
      if (res != z_policy_none)
        rc = z_proxy_stack_object(&self->super, res, stacked);
    }
  else if (called)
    {
      rc = FALSE;
    }
  z_policy_var_unref(res);
  z_policy_unlock(self->super.thread);
  z_proxy_return(self, rc);
}

static gboolean
pssl_config(ZProxy *s)
{
  PsslProxy *self = Z_CAST(s, PsslProxy);
  gboolean success;

  z_proxy_enter(self);
  pssl_config_set_defaults(self);
  pssl_register_vars(self);
  success = Z_SUPER(s, ZProxy)->config(s);
  z_proxy_return(self, success);
}

static void
pssl_main(ZProxy *s)
{
  PsslProxy *self = Z_CAST(s, PsslProxy);
  ZStackedProxy *stacked;
  ZPlugSession *session = NULL;

  z_proxy_enter(self);
  /* this sets the server side endpoint if successful */
  if (!z_proxy_connect_server(&self->super, NULL, 0) ||
      !pssl_init_streams(self))
    goto error;
  
  if (!pssl_request_stack_event(self, &stacked))
    goto error;
    
  session = z_plug_session_new(&self->session_data, self->super.endpoints[EP_CLIENT], self->super.endpoints[EP_SERVER], stacked, (gpointer)&self->super);
  if (!session)
    {
      z_stacked_proxy_destroy(stacked);
      goto error;
    }

  if (!z_plug_session_start(session, self->poll))
    {
      goto error;
    }
    
  while (z_poll_is_running(self->poll) && z_poll_iter_timeout(self->poll, -1))
    {
      if (!z_proxy_loop_iteration(s))
        {
          break;
        }
    }
 error:
  if (session)
    {
      z_plug_session_cancel(session);
      z_plug_session_unref(session);
    }
  z_proxy_leave(self);
}

static ZProxy *
pssl_proxy_new(ZProxyParams *params)
{
  PsslProxy *self;

  z_enter();
  self = Z_CAST(z_proxy_new(Z_CLASS(PsslProxy), params), PsslProxy);
  z_return(&self->super);
}

static void
pssl_proxy_free(ZObject *s)
{
  PsslProxy *self = Z_CAST(s, PsslProxy);

  z_enter();
  if (self->poll)
    {
      z_poll_unref(self->poll);
      self->poll = NULL;
    }
  ERR_remove_state(0);
  z_proxy_free_method(s);
  z_return();
}

ZProxyFuncs pssl_proxy_funcs =
{
  {
    Z_FUNCS_COUNT(ZProxy),
    pssl_proxy_free,
  },
  .config = pssl_config,
  .main = pssl_main,
  NULL
};


ZClass PsslProxy__class =
{
  Z_CLASS_HEADER,
  &ZProxy__class,
  "PsslProxy",
  sizeof(PsslProxy),
  &pssl_proxy_funcs.super
};

gint
zorp_module_init(void)
{
  z_registry_add("pssl", ZR_PROXY, pssl_proxy_new);
  return TRUE;
}
