#include "usuals.h"
#include "httpauthd.h"
#include "md5.h"
#include "sha1.h"
#include "bd.h"

#include <sys/time.h>

/* Postgresql library */
#include <libpq-fe.h>

/* -------------------------------------------------------------------------------
 * Structures
 */

#define DB_PW_CLEAR      0
#define DB_PW_SHA1       1
#define DB_PW_MD5        2
#define DB_PW_CRYPT      3

/* Our hanler context */
typedef struct pgsql_context
{
    /* Base Handler ------------------------------------------------------ */
    bd_context_t bd;

    /* Settings ---------------------------------------------------------- */
    const char* host;         /* The connection host or path */
    const char* port;         /* The connection port */
    const char* user;         /* The pgsql user name */
    const char* password;     /* The pgsql password */
    const char* database;     /* The database name */
    const char* query;        /* The query */
    const char* pw_column;    /* The database query to retrieve a password */
    int pw_type;              /* The type of password encoded in database */
    const char* ha1_column;   /* The database query to retrieve a ha1 */

    int pgsql_max;            /* Number of open connections allowed */
    int pgsql_timeout;        /* Maximum amount of time to dedicate to a query */

    /* Context ----------------------------------------------------------- */
    PGconn** pool;            /* Pool of available connections */
    int pool_mark;          /* Amount of connections allocated */
}
pgsql_context_t;

/* Forward declarations for callbacks */
static int validate_digest(ha_request_t* rq, const char* user, digest_context_t* dg);
static int validate_basic(ha_request_t* rq, const char* user, const char* password);
static void escape_pgsql(const ha_request_t* rq, ha_buffer_t* buf, const char* value);

/* The defaults for the context */
static const pgsql_context_t pgsql_defaults =
{
    BD_CALLBACKS(validate_digest,
        validate_basic, escape_pgsql),
    NULL,               /* host */
    0,                  /* port */
    NULL,               /* user */
    NULL,               /* password */
    NULL,               /* database */
    NULL,               /* query */
    NULL,               /* pw_attr */
    DB_PW_CLEAR,        /* pw_type */
    NULL,               /* ha1_attr */
    10,                 /* pgsql_max */
    30,                 /* pgsql_timeout */
    NULL,               /* pool */
    0                   /* pool_mark */
};


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

static void escape_pgsql(const ha_request_t* rq, ha_buffer_t* buf, const char* value)
{
    size_t len;
    char* t;

    ASSERT(value);

    len = strlen(value);

    t = (char*)malloc((len * 2) + 1);
    if(t != NULL)
    {
        PQescapeString(t, value, len);
        ha_bufcpy(buf, t);
        free(t);
    }
}

static int dec_pgsql_binary(const ha_request_t* rq, const char* enc,
                            unsigned char* bytes, size_t len)
{
    size_t enclen;
    void* d;

    ASSERT(rq && enc && bytes && len);
    enclen = strlen(enc);

    /* Hex encoded */
    if(enclen == (len * 2))
    {
        d = ha_bufdechex(rq->buf, enc, &enclen);

        if(d && enclen == len)
        {
            ha_messagex(rq, LOG_DEBUG, "found value in hex encoded format");
            memcpy(bytes, d, len);
            return HA_OK;
        }
    }

    /* Raw binary postgres encoded */
    d = PQunescapeBytea(enc, &enclen);
    if(d != NULL)
    {
        if(enclen == len)
            memcpy(bytes, d, len);

        PQfreemem(d);

        if(enclen == len)
        {
            ha_messagex(rq, LOG_DEBUG, "found value in raw binary format");
            return HA_OK;
        }
    }

    /* B64 Encoded */
    enclen = strlen(enc);
    d = ha_bufdec64(rq->buf, enc, &enclen);

    if(d && len == enclen)
    {
        ha_messagex(rq, LOG_DEBUG, "found value in b64 encoded format");
        memcpy(bytes, d, len);
        return HA_OK;
    }

    return ha_buferr(rq->buf) ? HA_CRITERROR : HA_FALSE;
}

