main

#!/usr/bin/env python3

from datetime import datetime, timezone
from dateutil import tz
from scapy.all import *
from time import sleep
from threading import Thread
from pprint import pprint
import atexit
import argparse
import csv
import dateutil.parser
import gps
import os
import sqlite3

parser = argparse.ArgumentParser()

parser.add_argument("--interface",
                      default='',
                      required=True,
                      help="Interface of choice")
parser.add_argument("--nomanagement",
                      default=False,
                      action='store_true', 
                      help="Assume another program is already handling int setup/restoration and channel switching. We're just along for the ride.")
parser.add_argument("--importcsv",
                      default='', 
                      help="CSV file to import")
parser.add_argument("--exportcsv",
                      default='', 
                      help="CSV filename to export with. If --google used this is a basename and CSVs get saved to your CWD")
parser.add_argument("--google",
                      default='',
                      action='store_true', 
                      help="Break CSV export into multiple files of 2000 data point each. For Google's 'My Maps' which has a 2000-per-csv limit.")
parser.add_argument("--skipgps",
                      default='', 
                      action='store_true', 
                      help="Don't wait for GPS info. For developing this script without a GPS present.")
parser.add_argument("--loglevel",
                      default=0,
                      type=int,
                      help="Optional verbosity. Verbose=1, Debug=2")
parser.add_argument("--debug",
                      default='',
                      action='store_const',
                      dest='loglevel',
                      const=2,
                      help="Quick --loglevel 2")
parser.add_argument("--verbose",
                      default='',
                      action='store_const',
                      dest='loglevel',
                      const=1,
                      help="Quick --loglevel 1")


args = parser.parse_args()

#   ____ _
#  / ___| | __ _ ___ ___  ___  ___
# | |   | |/ _` / __/ __|/ _ \/ __|
# | |___| | (_| \__ \__ \  __/\__ \
#  \____|_|\__,_|___/___/\___||___/
#

class GPS(): # Generate GPS data for other threads.
  def __init__(self):
    self.gpsData = { "gpsTimestamp": None, "gpsTimestampAge": None, "Longitude": None, "Latitude": None } # Init gpsData
      # Longitude/Latitude   The location data received at gpsTimestamp.
      # gpsTimestamp         The timestamp with Lon / Lat from GPS
      # gpsTimestampAge      The difference in seconds between your host's time and the gpsTimestamp stamp.
      #                      Useful to know if your GPS has lost its lock when writing results.

    try:
      print('Initializing GPS...')
      self.session = gps.gps(mode=gps.WATCH_ENABLE)
    except ConnectionRefusedError:
      print('Local gpsd refused our connection. Is it running?')
      print('If you are debugging with no GPS, you can launch again using --skipgps')
      exit(1)
    except Exception as e:
      print('Something else went wrong trying to speak with gpsd:')
      print(e)
      exit(1)
    finally:
      print('Initialized.')

  def readGps(self):
    time=lat=lon=None

      # Do initial read from gpsd
    self.session.read()

      # If invalid TPV result, keep trying until we get one...
      # May be inappropriate to gpsd.
    while not (gps.MODE_SET & self.session.valid):
      self.session.read() # Try reading from gpsd.


      # Capture GPS state for data purposes.
      # Not Yet Implemented!!!!!!!!
    #gpsFixStatus = ("Invalid", "NO_FIX", "2D", "3D")[session.fix.mode]

    if gps.TIME_SET & self.session.valid:
      time = self.session.fix.time

    if ((gps.isfinite(self.session.fix.latitude) and gps.isfinite(self.session.fix.longitude))):
      lat = self.session.fix.latitude
      lon = self.session.fix.longitude

    if time and lat and lon:
      self.gpsData["gpsTimestamp"] = time
      self.gpsData["Latitude"]     = lat
      self.gpsData["Longitude"]    = lon

    if self.gpsData['gpsTimestamp']: # Update staleness periodically
      gpsTime = datetime.strptime(str(self.gpsData['gpsTimestamp']),'%Y-%m-%dT%H:%M:%S.%f%z')
      gpsTime = gpsTime.replace(tzinfo=tz.tzutc()).astimezone(tz=tz.tzlocal())
      localTime = datetime.now().astimezone()
      self.gpsData["gpsTimestampAge"] = ((localTime-gpsTime).total_seconds())

  def _gpsThread(self):
    while True:
      self.readGps()

  def startGpsThread(self):
    try:
      self.gpsThread = Thread(target=self._gpsThread)
      self.gpsThread.daemon = True
      self.gpsThread.start()
    finally:
      print('GPS Data thread started.')

  def waitForLock(self):
    if args.skipgps:
      self.gpsData['gpsTimestamp']    = 'NA';
      self.gpsData['gpsTimestampAge'] = 'NA';
      self.gpsData['Latitude']        = 0;
      self.gpsData['Longitude']       = 0;
      print("Skipping GPS wait and set dummy values: " + str(list(gpsData.values())))
    else:
      print('Waiting for GPS lock...')
    
      while not self.gpsData['gpsTimestamp']: # Wait for data
        self.readGps()

      print("Got a lock: " + str(list(self.gpsData.values())))



