nse_main.cc

来自「Ubuntu packages of security software。 相」· CC 代码 · 共 680 行 · 第 1/2 页

#include "nse_main.h"

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

#include "nse_init.h"
#include "nse_nsock.h"
#include "nse_nmaplib.h"
#include "nse_debug.h"
#include "nse_macros.h"
#include "nse_string.h"

#include "nmap.h"
#include "nmap_error.h"
#include "portlist.h"
#include "nsock.h"
#include "NmapOps.h"
#include "timing.h"
#include "Target.h"
#include "nmap_tty.h"

extern NmapOps o;

struct run_record {
	short type; // 0 - hostrule; 1 - portrule
	unsigned int index; // index in the corresponding table
	Port* port;
	Target* host;
};

struct thread_record {
	lua_State* thread;
	int resume_arguments;
	unsigned int registry_idx; // index in the main state registry
	double runlevel;
	run_record* rr;
};

std::map<std::string, Target*> current_hosts;
std::list<std::list<struct thread_record> > torun_scripts;
std::list<struct thread_record> running_scripts;
std::list<struct thread_record> waiting_scripts;

class CompareRunlevels {
public:
	bool operator() (const struct thread_record& lhs, const struct thread_record& rhs) {
		return lhs.runlevel < rhs.runlevel;
	}
};

// prior execution
int process_preparerunlevels(std::list<struct thread_record> torun_threads);
int process_preparehost(lua_State* l, Target* target, std::list<struct thread_record>& torun_threads);
int process_preparethread(lua_State* l, struct run_record rr, struct thread_record* tr);

// helper functions
int process_getScriptId(lua_State* l, struct script_scan_result* ssr);
int process_pickScriptsForPort(
		lua_State* l, 
		Target* target, 
		Port* port,
		std::vector<run_record>& torun);

// execution
int process_mainloop(lua_State* l);
int process_waiting2running(lua_State* l, int resume_arguments);
int process_finalize(lua_State* l, unsigned int registry_idx);

int script_updatedb() {
	int status;
	lua_State* l;
	
	SCRIPT_ENGINE_VERBOSE(
		log_write(LOG_STDOUT, "%s: Updating rule database.\n", 
			SCRIPT_ENGINE);
	)

	l = lua_open();
	if(l == NULL) {
		error("%s: Failed lua_open()", SCRIPT_ENGINE);
		return 0;
	}

	status = init_lua(l);
	if(status != SCRIPT_ENGINE_SUCCESS) {
		goto finishup;
	}

	status = init_updatedb(l);
	if(status != SCRIPT_ENGINE_SUCCESS) {
		goto finishup;
	}

	log_write(LOG_STDOUT, "NSE script database updated successfully.\n");

finishup:
	lua_close(l);
	if(status != SCRIPT_ENGINE_SUCCESS) {
		error("%s: Aborting database update.\n", SCRIPT_ENGINE);
		return SCRIPT_ENGINE_ERROR;
	} else {
		return SCRIPT_ENGINE_SUCCESS;
	}
}

//int check_scripts(){
//}
/* check the script-arguments provided to nmap (--script-args) before
 * scanning starts - otherwise the whole scan will run through and be
 * aborted before script-scanning 
 */
int script_check_args(){
	lua_State* l;
	const char *argbuf;
	size_t argbuflen;

	l= lua_open();
	if(l==NULL){
		fatal("Error opening lua, for checking arguments\n");
	}
	/* set all global libraries (we'll need the string-lib) */
	SCRIPT_ENGINE_TRY(init_lua(l));
	SCRIPT_ENGINE_TRY(init_parseargs(l));
	lua_pushstring(l,"t={");
	lua_insert(l,-2);
	lua_pushstring(l,"}");
	lua_concat(l,3);
	argbuf=lua_tolstring(l,-1,&argbuflen);
	luaL_loadbuffer(l,argbuf,argbuflen,"Script-Arguments-prerun");
	SCRIPT_ENGINE_TRY(lua_pcall(l,0,0,0));

	lua_close(l);
	return SCRIPT_ENGINE_SUCCESS;
}
/* open a lua instance
 * open the lua standard libraries
 * open all the scripts and prepare them for execution
 * 	(export nmap bindings, add them to host/port rulesets etc.)
 * apply all scripts on all hosts
 * */
