#!/usr/bin/env python3
"""Wifi Exploit Finder

Scan wifi channels using an 802.11x adapter capable of monitor mode. If an exploitable device is found, an exploit message is sent to the hiprfisr.

For FISSURE database identification, add the following row to the `attacks` table.

"protocol","attack_name","modulation_type","hardware","attack_type","filename","category_name","version"
"802.11x","Exploit_Finder","Default","802.11x Adapter","Python3 Script","Wifi_Exploit_Finder.py","Sniffing/Snooping","maint-3.10"

References:

- [1] https://community.cisco.com/t5/wireless-mobility-knowledge-base/802-11-frames-a-starter-guide-to-learn-wireless-sniffer-traces/ta-p/3110019
"""
import sys
import subprocess
import time
import csv
import numpy as np
import json
from datetime import datetime, timezone
import os
from fissure.utils import FLOW_GRAPH_LIBRARY_3_8
from typing import List


DURATION = -1 # to run until interrupted
DWELL = 1 # 1 second per channel
NOTES = 'Wifi exploitable device finder'

# vendor plain text (reduced from variants)
# (search term, plain output)
VENDOR_PLAIN = [
    ('d-link', 'd-link'),
    ('tp-link', 'tp-link')
]

def apply_change(dev:str,commands:List[List[str]]):
    subprocess.run(['sudo','ip','link','set',dev,'down'])
    for command in commands:
        subprocess.run(command)
    subprocess.run(['sudo','ip','link','set',dev,'up'])

def set_monitor_mode(dev:str):
    apply_change(dev,[
        ['sudo','iwconfig',dev,'mode','Monitor']
    ])

def set_freq(dev:str,freq:float):
    apply_change(dev,[
        ['sudo','iwconfig',dev,'freq',str(freq)]
    ])

def set_channel(dev:str,channel:int):
    apply_change(dev,[
        ['sudo','iwconfig',dev,'channel',str(channel)]
    ])

def get_channels(dev:str)->dict:
    output = subprocess.run(['sudo','iwlist',dev,'frequency'], capture_output=True)
    output = output.stdout.decode('utf-8')
    output = output.split('\n')
    channels = {}
    for line in output:
        if 'Current' in line:
            continue
        if 'Channel' in line:
            line = line[line.index('Channel')+7:]
            channel = int(line[:line.index(':')])
            frequency = float(line[line.index(':')+1:line.index('GHz')])
            channels[channel] = frequency
    return channels

class OUILookup(object):
    def __init__(self):
        self._table = {}
        oui = os.path.join(FLOW_GRAPH_LIBRARY_3_8, "Single-Stage Flow Graphs", "resources", "oui.csv")
        with open(oui, "r") as ouifile:
            ouireader = csv.reader(ouifile)
            for lines in ouireader:
                self._table[lines[1]] = lines[2]

    def match(self, mac):
        tmp_mac = mac.replace(":", "").upper()
        found_mac = tmp_mac[:6]
        entry = self._table.get(found_mac)
        return entry

