/* lrexlib.c - PCRE regular expression library */
/* Reuben Thomas   nov00-18dec04 */
/* Shmuel Zeigerman   may04-18dec04 */


#include <stdio.h>
#include <stdlib.h>
#include <string.h>

extern "C" {
  #include "lua.h"
  #include "lauxlib.h"
}

#include <locale.h>

#include "nbase.h"
#include "nmap_error.h"

#ifdef HAVE_PCRE_PCRE_H
# include <pcre/pcre.h>
#else
# include <pcre.h>
#endif

#include "nse_pcrelib.h"

static int get_startoffset(lua_State *L, int stackpos, size_t len)
{
        int startoffset = luaL_optint(L, stackpos, 1);
        if(startoffset > 0)
                startoffset--;
        else if(startoffset < 0) {
                startoffset += (int) len;
                if(startoffset < 0)
                        startoffset = 0;
        }
        return startoffset;
}

static int udata_tostring (lua_State *L, const char* type_handle,
                const char* type_name)
{
        char buf[256];
        void *udata = luaL_checkudata(L, 1, type_handle);

        if(udata) {
                (void)Snprintf(buf, 255, "%s (%p)", type_name, udata);
                lua_pushstring(L, buf);
        }
        else {
                (void)Snprintf(buf, 255, "must be userdata of type '%s'", type_name);
                (void)luaL_argerror(L, 1, buf);
        }

        free(udata);
        return 1;
}

typedef struct { const char* key; lua_Number val; } flags_pair;

static int get_flags (lua_State *L, const flags_pair *arr)
{
        const flags_pair *p;
        lua_newtable(L);
        for(p=arr; p->key != NULL; p++) {
                lua_pushstring(L, p->key);
                lua_pushnumber(L, p->val);
                lua_rawset(L, -3);
        }
        return 1;
}

const char pcre_handle[] = "pcre_regex_handle";
const char pcre_typename[] = "pcre_regex";

typedef struct {
        pcre *pr;
        pcre_extra *extra;
        int *match;
        int ncapt;
        const unsigned char *tables;
} pcre2;      /* a better name is needed */

static const unsigned char *Lpcre_maketables(lua_State *L, int stackpos)
{
        const unsigned char *tables;
        char old_locale[256];
        char *locale = strdup(luaL_checkstring(L, stackpos));

        if(locale == NULL)
                luaL_error(L, "cannot set locale");

        strncpy(old_locale, setlocale(LC_CTYPE, NULL), 255); /* store the locale */

        if(setlocale(LC_CTYPE, locale) == NULL)        /* set new locale */
                luaL_error(L, "cannot set locale");

        tables = pcre_maketables();              /* make tables with new locale */
        (void)setlocale(LC_CTYPE, old_locale);         /* restore the old locale */

        free(locale);
        return tables;
}

static int Lpcre_comp(lua_State *L)
{
        char buf[256];
        const char *error;
        int erroffset;
        pcre2 *ud;
        const char *pattern = luaL_checkstring(L, 1);
        int cflags = luaL_optint(L, 2, 0);
        const unsigned char *tables = NULL;

        if(lua_gettop(L) > 2 && !lua_isnil(L, 3))
                tables = Lpcre_maketables(L, 3);
        if(tables == NULL)
                luaL_error(L, "PCRE compilation failed");

        ud = (pcre2*)lua_newuserdata(L, sizeof(pcre2));
        luaL_getmetatable(L, pcre_handle);
        (void)lua_setmetatable(L, -2);
        ud->match = NULL;
        ud->extra = NULL;
        ud->tables = tables; /* keep this for eventual freeing */

        ud->pr = pcre_compile(pattern, cflags, &error, &erroffset, tables);
        if(!ud->pr) {
                (void)Snprintf(buf, 255, "%s (pattern offset: %d)", error, erroffset+1);
                /* show offset 1-based as it's common in Lua */
                luaL_error(L, buf);
        }

        ud->extra = pcre_study(ud->pr, 0, &error);
        if(error) luaL_error(L, error);

        pcre_fullinfo(ud->pr, ud->extra, PCRE_INFO_CAPTURECOUNT, &ud->ncapt);
        /* need (2 ints per capture, plus one for substring match) * 3/2 */
        ud->match = (int *) safe_malloc((ud->ncapt + 1) * 3 * sizeof(int));

        return 1;
}

static void Lpcre_getargs(lua_State *L, pcre2 **pud, const char **text,
                size_t *text_len)
{
        *pud = (pcre2 *)luaL_checkudata(L, 1, pcre_handle);
        if(*pud == NULL)
                (void)luaL_argerror(L, 1, ("compiled regexp expected"));
        *text = luaL_checklstring(L, 2, text_len);
}

