/* TODO: Include attribution for ideas, and code from mod_ntlm */

#include "usuals.h"
#include "httpauthd.h"
#include "hash.h"
#include "defaults.h"
#include "md5.h"
#include "basic.h"

#include <syslog.h>

/* The NTLM headers */
#include "ntlmssp.h"

/* -------------------------------------------------------------------------------
 * Defaults and Constants
 */

#define NTLM_HASH_KEY_LEN MD5_LEN
#define NTLM_ESTABLISHED  (void*)1


/* -------------------------------------------------------------------------------
 * Structures and Globals
 */

/* A pending connection */
typedef struct ntlm_connection
{
  void* handle;
  char nonce[NONCE_LEN];
  unsigned int flags;
}
ntlm_connection_t;


/* The main context */
typedef struct ntlm_context
{
  /* Settings ---------------------------------------------------------- */
  const char* server;       /* Server to authenticate against */
  const char* domain;       /* NTLM domain to authenticate against */
  const char* backup;       /* Backup server if primary is down */
  int pending_max;          /* Maximum number of connections at once */
  int pending_timeout;      /* Timeout for authentication (in seconds) */

  /* Context ----------------------------------------------------------- */
  hash_t* pending;        /* Pending connections */
  hash_t* established;    /* Established connections */
}
ntlm_context_t;


/* The default context settings */
static const ntlm_context_t ntlm_defaults =
{
  NULL, NULL, NULL, DEFAULT_PENDING_MAX, DEFAULT_PENDING_TIMEOUT,
  NULL, NULL
};


/* Mutexes for accessing non-thread-safe smblib */
static pthread_mutex_t g_smblib_mutex;
static pthread_mutexattr_t g_smblib_mutexattr;


/* -------------------------------------------------------------------------------
 * Internal Functions
 */

static ntlm_connection_t* makeconnection(ntlm_context_t* ctx)
{
  ntlm_connection_t* conn;

  ASSERT(ctx);

  conn = (ntlm_connection_t*)malloc(sizeof(ntlm_connection_t));
  if(!conn)
  {
    ha_messagex(LOG_CRIT, "out of memory");
    return NULL;
  }

  memset(conn, 0, sizeof(*conn));

  /*
   * Open a connection to to the domain controller. I don't think
   * we can cache these connections or use them again as opening
   * a connection here results in an nonce being generated.
   */
  conn->handle = ntlmssp_connect(ctx->server, ctx->backup,
                                 ctx->domain, conn->nonce);
  if(!conn->handle)
  {
    ha_messagex(LOG_ERR, "ntlm: couldn't connect to the domain server %s (backup: %s)",
             ctx->server, ctx->backup ? ctx->backup : "none");
    free(conn);
    return NULL;
  }

  ha_messagex(LOG_INFO, "ntlm: established connection to server");
  return conn;
}

static void freeconnection(ntlm_connection_t* conn)
{
  ASSERT(conn);

  if(conn->handle)
  {
    ha_messagex(LOG_DEBUG, "ntlm: disconnected from server");
    ntlmssp_disconnect(conn->handle);
    conn->handle = NULL;
  }

  free(conn);
}

static void free_hash_object(void* arg, void* val)
{
  if(val)
  {
    ASSERT(val != NTLM_ESTABLISHED);
    freeconnection((ntlm_connection_t*)val);
  }
}

static ntlm_connection_t* getpending(ntlm_context_t* ctx, const void* key)
{
  ntlm_connection_t* ret;

  ASSERT(ctx && key);

  ha_lock(NULL);

    ret = (ntlm_connection_t*)hash_rem(ctx->pending, key);

  ha_unlock(NULL);

  return ret;
}

static int putpending(ntlm_context_t* ctx, const void* key, ntlm_connection_t* conn)
{
  int r = 0;

  ASSERT(ctx && key && conn);
  ASSERT(conn->handle);

  if(!hash_get(ctx->pending, key))
  {
    ha_lock(NULL);

      if(!hash_set(ctx->pending, key, (void*)conn))
      {
        free_hash_object(NULL, conn);
        ha_messagex(LOG_ERR, "out of memory");
        r = -1;
      }

    ha_unlock(NULL);
  }

  return r;
}

