#!/usr/bin/env python

# 2012 Laurent Ghigonis <laurent@p1sec.com>
# 2012 Pierre-Olivier Vauboin <po@p1sec.com>

# TODO
# * [TOTEST] detect AP only if they send more bea than other
# * display UTF* or other charset bssid
# * display hex value when bssid is non-printable
# * Network View (who is conn, has ap ...)
# * store in redis Nodes and Networks
# * read pcap file
# * better UI
#   * implement a time to forget about nodes
#   * forget nodes on demand

import sys
import datetime
import getopt
import pdb #pdb.set_trace()
import pprint
import netaddr
from scapy.all import *

node_list = {}
# Dot11 Type and Subtype put to hex and concatenated
dot11_types = {
    "00": ["arq", "Association request"],
    "01": ["ars", "Association response"],
    "04": ["prq", "Probe request"],
    "05": ["prs", "Probe response"],
    "08": ["bea", "Beacon frame"],
    "0b": ["aut", "Authentication"],
    "0c": ["dea", "Deauthentication"],
    "0d": ["act", "Action"],
    "19": ["bac", "Block Ack"],
    "1b": ["cts", "Clear to Send"],
    "1d": ["ack", "Acknowledgement"],
    "20": ["dat", "Data"],
    "24": ["nul", "Null function"],
    "28": ["qda", "QoS Data"],
}
packet_counter=0
start_time = datetime.datetime.now()
last_time = 0
option_ap = False
option_debug = False
option_iface = "wlan0"
option_update_time = 3
option_mac_target = ""

def usage():
  print "usage: %s [-ad] [-m mac_target] iface" % sys.argv[0]

def dot11_type_name(num):
  if num in dot11_types:
    return dot11_types[num][0]
  else:
    return "%s " % num

def dump_(obj, classkey=None):
  if isinstance(obj, dict):
      for k in obj.keys():
          obj[k] = dump_(obj[k], classkey)
      return obj
  elif hasattr(obj, "__iter__"):
      return [dump_(v, classkey) for v in obj]
  elif hasattr(obj, "__dict__"):
      data = dict([(key, dump_(value, classkey)) 
          for key, value in obj.__dict__.iteritems() 
          if not callable(value) and not key.startswith('_')])
      if classkey is not None and hasattr(obj, "__class__"):
          data[classkey] = obj.__class__.__name__
      return data
  else:
      return obj
def dump(obj, classkey=None):
  pprint.pprint(dump_(obj, classkey))

def mac_lookup(mac):
  mac_prefix = '-'.join(mac.replace(':', '-').split('-')[0:3])
  try:
    oui = netaddr.OUI(mac_prefix)
    if oui.records:
      vendor = oui.records[0]['org']
    else:
      vendor = 'Unknown'
  except netaddr.NotRegisteredError:
    vendor = 'Unknown'
  return vendor

class Node:
  def __init__(self, mac):
    self.mac = mac
    self.types = {}
    self.pkt_recv_count = 0
  def addpkt(self, pkt):
    self.pkt_recv_count += 1
    type1 = pkt[Dot11].type
    type2 = pkt[Dot11].subtype
    ftype = "%x%x" % (type1, type2)
    if ftype not in self.types:
      self.types[ftype] = {"count": 0}
    self.types[ftype]["count"] += 1
    if ftype in ["00", "04", "05", "08"]:
      if "ssid" not in self.types[ftype]:
        self.types[ftype]["ssid"] = set()
      ssid = pkt[Dot11Elt].info if Dot11Elt in pkt else "#unknown"
      if (pkt.addr2):
        if (pkt.addr3):
          Network.update(ssid, pkt.addr3)
        else:
          Network.update(ssid, pkt.addr2)
      self.types[ftype]["ssid"].add(ssid)
    else:
      bssid = pkt[Dot11].addr3
      if bssid:
        if "bssid" not in self.types[ftype]:
          self.types[ftype]["bssid"] = set()
        self.types[ftype]["bssid"].add(bssid)
  def is_ap(self):
    if '08' in self.types.keys():
      na = len(self.types['08'])
      for t in self.types.keys():
        if len(self.types[t]) > na:
          return False
      return True
    return False
  def __str__(self):
    s = ""
    s += self.mac + " %s (%s)" % (mac_lookup(self.mac), self.pkt_recv_count)
    l = sorted(self.types.keys(), key=lambda x: self.types[x]["count"], reverse=True)
    for t in l:
      s += "\n"
      s += "  " + str(self.types[t]["count"]) + "\t" + dot11_type_name(t) + " "
      for i in self.types[t]:
        if i == "count": continue
        if i == "bssid":
          s += ','.join(map(Network.resolv, self.types[t][i]))
        else:
          s += ','.join(self.types[t][i])
    return s

class Network:
  _nets = {}
  @classmethod
  def resolv(cls, bssid):
    for n in cls._nets.values():
      if bssid in n.bssid:
        return n.name
    return bssid
  @classmethod
  def update(cls, name, bssid):
    if name in cls._nets.keys():
      cls._nets[name].bssid.add(bssid)
    else:
      n = Network(name)
      n.addbssid(bssid)
      cls._nets[name] = n
  def __init__(self, name):
    self.name = name
    self.bssid = set()
    Network._nets[name] = self
  def addbssid(self, bssid):
    self.bssid.add(bssid)

def analyse_packet(pkt):
  # pkt.addr1 : bssid
  # pkt.addr2 : station
  # pkt.addr3 : 
  global packet_counter
  global start_time, last_time
  packet_counter+=1
  try:
    if pkt.addr2 is None:
      if option_debug:
        print "source None !!!"
      return
  except:
    if option_debug:
      print "no addr2 !!!"
    return
  if option_mac_target == "" or option_mac_target == pkt.addr2:
    if pkt.addr2 not in node_list:
      n = Node(pkt.addr2)
      node_list[pkt.addr2] = n
    node_list[pkt.addr2].addpkt(pkt)
  #pdb.set_trace()
  if packet_counter % 20 == 0:
    time = datetime.datetime.now()
    if time.second != last_time:
      delta = time - start_time
      if delta.seconds % option_update_time == 0:
        last_time = time.second
        print_stats()

def print_stats():
  print "==================================== %d pkts, %d nodes" % (packet_counter, len(node_list))
  l = sorted(node_list.keys(), key=lambda x: node_list[x].pkt_recv_count, reverse=True)
  for cli in l:
    if not option_ap and node_list[cli].is_ap():
      continue
    print node_list[cli]

if __name__ == '__main__':
  try:
    opts, args = getopt.getopt(sys.argv[1:], "adm:")
  except getopt.GetoptError, err:
    usage()
    sys.exit(1)
  for o, a in opts:
    if o == "-a": option_ap=True
    elif o == "-d": option_debug=True
    elif o == "-m": option_mac_target=a
    else:
      print "unknown option %s" % o
      sys.exit(1)
  if len(args) < 1:
    usage()
    sys.exit(1)

  option_iface = args[0]
  print "Be sure that %s is in monitor mode." % option_iface
  print "Use airodump at the same time if you want channel hopping. hop hop !"
  sniff(iface=option_iface, prn=analyse_packet, store=0)