static int validate_ha1(ha_request_t* rq, pgsql_context_t* ctx, const char* user,
                        const char* clearpw, const char* dbpw)
{
    unsigned char dbha1[MD5_LEN];
    unsigned char ha1[MD5_LEN];
    const char* p;

    int r = dec_pgsql_binary(rq, dbpw, dbha1, MD5_LEN);

    if(r < 0)
        return r;

    if(r == HA_OK)
    {
        digest_makeha1(ha1, user, rq->context->realm, clearpw);
        if(memcmp(ha1, dbha1, MD5_LEN) == 0)
            return HA_OK;
    }

    return HA_FALSE;
}

static int validate_password(ha_request_t* rq, pgsql_context_t* ctx, const char* user,
                             const char* clearpw, const char* dbpw)
{
    unsigned char buf[SHA1_LEN];
    const char* p;
    int r;

    ASSERT(SHA1_LEN > MD5_LEN);

    switch(ctx->pw_type)
    {

    /* Clear text */
    case DB_PW_CLEAR:

        if(strcmp(clearpw, dbpw) == 0)
        {
            ha_messagex(rq, LOG_DEBUG, "found matching clear text password");
            return HA_OK;
        }

        break;

    /* Crypt pw */
    case DB_PW_CRYPT:

        ha_lock(NULL);
            p = (const char*)crypt(clearpw, dbpw);
        ha_unlock(NULL);

        if(p && strcmp(dbpw, p) == 0)
        {
            ha_messagex(rq, LOG_DEBUG, "found matching crypt password");
            return HA_OK;
        }

        break;

    /* MD5 */
    case DB_PW_MD5:

        r = dec_pgsql_binary(rq, dbpw, buf, MD5_LEN);
        if(r < 0) return r;

        if(r == HA_OK && md5_strcmp(buf, clearpw) == 0)
        {
            ha_messagex(rq, LOG_DEBUG, "found matching md5 password");
            return HA_OK;
        }

        break;


    /* SHA1 */
    case DB_PW_SHA1:

        r = dec_pgsql_binary(rq, dbpw, buf, SHA1_LEN);
        if(r < 0) return r;

        if(r == HA_OK && sha1_strcmp(buf, clearpw) == 0)
        {
            ha_messagex(rq, LOG_DEBUG, "found matching sha1 password");
            return HA_OK;
        }

        break;

    default:
        ASSERT(0 && "invalid password type");
        break;
    };

    return HA_FALSE;
}

static PGconn* get_pgsql_connection(const ha_request_t* rq, pgsql_context_t* ctx)
{
    PGconn* pg;
    int i;

    ASSERT(ctx);

    for(i = 0; i < ctx->pgsql_max; i++)
    {
        /* An open connection in the pool */
        if(ctx->pool[i])
        {
            ha_messagex(rq, LOG_DEBUG, "using cached pgsql connection");
            pg = ctx->pool[i];
            ctx->pool[i] = NULL;
            return pg;
        }
    }

    if(ctx->pool_mark >= ctx->pgsql_max)
    {
        ha_messagex(rq, LOG_ERR, "too many open pgsql connections");
        return NULL;
    }

    pg = PQsetdbLogin(ctx->host, ctx->port, NULL, NULL, ctx->database,
                      ctx->user, ctx->password);
    if(!pg)
    {
        ha_messagex(rq, LOG_CRIT, "internal error in postgres library");
        return NULL;
    }

    if(PQstatus(pg) == CONNECTION_BAD)
    {
        ha_messagex(rq, LOG_ERR, "error opening pgsql connection: %s", PQerrorMessage(pg));
        PQfinish(pg);
        return NULL;
    }

    ctx->pool_mark++;
    ha_messagex(rq, LOG_DEBUG, "opened new pgsql connection (total %d)", ctx->pool_mark);
    return pg;
}