class Scan(object):
    def __init__(self, dev:str, channels:List[int]=None, duration:int=DURATION, dwell:int=DWELL, power_limit:float=-np.inf):
        self.dev = dev
        self.duration = duration
        self.dwell = dwell
        self.power_limit = power_limit

        if channels is None:
            # get available channels
            self.channels = list(get_channels(dev).keys())
        else:
            self.channels = channels

        # set monitor mode
        set_monitor_mode(dev)

        # initialize wireless ap address lookup
        self.aps = {}

        # create oui lookup table
        self.oui = OUILookup()

        self.fields_ordered = [
            'frame.time_epoch',
            'radiotap.channel.freq',
            'wlan.sa',
            'wlan.ta',
            'wlan.ra',
            'wlan.da',
            'radiotap.dbm_antsignal',
            'wlan.ssid',
            'wlan.fc.type_subtype'
        ]
        self.nfields = len(self.fields_ordered)
        self.fields = {field.split('.')[-1]: i for (i, field) in enumerate(self.fields_ordered)}
        self.cmd = ['sudo', 'tshark', '-i', dev, '-E', 'separator=,', '-E', 'occurrence=f', '-a', 'duration:' + str(dwell), '-Tfields']
        for field in self.fields_ordered:
            self.cmd += ['-e', field]

    def add_ap(self, mac, freq, row):
        if not mac in list(self.aps.keys()):
            # get vendor
            vendor = self.oui.match(mac)

            # get ssid
            ssid_raw = row[self.fields.get('ssid')]
            if ssid_raw == '<MISSING>':
                ssid = '<MISSING>'
            else:
                ssid = ''
                for i in range(0,len(ssid_raw),2):
                    ssid += chr(int(ssid_raw[i:i+2], 16))

            # get vendor plain text
            vendor_plain = None
            if not vendor is None:
                for vplain in VENDOR_PLAIN:
                    if vplain[0] in vendor.lower():
                        vendor_plain = vplain[1]
                        # add to AP table (unindent to avoid filtering on these)
                        self.aps[mac] = {
                            'vendor': vendor,
                            'vendor_plain': vendor_plain,
                            'ssid': ssid,
                            'stations': {}
                        }
                        sys.stdout.write(json.dumps({
                            'msg': 'alert',
                            'iface': self.dev,
                            'text': f'Exploitable {vendor_plain} Wifi AP Detected: AP MAC={mac}'
                        }) + '\n')
                        sys.stdout.flush()

                        if vendor_plain == 'tp-link':
                            sys.stdout.write(json.dumps({
                                'msg': 'exploit',
                                'protocol': '802.11x',
                                'modulation': 'Default',
                                'hardware': '802.11x Adapter',
                                'type': 'Python3 Script',
                                'attack': 'TP-Link Archer A7 Exploit',
                                'variables': json.dumps({}, separators=(',', ':'))
                            }) + '\n')
                            sys.stdout.write(json.dumps({
                                'msg': 'exploit',
                                'protocol': '802.11x',
                                'modulation': 'Default',
                                'hardware': '802.11x Adapter',
                                'type': 'Python3 Script',
                                'attack': 'TP-Link Archer A7 Reboot',
                                'variables': json.dumps({}, separators=(',', ':'))
                            }) + '\n')
                            sys.stdout.write(json.dumps({
                                'msg': 'tak',
                                'uid': mac,
                                'lat': '%(latitude)f',
                                'lon': '%(longitude)f',
                                'alt': '%(altitude)f',
                                'time': datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.%fZ'),
                                'remarks': f'Exploitable TP-Link router detected: AP MAC={mac}, Channel={int(freq)}M'
                            }) + '\n')
                            # tactical report
                            sys.stdout.write(json.dumps({
                                'msg': 'snreport',
                                'text': [ # https://www.globalsecurity.org/intell/library/policy/army/fm/34-35/appc.htm#figc_5
                                    'UNCLAS',
                                    f"MSGID/TACREP/FISSURE REPORT//", # TACREP=Tactical Report
                                    f"GNDOP/{datetime.now(timezone.utc).strftime('%H%M%SZ')}/1/US/ENEMY COMBATANT/ROUTER:{mac}//", # GNDOP=Ground Operations
                                    'LOCATION/TAMPA,FL/%(latitude_ddm)s%(longitude_ddm)s//', # City as city,state or city,country and Location in DM format
                                    f"COMEW/{freq:0.0f}MHZ/20MHZ//", # COMEW=Communications Electronic Warfare
                                ]
                            }) + '\n')
                            sys.stdout.flush()
                        break

    def add_station(self, mac, mac_ap, freq, row):
        if len(mac)==17 and\
        mac_ap in self.aps.keys() and\
        not mac == 'ff:ff:ff:ff:ff:ff' and\
        not mac in self.aps[mac_ap]['stations'].keys():
            # found new station
            self.aps[mac_ap]['stations'][mac] = freq

            sys.stdout.write(json.dumps({
                'msg': 'alert',
                'iface': self.dev,
                'text': f"Wifi {self.aps[mac_ap].get('vendor_plain')} Exploit Detected: AP MAC={mac_ap} Station MAC={mac}"
            }) + '\n')
            sys.stdout.flush()

            if self.aps[mac_ap].get('vendor_plain') == 'd-link':
                sys.stdout.write(json.dumps({
                    'msg': 'exploit',
                    'protocol': '802.11x',
                    'modulation': 'Default',
                    'hardware': '802.11x Adapter',
                    'type': 'Python3 Script',
                    'attack': 'DIR-815 Exploit',#'Wifi_DIR_815_Exploit.py',
                    'variables': json.dumps({
                        'ap_mac': ''.join(mac_ap.split(':')),
                        'sta_source_mac': ''.join(mac.split(':')),
                        'channel': f'{int(freq)}M'
                    }, separators=(',', ':'))
                }) + '\n')
                sys.stdout.write(json.dumps({
                    'msg': 'tak',
                    'uid': mac_ap,
                    'lat': '%(latitude)f',
                    'lon': '%(longitude)f',
                    'alt': '%(altitude)f',
                    'time': datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.%fZ'),
                    'remarks': f'Exploitable D-Link router detected: AP MAC={mac_ap}, Station MAC={mac}, Channel={int(freq)}M'
                }) + '\n')
            sys.stdout.flush()

    def run(self):
        tend = np.inf if self.duration == -1 else time.time() + self.duration
        while True:
            newta = False
            for channel in self.channels:
                # set channel
                set_channel(self.dev, channel)

                if time.time() + self.dwell > tend:
                    # stop scan
                    return self.aps

                # capture
                output = subprocess.run(self.cmd, capture_output=True)
                output = output.stdout.decode('utf-8')
                reader = csv.reader(output.split('\n'),escapechar="\\")
                for row in reader:
                    if len(row) == self.nfields: # populated line
                        power = row[self.fields.get('dbm_antsignal')]
                        power = self.power_limit - 1 if len(power)==0 else int(power.split(',')[0])
                        if power < self.power_limit:
                            continue

                        mac_sa = row[self.fields.get('sa')]
                        mac_ta = row[self.fields.get('ta')]
                        mac_ra = row[self.fields.get('ra')]
                        mac_da = row[self.fields.get('da')]
                        freq = float(row[self.fields.get('freq')])

                        type_subtype = int(row[self.fields.get('type_subtype')], 16)
                        if type_subtype == 8: # beacon
                            self.add_ap(mac_sa, freq, row)
                            continue

                        elif type_subtype == 0: # association request
                            self.add_ap(mac_da, freq, row)

                        elif type_subtype >= 24: # non-management frame
                            self.add_station(mac_sa, mac_da, freq, row)
                            self.add_station(mac_sa, mac_ra, freq, row)
                            self.add_station(mac_ta, mac_da, freq, row)
                            self.add_station(mac_ta, mac_ra, freq, row)
                            self.add_station(mac_ra, mac_sa, freq, row)
                            self.add_station(mac_ra, mac_ta, freq, row)
                            self.add_station(mac_da, mac_sa, freq, row)
                            self.add_station(mac_da, mac_ta, freq, row)

            if not newta:
                sys.stdout.write(time.strftime("%Y-%m-%d %H:%M:%S") + ' - No new mac addresses detected in channel scan\n')
                sys.stdout.flush()

def getArguments():
    iface = None
    duration=DURATION
    dwell=DWELL
    power = -100
    channels = None
    notes = NOTES

    return (
        ['iface', 'duration', 'dwell', 'power', 'channels', 'notes'],
        [iface, duration, dwell, power, channels, notes]
    )

if __name__ == '__main__':
    # get default values
    (iface, duration, dwell, power, channels, notes) = getArguments()[1]

    # handle input
    nargs = len(sys.argv)
    iface = sys.argv[1] if nargs > 1 else iface
    duration = int(sys.argv[2]) if nargs > 2 else duration
    dwell = int(sys.argv[3]) if nargs > 3 else dwell
    power = int(sys.argv[4]) if nargs > 4 else power
    channels = sys.argv[5] if nargs > 5 else channels
    channels = None if channels in [None, 'None'] else [int(c) for c in channels.split(',')]

    '''while True:
        sys.stdout.write(json.dumps({
            'msg': 'alert',
            'iface': iface,
            'text': f'Test'
        }) + '\n')
        sys.stdout.flush()
        time.sleep(1)'''

    # run
    scan = Scan(iface, channels, duration, dwell, power)
    scan.run()