int ntlm_auth_basic(ntlm_context_t* ctx, char* key, const char* header,
                    ha_response_t* resp, ha_buffer_t* buf)
{
  ntlm_connection_t* conn;
  char* t;
  basic_header_t basic;
  const char* domain = NULL;
  int found = 0;
  int r;

  ASSERT(ctx && key && header && resp && buf);

  /*
   * We're doing basic authentication on the connection
   * which invalidates any NTLM authentication we've started
   * or done on this connection.
   */
  conn = getpending(ctx, key);
  if(conn)
  {
    ha_messagex(LOG_WARNING, "ntlm: basic auth killed a pending ntlm auth in progress");
    freeconnection(conn);
  }

  if((r = basic_parse(header, buf, &basic)) < 0)
    return r;

  /* Check and see if this connection is in the cache */
  ha_lock(NULL);

    if(hash_get(ctx->established, basic.key) == NTLM_ESTABLISHED)
      found = 1;

  ha_unlock(NULL);

  if(found)
    ha_messagex(LOG_NOTICE, "ntlm: validated basic user against cache: %s", basic.user);

  else
  {
    /* Try to find a domain in the user */
    if((t = strchr(basic.user, '\\')) != NULL ||
       (t = strchr(basic.user, '/')) != NULL)
    {
      /* Break at the domain */
      domain = basic.user;
      basic.user = t + 1;
      *t = 0;

      /* Make sure this is our domain */
      if(strcasecmp(domain, ctx->domain) != 0)
         domain = NULL;
    }

    if(!domain)
    {
      /* Use the default domain if none specified */
      domain = ctx->domain;
    }

    /* Make sure above did not fail */
    if(basic.user && basic.user[0] && basic.password &&
       domain && domain[0])
    {
      ha_messagex(LOG_DEBUG, "ntlm: checking user against server: %s", basic.user);

      /* We need to lock to go into smblib */
      ha_lock(&g_smblib_mutex);

        /* Found in smbval/valid.h */
        if(ntlmssp_validuser(basic.user, basic.password, ctx->server,
                             ctx->backup, domain) == NTV_NO_ERROR)
        {
           /* If valid then we return */
           found = 1;
        }

      ha_unlock(&g_smblib_mutex);
    }

    if(found)
      ha_messagex(LOG_NOTICE, "ntlm: validated basic user against server: %s", basic.user);
  }

  if(found)
  {
    int r;
    resp->code = HA_SERVER_OK;
    resp->detail = basic.user;

    ha_lock(NULL);

      /* We put this connection into the successful connections */
      r = hash_set(ctx->established, basic.key, NTLM_ESTABLISHED);

    ha_unlock(NULL);

    if(!r)
    {
      ha_messagex(LOG_CRIT, "out of memory");
      return HA_CRITERROR;
    }

    return HA_OK;
  }

  return HA_FALSE;
}

