#!/usr/bin/python

'''
Transmit beacon request frames to the broadcast address while
channel hopping to identify ZigBee Coordinator/Router devices.
'''

import sys
import os
import signal
import time
import argparse

from killerbee import *

txcount = 0
rxcount = 0
stumbled = {}

def display_details(routerdata):
    global args, csvfile
    stackprofile_map = {0:"Network Specific",
                        1:"ZigBee Standard",
                        2:"ZigBee Enterprise"}
    stackver_map = {0:"ZigBee Prototype",
                    1:"ZigBee 2004",
                    2:"ZigBee 2006/2007"}
    spanid, source, extpanid, stackprofilever, channel = routerdata
    stackprofile = stackprofilever & 0x0f
    stackver = (stackprofilever & 0xf0) >>4

    print("New Network: PANID 0x{0:02X}{1:02X} Source 0x{2:02X}{3:02X}".format(spanid[0], spanid[1], source[0], source[1]))

    try:
        extpanidstr=""
        for ind in range(0,7):
            extpanidstr += "%02x:"%extpanid[ind]
        extpanidstr += "%02X"%extpanid[-1]
        sys.stdout.write("\tExt PANID: " + extpanidstr)
    except IndexError:
        sys.stdout.write("\tExt PANID: Unknown")

    try:
        print("\tStack Profile: {0}".format(stackprofile_map[stackprofile]))
        stackprofilestr = stackprofile_map[stackprofile]
    except KeyError:
        print("\tStack Profile: Unknown ({0})".format(stackprofile))
        stackprofilestr = "Unknown (%d)"%stackprofile

    try:
        print(("\tStack Version: {0}".format(stackver_map[stackver])))
        stackverstr = stackver_map[stackprofile]
    except KeyError:
        print(("\tStack Version: Unknown ({0})".format(stackver)))
        stackverstr = "Unknown (%d)"%stackver

    print(("\tChannel: {0}".format(channel)))

    if args.csvfile is not None:
        csvfile.write("0x%02X%02X,0x%02X%02X,%s,%s,%s,%d\n"%(spanid[0], spanid[1], source[0], source[1], extpanidstr, stackprofilestr, stackverstr, channel))


def response_handler(stumbled, packet, channel):
    global args
    d154 = Dot154PacketParser()
    # Chop the packet up
    pktdecode = d154.pktchop(packet)

    # Byte-swap the frame control field
    fcf = struct.unpack("<H", pktdecode[0])[0]

    # Check if this is a beacon frame
    if (fcf & DOT154_FCF_TYPE_MASK) == DOT154_FCF_TYPE_BEACON:
        # The 6th element offset in the Dot154PacketParser.pktchop() method
        # contains the beacon data in its own list.  Extract the Ext PAN ID.
        spanid = pktdecode[4][::-1]
        source = pktdecode[5][::-1]
        beacondata = pktdecode[6]
        extpanid = beacondata[6][::-1]
        stackprofilever = beacondata[4]
        assocPermit = struct.unpack("<H", beacondata[0])[0] & 0x8000

        key = b''.join([spanid, source])
        value = [spanid, source, extpanid, stackprofilever, channel]
        if not key in stumbled:
            if args.verbose:
                if assocPermit:
                    print("Beacon represents new network - ### Permitting new associations ###.")
                else:
                    print("Beacon represents new network - not accepting new associations.")
            stumbled[key] = value
            display_details(value)
        else:
            if args.verbose:
                if assocPermit:
                    print("Received frame is a beacon - ### Permitting new associations ###.")
                else:
                    print("Received frame is a beacon - not accepting new associations.")
        return value

    if args.verbose:
        print("Received frame is not a beacon (FCF={0}).".format(pktdecode[0]))

    return None

def interrupt(signum, frame):
    global txcount, rxcount
    global kb
    global args, csvfile
    if args.csvfile is not None:
        csvfile.close()
    kb.close()
    print(("\n{0} packets transmitted, {1} responses.".format(txcount, rxcount)))
    sys.exit(0)

if __name__ == '__main__':
    # Command-line arguments
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('-i', '--iface', '--dev', action='store', dest='devstring')
    parser.add_argument('-g', '--gps', '--ignore', action='append', dest='ignore')
    parser.add_argument('-s', '--delay', action='store', type=int, dest='delay', default=2)
    parser.add_argument('-v', '--verbose', action='store_true')
    parser.add_argument('-c', '--channel', action='store', type=int, default=None)
    parser.add_argument('-w', '--file', action='store', dest='csvfile', default=None)
    parser.add_argument('-D', action='store_true', dest='showdev')
    args = parser.parse_args()

    if args.showdev:
        show_dev()
        sys.exit(0)

    if args.csvfile is not None:
        try:
            csvfile = open(args.csvfile, 'w')
        except Exception as e:
            print(("Issue opening CSV output file: {0}.".format(e)))
        csvfile.write("panid,source,extpanid,stackprofile,stackversion,channel\n")

    # Beacon frame
    beacon = b"\x03\x08\x00\xff\xff\xff\xff\x07"
    # Immutable strings - split beacon around sequence number field
    beaconp1 = beacon[0:2]
    beaconp2 = beacon[3:]

    try:
        kb = KillerBee(device=args.devstring)
    except KBInterfaceError as e:
        sys.stderr.write("Interface Error: {0}".format(e))
        sys.exit(-1)

    signal.signal(signal.SIGINT, interrupt)
    print(("zbstumbler: Transmitting and receiving on interface \'{0}\'".format(kb.get_dev_info()[0])))

    # Sequence number of beacon request frame
    seqnum = 0
    if args.channel:
        channel = args.channel
        kb.set_channel(channel)
    else:
        channel = 11

    # Loop injecting and receiving packets
    while 1:
        if channel > 26:
            channel = 11

        if seqnum > 255:
            seqnum = 0

        if not args.channel:
            if args.verbose:
                print(("Setting channel to {0}.".format(channel)))
            try:
                kb.set_channel(channel)
            except Exception as e:
                sys.stderr.write("ERROR: Failed to set channel to {0}. ({1})".format(channel, e))
                sys.exit(-1)

        if args.verbose:
            print("Transmitting beacon request.")

        beaconinj = b''.join([beaconp1, b"%c" % seqnum, beaconp2])

        # response frame.
        try:
            txcount+=1
            kb.inject(beaconinj)
        except Exception as e:
            sys.stderr.write("ERROR: Unable to inject packet: {0}".format(e))
            sys.exit(-1)

        recvpkt = kb.pnext(args.delay)
        # Check for empty packet (timeout) and valid FCS
        if recvpkt is not None and recvpkt[1]:
            rxcount += 1
            if args.verbose:
                print("Received frame.")
            networkdata = response_handler(stumbled, recvpkt[0], channel)

        kb.sniffer_off()
        seqnum += 1
        if not args.channel:
            channel += 1

