From b61ed83e147844af7cfa10ad88a68f3637dbe83f Mon Sep 17 00:00:00 2001
From: batrick <batrick@e0a8ed71-7df4-0310-8962-fdc924857419>
Date: Fri, 28 Aug 2009 02:48:56 +0000
Subject: [PATCH] [NSE] Patch to allow virtual hosts (specified by name on the
 command line) which resolve to the same IP to have script output placed under
 the correct (virtual) host.  Previously, all script output would be "randomly
 but deterministically" placed under one of these hosts.  Other problems
 include having port information changed for only one of the virtual hosts.

---
 nse_main.cc    | 13 ++++++++++---
 nse_nmaplib.cc | 29 ++++++++++++++++++++++-------
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/nse_main.cc b/nse_main.cc
index 55aa26ba0..4571dc608 100644
--- a/nse_main.cc
+++ b/nse_main.cc
@@ -445,12 +445,19 @@ static int run_main (lua_State *L)
   for (std::vector<Target *>::iterator ti = targets->begin();
        ti != targets->end(); ti++)
   {
+    Target *target = (Target *) *ti;
+    const char *TargetName = target->TargetName();
+    const char *targetipstr = target->targetipstr();
     lua_newtable(L);
-    set_hostinfo(L, (Target *) *ti);
+    set_hostinfo(L, target);
     lua_rawseti(L, 3, lua_objlen(L, 3) + 1);
     lua_rawgeti(L, LUA_REGISTRYINDEX, current_hosts);
-    lua_pushlightuserdata(L, (void *) *ti);
-    lua_setfield(L, -2, (*ti)->targetipstr());
+    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 */
   }
 
diff --git a/nse_nmaplib.cc b/nse_nmaplib.cc
index d57d6d0be..dd6bc33f6 100644
--- a/nse_nmaplib.cc
+++ b/nse_nmaplib.cc
@@ -328,18 +328,33 @@ static int l_mutex (lua_State *L)
 
 Target *get_target (lua_State *L, int index)
 {
+  int top = lua_gettop(L);
   Target *target;
   luaL_checktype(L, index, LUA_TTABLE);
+  lua_getfield(L, index, "targetname");
   lua_getfield(L, index, "ip");
-  if (!lua_isstring(L, -1))
-    luaL_error(L, "host table does not contain 'ip' string field");
+  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);
-  lua_pushvalue(L, -2); /* target ip string */
-  lua_rawget(L, -2);
-  if (!lua_islightuserdata(L, -1))
-    luaL_argerror(L, 1, "host is not being processed right now");
+  if (lua_isstring(L, -3)) /* targetname */
+  {
+    lua_pushvalue(L, -3);
+    lua_rawget(L, -2);
+    if (lua_islightuserdata(L, -1))
+      goto done;
+    else
+      lua_pop(L, 1);
+  }
+  if (lua_isstring(L, -2)) /* ip */
+  {
+    lua_pushvalue(L, -2); /* ip */
+    lua_rawget(L, -2);
+    if (!lua_islightuserdata(L, -1))
+      luaL_argerror(L, 1, "host is not being processed right now");
+  }
+done:
   target = (Target *) lua_touserdata(L, -1);
-  lua_pop(L, 3); /* target ip string, current_hosts, target luserdata */
+  lua_settop(L, top); /* reset stack */
   return target;
 }