diff --git a/nse_main.cc b/nse_main.cc index ab9006ff7..dfddfb5d1 100644 --- a/nse_main.cc +++ b/nse_main.cc @@ -29,13 +29,12 @@ #define NSE_WAITING_TO_RUNNING "NSE_WAITING_TO_RUNNING" #define NSE_DESTRUCTOR "NSE_DESTRUCTOR" #define NSE_SELECTED_BY_NAME "NSE_SELECTED_BY_NAME" +#define NSE_CURRENT_HOSTS "NSE_CURRENT_HOSTS" #define MAX_FILENAME_LEN 4096 extern NmapOps o; -int current_hosts = LUA_NOREF; - static int timedOut (lua_State *L) { Target *target = get_target(L, 1); @@ -384,10 +383,8 @@ static int init_main (lua_State *L) luaL_openlibs(L); set_nmap_libraries(L); - /* Load current_hosts */ - luaL_unref(L, LUA_REGISTRYINDEX, current_hosts); lua_newtable(L); - current_hosts = luaL_ref(L, LUA_REGISTRYINDEX); + lua_setfield(L, LUA_REGISTRYINDEX, NSE_CURRENT_HOSTS); /* Load debug.traceback for collecting any error tracebacks */ lua_settop(L, 0); /* clear the stack */ @@ -435,9 +432,8 @@ static int run_main (lua_State *L) lua_settop(L, 0); /* New host group */ - luaL_unref(L, LUA_REGISTRYINDEX, current_hosts); lua_newtable(L); - current_hosts = luaL_ref(L, LUA_REGISTRYINDEX); + lua_setfield(L, LUA_REGISTRYINDEX, NSE_CURRENT_HOSTS); lua_getfield(L, LUA_REGISTRYINDEX, NSE_TRACEBACK); /* index 1 */ @@ -448,6 +444,7 @@ static int run_main (lua_State *L) * This has all the target names, 1-N, in a list. */ lua_createtable(L, targets->size(), 0); // stack index 3 + lua_getfield(L, LUA_REGISTRYINDEX, NSE_CURRENT_HOSTS); /* index 4 */ for (std::vector::iterator ti = targets->begin(); ti != targets->end(); ti++) { @@ -457,15 +454,14 @@ static int run_main (lua_State *L) lua_newtable(L); set_hostinfo(L, target); lua_rawseti(L, 3, lua_objlen(L, 3) + 1); - lua_rawgeti(L, LUA_REGISTRYINDEX, current_hosts); if (TargetName != NULL && strcmp(TargetName, "") != 0) lua_pushstring(L, TargetName); else lua_pushstring(L, targetipstr); lua_pushlightuserdata(L, target); - lua_rawset(L, -3); - lua_pop(L, 1); /* current_hosts */ + lua_rawset(L, 4); /* add to NSE_CURRENT_HOSTS */ } + lua_pop(L, 1); /* pop NSE_CURRENT_HOSTS */ if (lua_pcall(L, 1, 0, 1) != 0) lua_error(L); /* we wanted a traceback */ @@ -539,6 +535,15 @@ void nse_selectedbyname (lua_State *L) } } +void nse_gettarget (lua_State *L, int index) +{ + lua_pushvalue(L, index); + lua_getfield(L, LUA_REGISTRYINDEX, NSE_CURRENT_HOSTS); + lua_insert(L, -2); + lua_rawget(L, -2); + lua_replace(L, -2); +} + static lua_State *L_NSE = NULL; void open_nse (void) diff --git a/nse_main.h b/nse_main.h index 3f2e5251c..0e74d812e 100644 --- a/nse_main.h +++ b/nse_main.h @@ -36,6 +36,7 @@ void nse_restore (lua_State *, int); void nse_destructor (lua_State *, char); void nse_base (lua_State *); void nse_selectedbyname (lua_State *); +void nse_gettarget (lua_State *, int); void open_nse (void); void script_scan (std::vector &targets); diff --git a/nse_nmaplib.cc b/nse_nmaplib.cc index dd6bc33f6..f6b365dff 100644 --- a/nse_nmaplib.cc +++ b/nse_nmaplib.cc @@ -27,7 +27,6 @@ extern "C" { } extern NmapOps o; -extern int current_hosts; void set_version(lua_State *L, struct serviceDeductions sd) { SCRIPT_ENGINE_PUSHSTRING_NOTNULL(sd.name, "name"); @@ -335,11 +334,9 @@ Target *get_target (lua_State *L, int index) lua_getfield(L, index, "ip"); if (!(lua_isstring(L, -2) || lua_isstring(L, -1))) luaL_error(L, "host table does not have a 'ip' or 'targetname' field"); - lua_rawgeti(L, LUA_REGISTRYINDEX, current_hosts); if (lua_isstring(L, -3)) /* targetname */ { - lua_pushvalue(L, -3); - lua_rawget(L, -2); + nse_gettarget(L, -3); /* use targetname */ if (lua_islightuserdata(L, -1)) goto done; else @@ -347,8 +344,7 @@ Target *get_target (lua_State *L, int index) } if (lua_isstring(L, -2)) /* ip */ { - lua_pushvalue(L, -2); /* ip */ - lua_rawget(L, -2); + nse_gettarget(L, -2); /* use ip */ if (!lua_islightuserdata(L, -1)) luaL_argerror(L, 1, "host is not being processed right now"); }