int ntlm_auth_ntlm(ntlm_context_t* ctx, void* key, const char* header,
                   ha_response_t* resp, ha_buffer_t* buf)
{
  ntlmssp_info_rec ntlmssp;
  ntlm_connection_t* conn = NULL;
  unsigned int flags = 0;
  int ret = HA_FALSE;
  size_t len = 0;
  void* d;
  int r;

  ASSERT(ctx && key && header && resp && buf);

  /*
   * Retrieve and remove the connection from the pending bag.
   * We add it back again below if that's necessary.
   */
  conn = getpending(ctx, key);

  /*
   * We use the flags from an already established connection
   * if we've been pending and stuff
   */

  if(conn && conn->flags)
    flags = conn->flags;

  /*
   * First we figure out what kind of message the client
   * is sending us.
   */

  d = ha_bufdec64(buf, header, &len);

  if(!d || len == 0)
    goto finally;

  r = ntlmssp_decode_msg(&ntlmssp, d, len, &flags);
  if(r != 0)
  {
    ha_messagex(LOG_WARNING, "ntlm: decoding NTLMSSP message failed (error %d)", r);
    resp->code = HA_SERVER_BADREQ;
    goto finally;
  }


  switch(ntlmssp.msg_type)
  {

  /* An initial NTLM request? */
  case 1:
    {
      /* Win9x doesn't seem to send a domain or host */
      int win9x = !ntlmssp.host[0] && !ntlmssp.domain[0];

      /*
       * If we already have a connection to the domain controller
       * then we're in trouble. Basically this is the second
       * type 1 message we've received over this connection.
       *
       * TODO: Eventually what we want to do here is wait for the
       * other authentication request to complete, or something
       * like that.
       */
      if(conn)
      {
        /*
         * In this case we also add the connection back into the
         * pending stack so that the correct request will complete
         * properly when it comes through.
         */
        r = putpending(ctx, key, conn);
        conn = NULL;

        if(r < 0)
        {
          ret = HA_CRITERROR;
        }
        else
        {
          ha_messagex(LOG_ERR, "ntlm: received out of order NTLM request from client");
          resp->code = HA_SERVER_BADREQ;
        }

        goto finally;
      }


      /*
       * Check how many connections we have to the domain controller
       * and if too many then cut off here.
       */
      if(ctx->pending_max != -1)
      {
        ha_lock(NULL);

          if(hash_count(ctx->pending) >= ctx->pending_max)
            hash_bump(ctx->pending);

        ha_unlock(NULL);
      }


      /*
       * Open a connection to to the domain controller. I don't think
       * we can cache these connections or use them again as opening
       * a connection here results in an nonce being generated.
       */
      conn = makeconnection(ctx);

      if(!conn)
      {
        ret = HA_FAILED;
        goto finally;
      }

      /* Save away any flags given us by ntlm_decode_msg */
      conn->flags = flags;

      /* Start building the header */
      ha_bufcpy(buf, HA_PREFIX_NTLM);

      if(win9x)
      {
        struct ntlm_msg2_win9x msg_win9x;
        ntlmssp_encode_msg2_win9x(conn->nonce, &msg_win9x, (char*)ctx->domain, flags);
        ha_bufjoin(buf);
        ha_bufenc64(buf, (unsigned char*)&msg_win9x, sizeof(msg_win9x));
      }
      else
      {
        struct ntlm_msg2 msg;
        ntlmssp_encode_msg2(conn->nonce, &msg);
        ha_bufjoin(buf);
        ha_bufenc64(buf, (unsigned char*)&msg, sizeof(msg));
      }

      if(ha_buferr(buf))
        goto finally;

      /*
       * TODO: Our callers need to be able to keep alive
       * connections that have authentication going on.
       */

      /* Cache this connection in our pending set ... */
      r = putpending(ctx, key, conn);

      /*
       * By marking this as null, the cleanup code
       * won't free the connection since it's been
       * cached above.
       */
      conn = NULL;

      if(r < 0)
      {
        ret = HA_CRITERROR;
      }
      else
      {
        ha_messagex(LOG_DEBUG, "ntlm: sending ntlm challenge");
        ha_addheader(resp, "WWW-Authenticate", ha_bufdata(buf));
        resp->code = HA_SERVER_DECLINE;
      }
      goto finally;
    }

  /* A response to a challenge */
  case 3:
    {
      /*
       * We need to have a connection at this point or this whole thing
       * has come in in the wrong order. Actually it's a client error
       * for stuff to come in wrong. But since some web servers also
       * kill keep-alives and stuff, we forgive and just ask the client
       * for the authentication info again.
       */
      if(!conn || !conn->handle)
      {
        ha_messagex(LOG_WARNING, "ntlm: received out of order NTLM response from client");
        resp->code = HA_SERVER_BADREQ;
        goto finally;
      }

      if(!ntlmssp.user)
      {
        ha_messagex(LOG_WARNING, "ntlm: received NTLM response without user name");
        resp->code = HA_SERVER_BADREQ;
        goto finally;
      }

      /* We have to lock while going into smblib */
      ha_lock(&g_smblib_mutex);

        /* Now authenticate them against the DC */
        r = ntlmssp_auth(conn->handle, ntlmssp.user, ntlmssp.nt, 1,
                      ntlmssp.domain[0] ? (char*)ntlmssp.domain : (char*)ctx->domain);

      ha_unlock(&g_smblib_mutex);

      /* The connection gets disconnected below */

      if(r == NTV_LOGON_ERROR)
      {
        /*
         * Note that we don't set a code here. This causes our
         * caller to put in all the proper headers for us.
         */
        ha_messagex(LOG_WARNING, "ntlm: failed NTLM logon for user '%s'", ntlmssp.user);
        ret = HA_FALSE;
      }

      /* A successful login ends here */
      else
      {
        int r;
        resp->detail = ntlmssp.user;
        ha_messagex(LOG_NOTICE, "ntlm: validated ntlm user against server", ntlmssp.user);

        ha_lock(NULL);

          /* We put this connection into the successful connections */
          r = hash_set(ctx->established, key, NTLM_ESTABLISHED);

        ha_unlock(NULL);

        if(!r)
        {
          ha_messagex(LOG_CRIT, "out of memory");
          ret = HA_CRITERROR;
        }
        else
        {
          ret = HA_OK;
        }
      }

      goto finally;
    }

  default:
    ha_messagex(LOG_WARNING, "ntlm: received invalid NTLM message (type %d)", ntlmssp.msg_type);
    resp->code = HA_SERVER_BADREQ;
    goto finally;
  };


finally:
  if(ha_buferr(buf))
    ret = HA_CRITERROR;

  if(conn)
    freeconnection(conn);

  return ret;
}