int script_scan(std::vector<Target*> &targets) {
	int status;
	std::vector<Target*>::iterator target_iter;
	std::list<std::list<struct thread_record> >::iterator runlevel_iter;
	std::list<struct thread_record> torun_threads;
	lua_State* l;

	o.current_scantype = SCRIPT_SCAN;

	SCRIPT_ENGINE_VERBOSE(
			log_write(LOG_STDOUT, "%s: Initiating script scanning.\n", SCRIPT_ENGINE);
			)

	SCRIPT_ENGINE_DEBUGGING(
		unsigned int tlen = targets.size();
		if(tlen > 1)
			log_write(LOG_STDOUT, "%s: Script scanning %d hosts.\n", 
				SCRIPT_ENGINE, tlen);
		else
			log_write(LOG_STDOUT, "%s: Script scanning %s.\n", 
				SCRIPT_ENGINE, (*targets.begin())->HostName());
	)

	l = lua_open();
	if(l == NULL) {
		error("%s: Failed lua_open()", SCRIPT_ENGINE);
		return 0;
	}

	status = init_lua(l);
	if(status != SCRIPT_ENGINE_SUCCESS) {
		goto finishup;
	}
	//set the arguments - if provided
	status = init_setargs(l);
	if(status != SCRIPT_ENGINE_SUCCESS) {
		goto finishup;
	}

	status = init_rules(l, o.chosenScripts);
	if(status != SCRIPT_ENGINE_SUCCESS) {
		goto finishup;
	}

	SCRIPT_ENGINE_DEBUGGING(log_write(LOG_STDOUT, "%s: Matching rules.\n", SCRIPT_ENGINE);)

	for(target_iter = targets.begin(); target_iter != targets.end(); target_iter++) {
		std::string key = ((Target*) (*target_iter))->targetipstr();
		current_hosts[key] = (Target*) *target_iter;

		status = process_preparehost(l, *target_iter, torun_threads);
		if(status != SCRIPT_ENGINE_SUCCESS){
			goto finishup;
		}
	}

	status = process_preparerunlevels(torun_threads);
	if(status != SCRIPT_ENGINE_SUCCESS) {
		goto finishup;
	}

	SCRIPT_ENGINE_DEBUGGING(log_write(LOG_STDOUT, "%s: Running scripts.\n", SCRIPT_ENGINE);)
	
	for(runlevel_iter = torun_scripts.begin(); runlevel_iter != torun_scripts.end(); runlevel_iter++) {
		running_scripts = (*runlevel_iter);

		SCRIPT_ENGINE_DEBUGGING(log_write(LOG_STDOUT, "%s: Runlevel: %f\n", 
					SCRIPT_ENGINE,
					running_scripts.front().runlevel);)

		status = process_mainloop(l);
		if(status != SCRIPT_ENGINE_SUCCESS){
			goto finishup;
		}
	}
	

finishup:
	SCRIPT_ENGINE_DEBUGGING(
			log_write(LOG_STDOUT, "%s: Script scanning completed.\n", SCRIPT_ENGINE);
			)
	lua_close(l);
	current_hosts.clear();
	torun_scripts.clear();
	if(status != SCRIPT_ENGINE_SUCCESS) {
		error("%s: Aborting script scan.", SCRIPT_ENGINE);
		return SCRIPT_ENGINE_ERROR;
	} else {
		return SCRIPT_ENGINE_SUCCESS;
	}
}

int process_mainloop(lua_State* l) {
	int state;
	int unfinished = running_scripts.size() + waiting_scripts.size(); 
	struct script_scan_result ssr;
	struct thread_record current;
	ScanProgressMeter progress = ScanProgressMeter(SCRIPT_ENGINE);

	double total = (double) unfinished;
	double done = 0;

	std::list<struct thread_record>::iterator iter;
	struct timeval now;

	// while there are scripts in running or waiting state, we loop.
	// we rely on nsock_loop to protect us from busy loops when 
	// all scripts are waiting.
	while( unfinished > 0 ) {

		if(l_nsock_loop(50) == NSOCK_LOOP_ERROR) {
			error("%s: An error occured in the nsock loop", SCRIPT_ENGINE);
			return SCRIPT_ENGINE_ERROR;
		}

		unfinished = running_scripts.size() + waiting_scripts.size();

		if (keyWasPressed()) {
			done = 1.0 - (((double) unfinished) / total);
			if (o.verbose > 1 || o.debugging) {
				log_write(LOG_STDOUT, "Active NSE scripts: %d\n", unfinished);
				log_flush(LOG_STDOUT);
			}
			progress.printStats(done, NULL);
		}

		SCRIPT_ENGINE_VERBOSE(
			if(progress.mayBePrinted(NULL)) { 
				done = 1.0 - (((double) unfinished) / total);
				if(o.verbose > 1 || o.debugging)
					progress.printStats(done, NULL);
				else
					progress.printStatsIfNeccessary(done, NULL);
			})

		gettimeofday(&now, NULL);

		for(iter = waiting_scripts.begin(); iter != waiting_scripts.end(); iter++)
			if (iter->rr->host->timedOut(&now)) {
				running_scripts.push_front((*iter));
				waiting_scripts.erase(iter);
				iter = waiting_scripts.begin();
			}


		if(running_scripts.begin() == running_scripts.end())
			continue;

		current = *(running_scripts.begin());

		if (current.rr->host->timedOut(&now))
			state = LUA_ERRRUN;
		else
			state = lua_resume(current.thread, current.resume_arguments);

		if(state == LUA_YIELD) {
			// this script has performed a network io operation
			// we put it in the waiting
			// when the network io operation has completed,
			// a callback from the nsock library will put the
			// script back into the running state
			
			waiting_scripts.push_back(current);
			running_scripts.pop_front();
		} else if( state == 0) {
			// this script has finished
			// we first check if it produced output
			// then we release the thread and remove it from the
			// running_scripts list

			if(lua_isstring (current.thread, -1)) {
				SCRIPT_ENGINE_TRY(process_getScriptId(current.thread, &ssr));
				ssr.output = nse_printable
					(lua_tostring(current.thread, -1), lua_objlen(current.thread, -1));
				if(current.rr->type == 0) {
					current.rr->host->scriptResults.push_back(ssr);
				} else if(current.rr->type == 1) {
					current.rr->port->scriptResults.push_back(ssr);
					current.rr->host->ports.numscriptresults++;
				}
				lua_pop(current.thread, 2);
			}

			SCRIPT_ENGINE_TRY(process_finalize(l, current.registry_idx));
			SCRIPT_ENGINE_TRY(lua_gc(l, LUA_GCCOLLECT, 0));
		} else {
			// this script returned because of an error
			// print the failing reason if the verbose level is high enough	
			SCRIPT_ENGINE_DEBUGGING(
				const char* errmsg = lua_tostring(current.thread, -1);
				log_write(LOG_STDOUT, "%s: %s\n", SCRIPT_ENGINE, errmsg);
			)
			SCRIPT_ENGINE_TRY(process_finalize(l, current.registry_idx));
		}
	}

	progress.endTask(NULL, NULL);