#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : Lucifer
# @Time    : 2019-01-27 13:28
# @File    : hostexploit.py
# @Software: PyCharm

from gevent import monkey;monkey.patch_all()
from gevent.queue import Queue
from gevent.lock import BoundedSemaphore
from conf.config import GlobalConf
import gevent
import sys
import json
sys.path.append('..')
import pymysql
import importlib
from Heaven_Hell.database import db
from pyfancy.pyfancy import pyfancy
from Backtracking.SatanLogging import mylog
sys.tracebacklimit = 0

sem = BoundedSemaphore(1)

class hostpocfactory:
    def __init__(self, host, port, threads):
        self.host = host
        self.port = port
        self.threads = threads
        self.vulns = list()
        self.queue = Queue(self.threads)

    def loadmodule(self):
        requirements = ['redis', 'socket', 're', 'pymongo', 'sys', 'os',
                        'paramiko', 'binascii', 'struct', 'string', 'random',
                        'codecs', 'time', 'select', 'ssl', 'telnetlib', 'hashlib',
                        'io', 'string', 'urllib']
        imported_libs = {lib: importlib.import_module(lib) for lib in requirements}
        globals().update(imported_libs)

    def Consumer(self, pocstr):
        while not self.queue.empty():
            data = self.queue.get()
            sem.acquire()
            exec(data)
            gevent.sleep(0)
            sem.release()

    def runpocwithsysname(self, keyword):
        try:
            poclist = list()
            self.loadmodule()
            sql = 'SELECT poc from hostexploit WHERE vulname like "%{}%"'.format(keyword)
            res = db().execute(sql)
            for item in res:
                poclist.append(item['poc'])
                self.queue.put_nowait(item['poc'])
            mylog('hostexploit', True).log.info(pyfancy().green('[+]针对目标:{0}:{1} 加载{2} hostpoc {3}个'.format(self.host, self.port, keyword, len(poclist))))
            threads = [gevent.spawn(self.Consumer, item) for item in poclist]
            gevent.joinall(threads)
            for vuln in self.vulns:
                sqlstr = 'INSERT INTO hostvulnlist (vulnhost, vulnport, vulnname, isvul, payload, proof, exception) VALUE ("{0}", "{1}", "{2}", "{3}", "{4}", "{5}", "{6}")'.format(vuln['vulnhost'], vuln['vulnport'], vuln['vulnname'], vuln['isvul'], pymysql.escape_string(str(vuln['payload'])), vuln['proof'], pymysql.escape_string(str(vuln['exception'])))
                db().execute(sqlstr)
                vuln = json.dumps(vuln, indent=4)
                mylog('hostexploit').log.debug(pyfancy().magenta('[*] {0}'.format(vuln)))
            self.vulns = []
        except Exception as e:
            mylog('hostexploit').log.critical(pyfancy().red(e))