/* -------------------------------------------------------------------------------
 * Handler Functions
 */

int ntlm_config(ha_context_t* context, const char* name, const char* value)
{
  ntlm_context_t* ctx = (ntlm_context_t*)(context->data);

  ASSERT(name && value && value[0]);

  if(strcmp(name, "ntlmserver") == 0)
  {
    ctx->server = value;
    return HA_OK;
  }

  else if(strcmp(name, "ntlmbackup") == 0)
  {
    ctx->backup = value;
    return HA_OK;
  }

  else if(strcmp(name, "ntlmdomain") == 0)
  {
    ctx->domain = value;
    return HA_OK;
  }

  else if(strcmp(name, "pendingmax") == 0)
  {
    return ha_confint(name, value, 1, 256, &(ctx->pending_max));
  }

  else if(strcmp(name, "pendingtimeout") == 0)
  {
    return ha_confint(name, value, 1, 86400, &(ctx->pending_timeout));
  }

  return HA_FALSE;
}

int ntlm_init(ha_context_t* context)
{
  /* Per context initialization */
  if(context)
  {
    ntlm_context_t* ctx = (ntlm_context_t*)(context->data);

    ASSERT(ctx);

    /* Make sure there are some types of authentication we can do */
    if(!(context->opts->types & (HA_TYPE_BASIC | HA_TYPE_NTLM)))
    {
      ha_messagex(LOG_ERR, "NTLM module configured, but does not implement any "
                           "configured authentication type.");
      return HA_FAILED;
    }

    /* Check for mandatory configuration */
    if(!(ctx->server) || !(ctx->domain))
    {
      ha_messagex(LOG_ERR, "NTLM configuration incomplete. "
                           "Must have NTLMServer and NTLMDomain configured.");
      return HA_FAILED;
    }

    ASSERT(!ctx->pending);
    ASSERT(!ctx->established);

    /* Initialize our tables */
    if(!(ctx->pending = hash_create(NTLM_HASH_KEY_LEN, free_hash_object, NULL)) ||
       !(ctx->established = hash_create(NTLM_HASH_KEY_LEN, NULL, NULL)))
    {
      ha_messagex(LOG_CRIT, "out of memory");
      return HA_CRITERROR;
    }

    ha_messagex(LOG_INFO, "ntlm: initialized handler");
  }

  /* Global Initialization */
  else
  {
    /* Create the smblib mutex */
    if(pthread_mutexattr_init(&g_smblib_mutexattr) != 0 ||
       pthread_mutexattr_settype(&g_smblib_mutexattr, HA_MUTEX_TYPE) ||
       pthread_mutex_init(&g_smblib_mutex, &g_smblib_mutexattr) != 0)
    {
      ha_messagex(LOG_CRIT, "threading problem. can't create mutex");
      return HA_CRITERROR;
    }
  }

  return HA_OK;
}

void ntlm_destroy(ha_context_t* context)
{
  /* Per context destroy */
  if(context)
  {
    /* Note: We don't need to be thread safe here anymore */
    ntlm_context_t* ctx = (ntlm_context_t*)(context->data);

    if(ctx->pending)
      hash_free(ctx->pending);

    if(ctx->established)
      hash_free(ctx->established);

    ha_messagex(LOG_INFO, "ntlm: uninitialized handler");
  }

  /* Global Destroy */
  else
  {
    /* Close the mutex */
    pthread_mutex_destroy(&g_smblib_mutex);
    pthread_mutexattr_destroy(&g_smblib_mutexattr);
  }
}