class Database: # Our database object
  def __init__(self, databaseFile):
    try:
      self.con = sqlite3.connect(databaseFile,check_same_thread=False)
      self.cur = self.con.cursor()
        # Create our database and tables, including a secondary table to evict old data.

        # Dict-powered DB creation. Probably a stupid idea
      schema = {  "columns": {
                    "bssid": {
                      "type": "VARCHAR(17)"
                    },
                    "ssid": {
                      "type": ""
                     },
                    "channel": {
                      "type": "INT"
                    },
                    "security": {
                      "type": "TEXT"
                    },
                    "rssi": {
                      "type": "INT"
                    },
                    "timestamp": {
                      "type": "INT"
                    },
                    "timestampAge": {
                      "type": "REAL"
                    },
                    "lat": {
                      "type": "REAL"
                    },
                    "lon": {
                      "type": "REAL"
                    }
                  },
                 "tables": {
                   "wifi": {
                     "tableSpecial": "unique (bssid)"
                   },
                   "wifi_evicted": {
                     "tableSpecial": "" 
                   }
                 }
               }

      for table in list(schema['tables'].keys()): # Build the database
        tableQuery = "create table if not exists %s" %(table)

        columnCount = len(schema['columns'])
        counter = 1

        columnQuery = ''
        for column in schema['columns']:
          columnQuery += column + ' ' + schema['columns'][column]['type']
          if counter < 2:
            columnQuery = ' ' + columnQuery
          if counter < columnCount:
            columnQuery += ', '
          else:
            columnQuery += ' '

          counter += 1

        if schema['tables'][table]['tableSpecial']:
          columnQuery += ', ' + schema['tables'][table]['tableSpecial']

        tableQuery += ' (' + columnQuery + ')'
        self.exec(tableQuery)

    except Exception as e:
      print('Failed to prepare Database: ', e)
      exit(1)
    finally:
      print('Database connected')

  def exec(self, query):
    logger('About to execute this query: ' + '[' + query + ']',2)
    self.cur.execute(query)
    self.con.commit()

  def execFetchAll(self, query):
    self.exec(query)
    return(self.cur.fetchall())

  def checkBssidExists(self,bssid): # Check if we have a point for a given BSSID already.
    self.exec("select * from wifi where bssid = '%s' limit 1" %(bssid))
    result = self.cur.fetchone()
    if result: # Truthy check for whether we got any result at all.
      logger('BSSID Exists in database already',2)
      return(result)
    else:
      logger('BSSID is new',2)
      return(False)

  def checkNewPointBetter(self,storeCheck,bssid,ssid,rssi,channel,security,timestamp,timestampAge,lat,lon): # Check if we're replacing stored data.
    if int(rssi) > storeCheck[2]:
      logger('RSSI improved',2)
      return(True)
    else:
      logger('No RSSI improvement, leaving alone.',2)
      return(False)

  def retirePoint(self,bssid):
    self.exec("INSERT OR REPLACE INTO wifi_evicted select * from wifi where bssid = '%s'" %(bssid))
    self.exec("DELETE from wifi where bssid = '%s'" %(bssid))

  def writePoint(self,bssid,ssid,rssi,channel,security,timestamp,timestampAge,lat,lon):
    if args.skipgps:
      print('Skipping point write due to --skipgps: ' + bssid + ' ' + ssid)
    else:
      ssid = ssid.replace("'","''")
      self.exec("insert or ignore into wifi values ('%s','%s','%s','%s','%s','%s','%s','%s','%s')" %(bssid,ssid,rssi,channel,security,timestamp,timestampAge,lat,lon))

  def handleWifiPoint(self,bssid,ssid,rssi,channel,security,timestamp,timestampAge,lat,lon):
    storeCheck = self.checkBssidExists(bssid)
    if storeCheck:
        if self.checkNewPointBetter(storeCheck,bssid,ssid,rssi,channel,security,timestamp,timestampAge,lat,lon):
          logger('Updating point:\t' + bssid + '\tStaleness:\t' + str(timestampAge), 1)
          self.retirePoint(bssid)
          self.writePoint(bssid,ssid,rssi,channel,security,timestamp,timestampAge,lat,lon)
    else:
        logger('Writing new point:\t' + bssid + '\tStaleness:\t' + str(timestampAge), 1)
        self.writePoint(bssid,ssid,rssi,channel,security,timestamp,timestampAge,lat,lon)