static void discard_pgsql_connection(const ha_request_t* rq, pgsql_context_t* ctx, PGconn* pg)
{
    PQfinish(pg);
    ctx->pool_mark--;
    ha_messagex(rq, LOG_DEBUG, "discarding pgsql connection (total %d)", ctx->pool_mark);
}

static void save_pgsql_connection(const ha_request_t* rq, pgsql_context_t* ctx, PGconn* pg)
{
    int i, e;

    ASSERT(ctx);

    if(!pg)
        return;

    /* Make sure it's worth saving */
    if(PQstatus(pg) != CONNECTION_BAD)
    {
        for(i = 0; i < ctx->pgsql_max; i++)
        {
            /* An open connection in the pool */
            if(!ctx->pool[i])
            {
                ha_messagex(rq, LOG_DEBUG, "caching pgsql connection for later use");
                ctx->pool[i] = pg;
                pg = NULL;
                break;
            }
        }
    }

    if(pg != NULL)
        discard_pgsql_connection(rq, ctx, pg);
}

static int check_pgsql_result(ha_request_t* rq, PGresult* res)
{
    switch(PQresultStatus(res))
    {
    case PGRES_COMMAND_OK:
    case PGRES_EMPTY_QUERY:
        ha_messagex(rq, LOG_WARNING, "query did not return data");
        return HA_FALSE;
    case PGRES_TUPLES_OK:
        return HA_OK;
    case PGRES_BAD_RESPONSE:
        ha_messagex(rq, LOG_ERR, "error communicating with pgsql server");
        return HA_FAILED;
    case PGRES_NONFATAL_ERROR:
        ha_messagex(rq, LOG_ERR, "warning querying database: %s", PQresultErrorMessage(res));
        return HA_OK;
    case PGRES_FATAL_ERROR:
        ha_messagex(rq, LOG_ERR, "error querying database: %s", PQresultErrorMessage(res));
        return HA_FAILED;
    case PGRES_COPY_OUT:
    case PGRES_COPY_IN:
    default:
        ASSERT(0 && "unexpected response");
        return HA_FAILED;
    };

    return HA_OK;
}

static int resolve_column(PGresult* res, const char* column)
{
    int i;

    if(column)
    {
        for(i = 0; i < PQnfields(res); i++)
        {
            if(strcasecmp(column, PQfname(res, i)) == 0)
                return i;
        }
    }

    return -1;
}

static int retrieve_user_rows(ha_request_t* rq, pgsql_context_t* ctx,
                              const char* user, PGresult** results)
{
    PGconn* pg = NULL;
    PGresult* res = NULL;
    const char* query;
    int ret = HA_OK;

    ASSERT(rq && ctx && user && results);
    *results = NULL;

    pg = get_pgsql_connection(rq, ctx);
    if(!pg)
        RETURN(HA_FAILED);

    ASSERT(ctx->query);

    /* The map can have %u and %r to denote user and realm */
    query = bd_substitute(rq, user, ctx->query);
    if(!query)
        RETURN(HA_CRITERROR);

    ha_messagex(rq, LOG_DEBUG, "executing query: %s", query);
    res = PQexec(pg, query);


    ret = check_pgsql_result(rq, res);
    if(ret != HA_OK)
        RETURN(ret);

    if(PQntuples(res) == 0)
    {
        ha_messagex(rq, LOG_WARNING, "login failed. couldn't find user: %s", user);
        RETURN(HA_FALSE);
    }

    if(PQnfields(res) <= 0)
    {
        ha_messagex(rq, LOG_ERR, "query returned 0 columns: %s", query);
        RETURN(HA_FAILED);
    }

    *results = res;
    res = NULL;
    ha_messagex(rq, LOG_DEBUG, "received %d result rows", PQntuples(res));

finally:

    if(res != NULL)
        PQclear(res);

    /* According to libpg we can close/save the connection
     * before the returned results are freed, no worries there */
    if(pg != NULL)
        save_pgsql_connection(rq, ctx, pg);

    return ret;
}