int ntlm_process(ha_context_t* context, const ha_request_t* req,
                 ha_response_t* resp, ha_buffer_t* buf)
{
  ntlm_context_t* ctx = (ntlm_context_t*)(context->data);
  void* ntlm_connection_t = NULL;
  unsigned char key[NTLM_HASH_KEY_LEN];
  const char* header = NULL;
  time_t t = time(NULL);
  int ret, r;

  ASSERT(context && req && resp && buf);
  ASSERT(req->args[AUTH_ARG_CONN]);

  resp->code = -1;

  /* Hash the unique key */
  md5_string(key, req->args[AUTH_ARG_CONN]);


  ha_lock(NULL);

    /*
     * Purge out stale connection stuff. This includes
     * authenticated connections which have expired as
     * well as half open connections which expire.
     */
    r = hash_purge(ctx->pending, t - ctx->pending_timeout);
    r += hash_purge(ctx->established, t - context->opts->cache_timeout);

  ha_unlock(NULL);

  if(r > 0)
    ha_messagex(LOG_DEBUG, "ntlm: purged info from cache: %d", r);

  /* Look for a NTLM header */
  if(context->opts->types & HA_TYPE_NTLM)
  {
    header = ha_getheader(req, "Authorization", HA_PREFIX_NTLM);
    if(header)
    {
      /* Trim off for decoding */
      while(*header && isspace(*header))
        header++;

      ha_messagex(LOG_DEBUG, "ntlm: processing ntlm auth header");
      ret = ntlm_auth_ntlm(ctx, key, header, resp, buf);
      if(ret < 0)
        return ret;
    }
  }

  /* If basic is enabled, and no NTLM */
  if(!header && context->opts->types & HA_TYPE_BASIC)
  {
    /* Look for a Basic header */
    header = ha_getheader(req, "Authorization", HA_PREFIX_BASIC);
    if(header)
    {
      /* Trim off for decoding */
      while(*header && isspace(*header))
        header++;

      ha_messagex(LOG_DEBUG, "ntlm: processing basic auth header");
      ret = ntlm_auth_basic(ctx, key, header, resp, buf);
      if(ret < 0)
        return ret;
    }
  }

  /* The authorization header was not found */
  else
  {
    ha_lock(NULL);

      /*
       * NTLM trusts a connection after it's been authenticated
       * so just pass success for those. Note that we do this
       * in the absence of a authorization header so that we
       * allow connections to be re-authenticated.
       */

      if(hash_get(ctx->established, key) == NTLM_ESTABLISHED)
      {
        hash_touch(ctx->established, key);
        resp->code = HA_SERVER_OK;
      }

    ha_unlock(NULL);

    if(resp->code == HA_SERVER_OK)
      ha_messagex(LOG_NOTICE, "ntlm: validated user against connection cache");

    /* TODO: We need to be able to retrieve the user here somehow */
  }


  /* If nobody's set any other response then... */
  if(resp->code != -1)
  {
    /* If authentication failed tell the browser about it */
    resp->code = HA_SERVER_DECLINE;

    if(context->opts->types & HA_TYPE_NTLM)
    {
      ha_addheader(resp, "WWW-Authenticate", HA_PREFIX_NTLM);
      ha_messagex(LOG_DEBUG, "ntlm: sent ntlm auth request");
    }

    if(context->opts->types & HA_TYPE_BASIC)
    {
      ha_bufmcat(buf, HA_PREFIX_BASIC, "realm=\"", context->opts->realm, "\"", NULL);

      if(ha_buferr(buf))
        return HA_CRITERROR;

      ha_addheader(resp, "WWW-Authenticate", ha_bufdata(buf));
      ha_messagex(LOG_DEBUG, "ntlm: sent basic auth request");
    }
  }

  return ret;
}



/* -------------------------------------------------------------------------------
 * Handler Definition
 */

ha_handler_t ntlm_handler =
{
  "NTLM",                  /* The type */
  ntlm_init,               /* Initialization function */
  ntlm_destroy,            /* Uninitialization routine */
  ntlm_config,             /* Config routine */
  ntlm_process,            /* Processing routine */
  &ntlm_defaults,           /* Default settings */
  sizeof(ntlm_context_t)
};