# -*- coding: utf-8 -*-
import multiprocessing
import time
import sys
import imp
from os import listdir,path
import config
import os, signal
import json

import traceback

conn = config.dbconn()

def AttackTargets(module_name, request, response, method, key, is_array=False):
	signal.signal(signal.SIGALRM, timeout_handler)
	try:
		module = imp.load_source(module_name, config.EXPLOITS_PATH+ os.sep + module_name + '.py')
		timeout = module.timeout if module.timeout else config.EXPLOIT_TIMEOUT
		signal.alarm(timeout)
		ret = module.exploit(request, response, method, key, is_array)
		signal.alarm(0)
		return ret
	except TimeoutError:
		print '%s timeout' % module_name
	except Exception as e:
		traceback.print_exc()
		print ('`%s`: %s' % (module_name, e))
		signal.alarm(0)

class TimeoutError(Exception):
	pass

def timeout_handler(signum, frame):
	raise TimeoutError()

class Result:
	def __init__(self, module_name, requestid, method, key, url, host):
		self.requestid=requestid
		self.method=method
		self.key=key
		self.host=host
		self.url=url
		self.module_name = module_name


	def callback(self, ret):
		if not ret: return
		result = ret.get('result', 'unknown')
		info = ret.get('info', 'unknown')
		level = ret.get('level', 'low')
		hash = ret.get('hash', None)
		if check_hash(hash) and result=='unknown':
			result = 'vul'
		result = {'requestid': self.requestid,
				'method': self.method,
				'key': self.key,
				'host': self.host,
				'url': self.url,
				'level': level,
				'hash': hash,
				'info': info,
				'result': result,
				'exploit': self.module_name}
		config.dbconn().insert(table= 'result', data = result)

def get_one():
	r = conn.fetch_rows('requests', condition={'checked': 0}, order="id asc", limit="1", fetchone=True)
	if not r: return None
	r['http'] = conn.fetch_rows('http', condition={'id': r['requestid']}, fetchone=True)
	return r

def mark_checked(id):
	conn.update('requests', {'checked': 1}, {'id': id})

def check_special(module_name, hostname):
	if conn.fetch_rows('result', condition="exploit='%s' and result != 'continue' and `host`='%s'" % (module_name, hostname), limit="1", fetchone=True):
		return True
	else:
		return False

def check_hash(hash):
	if config.dbconn().fetch_rows('hashlog', condition={'hash': hash}, limit="1", fetchone=True):
		return True
	else:
		return False


pool = multiprocessing.Pool()

modules = []
for f in os.listdir(config.EXPLOITS_PATH):
	if f.endswith('.py') and not f.endswith('__init__.py'):
		modules.append(imp.load_source(f[:-3], config.EXPLOITS_PATH+ os.sep + f))

while True:

	r = get_one()
	if not r:
		time.sleep(3)
		continue

	request = json.loads(r['http']['req'])
	response = json.loads(r['http']['rsp'])

	if 'special' in r['type']:
		m = filter(lambda module: 'special' == module.vtype or ('special' in module.vtype and r['type'].replace('special|','') in module.vtype), modules)
		m = filter(lambda module: check_special(module.__name__, request['host']), m)
	elif 'rewrite' in r['type']:
		m = filter(lambda module: 'rewrite' == module.vtype or ('rewrite' in module.vtype and r['type'].replace('rewrite|','') in module.vtype), modules)
	else:
		m = filter(lambda module: (module.vtype == 'all' or r['type'] in module.vtype) and 'special' not in module.vtype and 'rewrite' not in module.vtype, modules)


	stype = conn.fetch_rows('siteinfo', condition={'host': request['host'], 'key': 'lang'}, limit="1", fetchone=True)
	if stype:
		m = filter(lambda module: module.stype == 'all' or stype['value'] == 'unknown' or stype['value'] in module.stype, m)

	method = r['method']
	key = r['key']
	is_array = r['type'] == 'array'
	#print r['type']
	for module in m:
		pool.apply_async(AttackTargets, args = (module.__name__, request, response, method, key, is_array,), callback = Result(module.__name__, r['requestid'], method, key, request['uri'], request['host'],).callback)
	mark_checked(r['id'])