static int validate_digest(ha_request_t* rq, const char* user, digest_context_t* dg)
{
    pgsql_context_t* ctx = (pgsql_context_t*)rq->context->ctx_data;
    PGresult* res = NULL;
    int ret = HA_FALSE;
    int pw_column = -1;
    int ha1_column = -1;
    int r, i, foundany = 0;

    ASSERT(rq && user && dg);

    ret = retrieve_user_rows(rq, ctx, user, &res);
    if(ret != HA_OK)
        RETURN(ret);

    ASSERT(res);

    pw_column = resolve_column(res, ctx->pw_column);
    ha1_column = resolve_column(res, ctx->ha1_column);

    if(pw_column == -1 && ha1_column == -1)
    {
        if(PQnfields(res) > 1)
            ha_messagex(rq, LOG_WARNING, "query returned more than 1 column, using first as password");

        pw_column = 0;
    }

    for(i = 0; i < PQntuples(res); i++)
    {
        if(pw_column != -1)
        {
            if(ctx->pw_type == DB_PW_CLEAR && !PQgetisnull(res, i, pw_column))
            {
                foundany = 1;

                digest_makeha1(dg->ha1, user, rq->context->realm, PQgetvalue(res, i, pw_column));
                ha_messagex(rq, LOG_DEBUG, "testing clear text password for digest auth");

                /* Run the actual check */
                ret = digest_complete_check(dg, rq->buf);

                if(ret != HA_FALSE)
                    RETURN(ret);
            }
        }

        if(ha1_column != -1)
        {
            if(!PQgetisnull(res, i, ha1_column))
            {
                ret = dec_pgsql_binary(rq, PQgetvalue(res, i, ha1_column), dg->ha1, MD5_LEN);
                if(ret < 0)
                    RETURN(ret)
                else if(ret == HA_FALSE)
                    continue;

                foundany = 1;

                /* Run the actual check */
                ret = digest_complete_check(dg, rq->buf);

                if(ret != HA_FALSE)
                    RETURN(ret);
            }
        }
    }

    if(!foundany)
        ha_messagex(rq, LOG_WARNING, "no clear password or ha1 present for user: %s", user);

finally:
    if(res)
        PQclear(res);

    return ret;
}

static int validate_basic(ha_request_t* rq, const char* user, const char* password)
{
    pgsql_context_t* ctx = (pgsql_context_t*)rq->context->ctx_data;
    PGresult* res = NULL;
    int ret = HA_FALSE;
    int pw_column = -1;
    int ha1_column = -1;
    int i, foundany = 0;

    ASSERT(rq && user && password);

    ret = retrieve_user_rows(rq, ctx, user, &res);
    if(ret != HA_OK)
        RETURN(ret);

    ASSERT(res);

    pw_column = resolve_column(res, ctx->pw_column);
    ha1_column = resolve_column(res, ctx->ha1_column);

    if(pw_column == -1 && ha1_column == -1)
    {
        if(PQnfields(res) > 1)
            ha_messagex(rq, LOG_WARNING, "query returned more than 1 column, using first as password");
        pw_column = 0;
    }


    for(i = 0; i < PQntuples(res); i++)
    {
        if(pw_column != -1)
        {
            if(!PQgetisnull(res, i, pw_column))
            {
                foundany = 1;
                ret = validate_password(rq, ctx, user, password, PQgetvalue(res, i, pw_column));
                if(ret != HA_FALSE)
                    RETURN(ret);
            }
        }

        if(ha1_column != -1)
        {
            if(!PQgetisnull(res, i, ha1_column))
            {
                foundany = 1;
                ret = validate_ha1(rq, ctx, user, password, PQgetvalue(res, i, ha1_column));
                if(ret != HA_FALSE)
                    RETURN(ret);
            }
        }
    }

    if(!foundany)
        ha_messagex(rq, LOG_WARNING, "no password present for user: %s", user);

finally:
    if(res)
        PQclear(res);

    return ret;
}


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