typedef void (*Lpcre_push_matches) (lua_State *L, const char *text, pcre2 *ud);

static void Lpcre_push_substrings (lua_State *L, const char *text, pcre2 *ud)
{
        unsigned int i, j;
        unsigned int namecount;
        unsigned char *name_table;
        int name_entry_size;
        unsigned char *tabptr;
        const int *match = ud->match;

        lua_newtable(L);
        for (i = 1; i <= (unsigned) ud->ncapt; i++) {
                j = i * 2;
                if (match[j] >= 0)
                        lua_pushlstring(L, text + match[j], (size_t)(match[j + 1] - match[j]));
                else
                        lua_pushboolean(L, 0);
                lua_rawseti(L, -2, i);
        }

        /* now do named subpatterns - NJG */
        (void)pcre_fullinfo(ud->pr, ud->extra, PCRE_INFO_NAMECOUNT, &namecount);
        if (namecount <= 0)
                return;
        (void)pcre_fullinfo(ud->pr, ud->extra, PCRE_INFO_NAMETABLE, &name_table);
        (void)pcre_fullinfo(ud->pr, ud->extra, PCRE_INFO_NAMEENTRYSIZE, &name_entry_size);
        tabptr = name_table;
        for (i = 0; i < namecount; i++) {
                unsigned int n = (tabptr[0] << 8) | tabptr[1]; /* number of the capturing parenthesis */
                if (n > 0 && n <= (unsigned) ud->ncapt) {   /* check range */
                        unsigned int j = n * 2;
                        lua_pushstring(L, (char*)tabptr + 2); /* name of the capture, zero terminated */
                        if (match[j] >= 0)
                                lua_pushlstring(L, text + match[j], match[j + 1] - match[j]);
                        else
                                lua_pushboolean(L, 0);
                        lua_rawset(L, -3);
                }
                tabptr += name_entry_size;
        }
}

static void Lpcre_push_offsets (lua_State *L, const char *text, pcre2 * ud)
{
        unsigned int i, j, k;
        if(text) {
                /* suppress compiler warning */
        }
        lua_newtable(L);
        for (i=1, j=1; i <= (unsigned) ud->ncapt; i++) {
                k = i * 2;
                if (ud->match[k] >= 0) {
                        lua_pushnumber(L, ud->match[k] + 1);
                        lua_rawseti(L, -2, j++);
                        lua_pushnumber(L, ud->match[k+1]);
                        lua_rawseti(L, -2, j++);
                }
                else {
                        lua_pushboolean(L, 0);
                        lua_rawseti(L, -2, j++);
                        lua_pushboolean(L, 0);
                        lua_rawseti(L, -2, j++);
                }
        }
}

static int Lpcre_match_generic(lua_State *L, Lpcre_push_matches push_matches)
{
        int res;
        const char *text;
        pcre2 *ud;
        size_t elen;
        int startoffset;
        int eflags = luaL_optint(L, 4, 0);

        Lpcre_getargs(L, &ud, &text, &elen);
        startoffset = get_startoffset(L, 3, elen);

        res = pcre_exec(ud->pr, ud->extra, text, (int)elen, startoffset, eflags,
                        ud->match, (ud->ncapt + 1) * 3);
        if (res >= 0) {
                lua_pushnumber(L, (lua_Number) ud->match[0] + 1);
                lua_pushnumber(L, (lua_Number) ud->match[1]);
                (*push_matches)(L, text, ud);
                return 3;
        }
        return 0;
}

static int Lpcre_match(lua_State *L)
{
        return Lpcre_match_generic(L, Lpcre_push_substrings);
}

static int Lpcre_exec(lua_State *L)
{
        return Lpcre_match_generic(L, Lpcre_push_offsets);
}

static int Lpcre_gmatch(lua_State *L)
{
        int res;
        size_t len;
        int nmatch = 0, limit = 0;
        const char *text;
        pcre2 *ud;
        int maxmatch = luaL_optint(L, 4, 0);
        int eflags = luaL_optint(L, 5, 0);
        int startoffset = 0;
        Lpcre_getargs(L, &ud, &text, &len);
        luaL_checktype(L, 3, LUA_TFUNCTION);

        if(maxmatch > 0) /* this must be stated in the docs */
                limit = 1;

        while (!limit || nmatch < maxmatch) {
                res = pcre_exec(ud->pr, ud->extra, text, (int)len, startoffset, eflags,
                                ud->match, (ud->ncapt + 1) * 3);
                if (res >= 0) {
                        nmatch++;
                        lua_pushvalue(L, 3);
                        lua_pushlstring(L, text + ud->match[0], ud->match[1] - ud->match[0]);
                        Lpcre_push_substrings(L, text, ud);
                        lua_call(L, 2, 1);
                        if(lua_toboolean(L, -1))
                                break;
                        lua_pop(L, 1);
                        startoffset = ud->match[1];
                } else
                        break;
        }
        lua_pushnumber(L, nmatch);
        return 1;
}