# _____                 _   _
#|  ___|   _ _ __   ___| |_(_) ___  _ __  ___
#| |_ | | | | '_ \ / __| __| |/ _ \| '_ \/ __|
#|  _|| |_| | | | | (__| |_| | (_) | | | \__ \
#|_|   \__,_|_| |_|\___|\__|_|\___/|_| |_|___/
#

def cleanup():
  print('Cleaning up...')
  manageInt(args,'managed')

def logger(content,loglevel=0):
  loglevelString = '[' + ("info", "verbose", "debug")[loglevel] + '] '
  if loglevel <= args.loglevel:
    print(loglevelString + content)

def date2unix(date):
  if type(date) == int: # If already a number, assume it's already a unix timestamp
    return(date)
  else: # Otherwise, try reading it
    try:
      result = str(dateutil.parser.parse(date).timestamp()).split('.')[0]
    except:
      return(False)
    finally:
      return(int(result))

def importCsv(args):
  if os.path.isfile(args.importcsv):
    with open(args.importcsv) as csvfile:
      reader = csv.DictReader(csvfile)
      for row in reader:
        if row['bssid'] and row['ssid'] and row['Latitude'] and row['Longitude']: # Check for these during import at a minimum

          print('Importing: ' + row['bssid'] + ' to our database.')
          if row['gpsTimestamp']: # Convert the timestamp to UNIX
            try:
              gpsTimestamp = date2unix(row['gpsTimestamp'])
            except:
              print('Failed to convert date...')
              gpsTimestamp = row['gpsTimestamp']

          Db.handleWifiPoint(row['bssid'],
                             row['ssid'],
                             row['rssi'],
                             row['channel'],
                             row['security'],
                             gpsTimestamp,
                             row['gpsTimestampAge'],
                             row['Latitude'],
                             row['Longitude'])

        else:
          print('Failed to import: ' + row['bssid'] + '/' + row['ssid'] + ' to our database.')
          print('We need at least bssid,ssid,Latitude,Longitide to import to the database.')
          print('If this datapoint has all of these, please check your CSV for errors!')

  else:
    print('Argument given is not a file')
    exit(1)

  exit(0)

def writeFile(filename,headers,items):
  try:
    with open(filename, 'w') as csvFile:
      w = csv.DictWriter(csvFile, headers)
      w = csv.writer(csvFile)

        # We can either select the data out converting the timestamp row (select datetime(timestamp, 'auto')) for example
        # Or we can take the tuple row here and convert it in python to a local timezone. This would easily specify the +TZ as well.
        # e.g. date = print(time.replace(tzinfo=tz.tzutc()).astimezone(tz=tz.tzlocal()))
        # Make a unix2date function?
      w.writerow(headers) # Write the header
      for row in items:
          # for csv consistency, export SSID as unicode escaped rather than interpreting newlines and linefeed chars literally.
        row = list(row) # To modify the tuple result
        row[1] = row[1].encode('unicode_escape')
        row[1] = row[1].decode('utf-8')
        w.writerow(row) # Write the data
  finally:
    print('Written: ' + filename)