int pgsql_config(ha_context_t* context, const char* name, const char* value)
{
    pgsql_context_t* ctx = (pgsql_context_t*)(context->ctx_data);

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

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

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

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

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

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

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

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

    if(strcmp(name, "dbpwtype") == 0)
    {
        if(strcasecmp(value, "clear") == 0)
            ctx->pw_type = DB_PW_CLEAR;
        else if(strcasecmp(value, "crypt") == 0)
            ctx->pw_type = DB_PW_CRYPT;
        else if(strcasecmp(value, "md5") == 0)
            ctx->pw_type = DB_PW_MD5;
        else if(strcasecmp(value, "sha1") == 0)
            ctx->pw_type = DB_PW_SHA1;
        else
        {
            ha_messagex(NULL, LOG_ERR, "invalid value for '%s' (must be 'clear', 'crypt', 'md5' or 'sha1')", name);
            return HA_FAILED;
        }

        return HA_OK;
    }

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

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

    else if(strcmp(name, "dbtimeout") == 0)
    {
        /* TODO: Implement database timeouts */
        return ha_confint(name, value, 0, 86400, &(ctx->pgsql_timeout));
    }

    return HA_FALSE;
}

int pgsql_init(ha_context_t* context)
{
    int r;

    if((r = bd_init(context)) != HA_OK)
        return r;

    /* Context specific initialization */
    if(context)
    {
        pgsql_context_t* ctx = (pgsql_context_t*)(context->ctx_data);
        ASSERT(ctx);

        /* Check for mandatory configuration */
        if(!ctx->database || !ctx->query)
        {
            ha_messagex(NULL, LOG_ERR, "pgsql configuration incomplete. "
                                       "Must have DBDatabase and DBQuery.");
            return HA_FAILED;
        }

        ASSERT(!ctx->pool);
        ASSERT(ctx->pgsql_max > 0);

        /*
         * Our connection pool. It's the size of our maximum
         * amount of pending connections as that's the max
         * we'd be able to use at a time anyway.
         */
        ctx->pool = (PGconn**)malloc(sizeof(PGconn*) * ctx->pgsql_max);
        if(!ctx->pool)
        {
            ha_messagex(NULL, LOG_CRIT, "out of memory");
            return HA_CRITERROR;
        }

        memset(ctx->pool, 0, sizeof(PGconn*) * ctx->pgsql_max);
        ha_messagex(NULL, LOG_INFO, "initialized pgsql handler");
    }

    return HA_OK;
}

void pgsql_destroy(ha_context_t* context)
{
    if(context)
    {
        /* Note: We don't need to be thread safe here anymore */
        pgsql_context_t* ctx = (pgsql_context_t*)(context->ctx_data);
        int i;

        ASSERT(ctx);

        if(ctx->pool)
        {
            /* Close any connections we have open */
            for(i = 0; i < ctx->pgsql_max; i++)
            {
                if(ctx->pool[i])
                    PQfinish(ctx->pool[i]);
            }

            /* And free the connection pool */
            free(ctx->pool);
        }
    }

    bd_destroy(context);
    ha_messagex(NULL, LOG_INFO, "uninitialized pgsql handler");
}



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

ha_handler_t pgsql_handler =
{
    "PGSQL",                 /* The type */
    pgsql_init,              /* Initialization function */
    pgsql_destroy,           /* Uninitialization routine */
    pgsql_config,            /* Config routine */
    bd_process,              /* Processing routine */
    &pgsql_defaults,          /* The context defaults */
    sizeof(pgsql_context_t)
};