static int Lpcre_gc (lua_State *L)
{
        pcre2 *ud = (pcre2 *)luaL_checkudata(L, 1, pcre_handle);
        if (ud) {
                if(ud->pr)      pcre_free(ud->pr);
                if(ud->extra)   pcre_free(ud->extra);
                if(ud->tables)  pcre_free((void *)ud->tables);
                if(ud->match)   free(ud->match);
        }
        return 0;
}

static int Lpcre_tostring (lua_State *L) {
        return udata_tostring(L, pcre_handle, pcre_typename);
}

static int Lpcre_vers (lua_State *L)
{
        lua_pushstring(L, pcre_version());
        return 1;
}

static flags_pair pcre_flags[] =
{
        { "CASELESS",        PCRE_CASELESS },
        { "MULTILINE",       PCRE_MULTILINE },
        { "DOTALL",          PCRE_DOTALL },
        { "EXTENDED",        PCRE_EXTENDED },
        { "ANCHORED",        PCRE_ANCHORED },
        { "DOLLAR_ENDONLY",  PCRE_DOLLAR_ENDONLY },
        { "EXTRA",           PCRE_EXTRA },
        { "NOTBOL",          PCRE_NOTBOL },
        { "NOTEOL",          PCRE_NOTEOL },
        { "UNGREEDY",        PCRE_UNGREEDY },
        { "NOTEMPTY",        PCRE_NOTEMPTY },
        { "UTF8",            PCRE_UTF8 },
#if PCRE_MAJOR >= 4
        { "NO_AUTO_CAPTURE", PCRE_NO_AUTO_CAPTURE },
        { "NO_UTF8_CHECK",   PCRE_NO_UTF8_CHECK },
#endif
#ifdef PCRE_AUTO_CALLOUT
        { "AUTO_CALLOUT",    PCRE_AUTO_CALLOUT },
#endif
#ifdef PCRE_PARTIAL
        { "PARTIAL",         PCRE_PARTIAL },
#endif
#ifdef PCRE_DFA_SHORTEST
        { "DFA_SHORTEST",    PCRE_DFA_SHORTEST },
#endif
#ifdef PCRE_DFA_RESTART
        { "DFA_RESTART",     PCRE_DFA_RESTART },
#endif
#ifdef PCRE_FIRSTLINE
        { "FIRSTLINE",       PCRE_FIRSTLINE },
#endif
#ifdef PCRE_DUPNAMES
        { "DUPNAMES",        PCRE_DUPNAMES },
#endif
#ifdef PCRE_NEWLINE_CR
        { "NEWLINE_CR",      PCRE_NEWLINE_CR },
#endif
#ifdef PCRE_NEWLINE_LF
        { "NEWLINE_LF",      PCRE_NEWLINE_LF },
#endif
#ifdef PCRE_NEWLINE_CRLF
        { "NEWLINE_CRLF",    PCRE_NEWLINE_CRLF },
#endif
#ifdef PCRE_NEWLINE_ANY
        { "NEWLINE_ANY",     PCRE_NEWLINE_ANY },
#endif
#ifdef PCRE_NEWLINE_ANYCRLF
        { "NEWLINE_ANYCRLF", PCRE_NEWLINE_ANYCRLF },
#endif
#ifdef PCRE_BSR_ANYCRLF
        { "BSR_ANYCRLF",     PCRE_BSR_ANYCRLF },
#endif
#ifdef PCRE_BSR_UNICODE
        { "BSR_UNICODE",     PCRE_BSR_UNICODE },
#endif
        { NULL, 0 }
};

static int Lpcre_get_flags (lua_State *L) {
        return get_flags(L, pcre_flags);
}

static const luaL_Reg pcremeta[] = {
        {"exec",       Lpcre_exec},
        {"match",      Lpcre_match},
        {"gmatch",     Lpcre_gmatch},
        {"__gc",       Lpcre_gc},
        {"__tostring", Lpcre_tostring},
        {NULL, NULL}
};

/* Open the library */
static const luaL_Reg pcrelib[] = {
        {"new",	Lpcre_comp},
        {"flags", Lpcre_get_flags},
        {"version", Lpcre_vers},
        {NULL, NULL}
};

LUALIB_API int luaopen_pcrelib(lua_State *L)
{
        luaL_newmetatable(L, pcre_handle);
        lua_pushliteral(L, "__index");
        luaL_newlib(L, pcremeta);
        lua_rawset(L, -3);
        lua_pop(L, 1);

        luaL_newlib(L, pcrelib);

        return 1;
}