def exportCsv(args):
  filename = args.exportcsv
  headerQuery = Db.execFetchAll('PRAGMA table_info(wifi)')
  headers = []
  for column in headerQuery: # Put headers together for the csv.
    headers.append(column[1])

    # Read out our data into memory
  wifiQuery = Db.execFetchAll('select * from wifi')

  if args.google: # Limit each csv to 2000 points each.
    split = 2000
    wifiQueryCount = len(wifiQuery) - 1
    iter = 0
    while iter < wifiQueryCount:
      print(str(iter) + ':' + str(wifiQueryCount+2000))
        # Insert an incrementing 'split' number into the desired filename before the suffix for visibility.
      splitFilename = str('.'.join(filename.split('.')[:-1]) + '-' + str(iter) + '.' + filename.split('.')[-1])
      writeFile(splitFilename,headers,wifiQuery[iter:iter+split])
      iter = iter + split
  else: # Just write the csv filename specified.
    writeFile(filename,headers,wifiQuery)

  exit(0)

def processPacket(packet):
  try:
    if packet.haslayer(Dot11Beacon): # If the packet looks good, proceed
      bssid = packet[Dot11].addr2
      ssid  = packet[Dot11Elt].info.decode().rstrip('\x00')
      if len(ssid) == 0: ssid = "<Hidden>"
      rssi = packet.dBm_AntSignal
  
      stats = packet[Dot11Beacon].network_stats()
      channel  = stats.get("channel")
      security = ''.join(stats.get("crypto"))
      
      # Handle new and seen APs. 
      # If new:  Add to array.
      # If seen: Compare RSSI, GPS Coordinates if better.
      
      logger('Processing: ' + bssid + str(ssid) + str(channel) + security + str(rssi),2)
      Db.handleWifiPoint(bssid,
                         ssid,
                         rssi,
                         channel,
                         security,
                         Gps.gpsData['gpsTimestamp'],
                         Gps.gpsData['gpsTimestampAge'],
                         Gps.gpsData['Latitude'],
                         Gps.gpsData['Longitude'])

  except Exception as e: # Ignore busted packets
    print('Packet issue: ' + str(e))
    pass

def executeCmd(cmd):
  try: 
    subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True, shell=True)
  except subprocess.CalledProcessError as e:
    print('Command [' + cmd + '] failed: ' + e.stderr.decode("utf-8"))

def manageInt(args,mode):
  if args.nomanagement:
    print('--nomanagement specified, not touching ' + args.interface)
    return
  if mode == 'monitor':
    txmode = 'fixed 3000'
  elif mode == 'managed':
    txmode = 'auto'

  try:
    executeCmd('ip link set ' + args.interface + ' down')
    executeCmd('iw ' + args.interface + ' set type ' + mode)
    executeCmd('ip link set ' + args.interface + ' up')
    # Not all interfaces like being told to go into auto
    executeCmd('iw ' + args.interface + ' set txpower ' + txmode)
  except Exception as e:
    print('Something went wrong prepping interface: ' + args.interface)
    print(e)
  finally:
    print(args.interface + ' is now a ' + mode + ' interface.')


def snifferThread(args):
  try:
    sniff(prn=processPacket, iface=args.interface)
  except PermissionError:
    print('No access to interface ' + args.interface + ' Please give this user access or run the script as root.')
    exit(1)
  except Exception as e:
    print('Interface trouble: ' + str(e))
    exit(1)

def channelSwitcher():
  print('Channel switching thread started...')
  ch = 1
  while True:
    executeCmd('iwconfig {interface} channel {ch}'.format(interface = args.interface, ch = ch))
    ch = ch % 14 + 1
    sleep(0.25)

#  ____
# / ___| ___
#| |  _ / _ \
#| |_| | (_) |
# \____|\___/
#

if __name__ == "__main__":
  scriptRoot = os.path.dirname(os.path.realpath(__file__))
  databaseFile = scriptRoot + '/wifi.db'

  Db = Database(databaseFile)

  if args.importcsv:
    importCsv(args)

  if args.exportcsv:
    exportCsv(args)

  logger('Loglevel is: ' + str(args.loglevel),1)

  Gps = GPS()
  Gps.waitForLock()    # Don't take off until we get a lock.

  Gps.startGpsThread() # Get started

  manageInt(args,'monitor')
  atexit.register(cleanup)

  if not args.nomanagement:
    channelSwitcher = Thread(target=channelSwitcher)
    channelSwitcher.daemon = True
    channelSwitcher.start()

  print('Beginning recon...')
  snifferThread = Thread(target=snifferThread(args))
  snifferThread.daemon = True
  snifferThread.start()