/******************************************************************************
**
** Attempts to establish multioboot with DS.
**
**   Author:  Gary Linscott (premandrake@hotmail.com)
**
**   Copyright (c) 2004 Gary Linscott, All Rights Reserved (see license info below). 
**
** Legal/Credits:
**
**   HUGE thanks to crazyc for the checksum code!
**   Heavily influenced by both the Airjack code and the Madwifi driver.
**
**   This program is free software; you can redistribute it and/or
**   modify it under the terms of the GNU General Public License
**   as published by the Free Software Foundation; either version 2
**   of the License, or (at your option) any later version.
**
**   This program is distributed in the hope that it will be useful,
**   but WITHOUT ANY WARRANTY; without even the implied warranty of
**   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
**   GNU General Public License for more details.
**
**   You should have received a copy of the GNU General Public License
**   along with this program; if not, write to the Free Software
**   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
**
******************************************************************************/

#include <arpa/inet.h>
#include <net/ethernet.h>
#include <net/if.h>
#include <net/if_arp.h>
#include <netinet/in.h>
#include <netpacket/packet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/uio.h>
#include <sys/ioctl.h>
#include <sys/timeb.h>
#include <errno.h>
#include <pcap.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef unsigned char __u8;
typedef unsigned short __u16;
typedef unsigned int __u32;
#include "80211.h"

struct llc_hdr {
    __u8        dsap;
    __u8        ssap;
    __u8        cntl;
    __u8        org_code[3];
    __u16       type;
} __attribute__ ((packed));

struct broadcastPacket {
    int length;
    char packet[512];
};

enum dsmbStates {
    DSMB_NOT_STARTED,
    DSMB_WAITING_FOR_AUTH,
    DSMB_WAITING_FOR_ASSOC,
    DSMB_WAITING_FOR_ACK,
    DSMB_SENDING_MB_LOADER
};

struct dsmbState {
    // Socket descriptor
    int sd;
    int ethIndex;
    // Are we currently advertising a game for download?
    int advertisingOn;
    int blockTime;
    int curTime;
    // What is the current beacon frame we are sending?
    int curBeacon;
    // The beacon frames
    struct broadcastPacket beaconPackets[10];
    // The current packet sequence number
    int seqNum;
    // The sockets to check for incoming data
    fd_set socketsToCheck;
    // Our MAC address
    __u8 MAC[6];
    // Client MAC address
    __u8 clientMAC[6];
    // Current state
    enum dsmbStates curState;
    
    int recvSendingACK;
    int packetsSentWithoutACK;
    int sentAuth,sentAssocReq;

    // Current state processing function
    void (*processState)(struct dsmbState *dsmb, __u8 *packet, int length);
};

// Packet description:
// First 24 bytes, 802.11B header.  0x14 = packet sequence # (2 bytes?)
// Next 12 bytes, fixed parameters, first 8 bytes, timestamp
// 50 bytes in, beginning of custom tag.  should be 0xDD (custom)
// next byte, length of tag, 0x88 for mario beacon packets

// From offset 52 from now on:
// 0x00: 6 bytes - 00 09 BF ?? ?? ?? - MAC?
// 0x06: 10 bytes - unknown
// 0x10: 2 bytes - Game ID (randomly initialized it appears, needed for client DS to authorize)
// 0x12: 10 bytes - unknown
// 0x1C: 1 byte - 0x02 on final packet, 0 otherwise
// 0x1D: 2 bytes - unknown
// 0x20: 1 byte - packet seq. number
// 0x21: 2 bytes - CRC
// 0x23: 1 byte - packet seq. number if not last packet, if last packet, number of players connected
// 0x24: 1 byte - total number of packets in message
// 0x25: 2 bytes - length of data segment (except in last packet for some reason)
// 0x27: x bytes - packet data

// Last packet contains name data for DS's connected to the game

struct mbPacketHeader {
    __u8 unk1[28];
    __u8 final;
    __u8 unk2[2];
    __u8 seqNum;
    __u16 checksum;
    __u8 seqNum2;
    __u8 numPackets;
    __u16 length;
    __u8 data[0x62];
} __attribute__ ((packed));

void initDsMbState(struct dsmbState *dsmb) {
    dsmb->sd = 0;
    dsmb->advertisingOn = 1;
    dsmb->curBeacon = 0;
    dsmb->seqNum = 40;
    dsmb->curState = DSMB_NOT_STARTED;
    
    const __u8 MAC[8] = {0x00, 0x09, 0xBF, 0x0B, 0x79, 0xB0};
    memcpy(dsmb->MAC, MAC, 6);
}

int initialize(char *device, int sd, int promisc) {
    struct ifreq ifr;
    
    strncpy(ifr.ifr_name, device, sizeof(ifr.ifr_name));
    if (ioctl(sd, SIOCGIFFLAGS, &ifr)<0) {
        perror("SIOCGIFFLAGS");
        return -1;
    }
    
    if (promisc)
        ifr.ifr_flags |= IFF_PROMISC;
    else
        ifr.ifr_flags &= ~IFF_PROMISC;
    
    if(ioctl(sd, SIOCSIFFLAGS, &ifr)<0) {
        perror("SIOCSIFFLAGS");
        return -1;
    }
    
    printf("Initialized device %s\n", device);
    
    if (ioctl(sd, SIOCGIFINDEX, &ifr)<0) {
        perror("SIOCGIFINDEX");
        return -1;
    }
    
    return ifr.ifr_ifindex;
}

void shutdownDevices() {
    printf("User shutdown\n");

    exit(0);
}

void printMAC(__u8 *mac) {
    int i;
    printf("%.2x", mac[0]);
    for (i = 1; i < ETH_ALEN; i++) {
        printf(":%.2x", mac[i]);
    }
}

int setupNetwork(struct dsmbState *dsmb) {
    // Set up raw sockets for sending
    dsmb->sd = socket(PF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
    if (dsmb->sd < 0) {
        printf("Couldn't create packet socket :(\n");
        return 0;
    }

    if ((dsmb->ethIndex = initialize("ath0", dsmb->sd, 1)) < 0) {
        printf("Couldn't find adapter.\n");
        return 0;
    }

    // Bind the socket to the interface
    struct sockaddr_ll etherSocket;
    memset(&etherSocket, 0, sizeof(etherSocket));
    etherSocket.sll_family = AF_PACKET;
    etherSocket.sll_protocol = htons(ETH_P_ALL);
    etherSocket.sll_ifindex = dsmb->ethIndex;

    if (bind(dsmb->sd, (struct sockaddr *)&etherSocket, sizeof(etherSocket)) < 0) {
        printf("Couldn't bind socket.\n");
        return 0;
    }
    
    return 1;
}

__u16 calcChecksum(struct broadcastPacket *packet) {
    int i;
    __u16 *p = (__u16*)&(packet->packet);
   
    __u32 checksum = 0;
    for (i = 0x11; i < 0x44; i++) {
//    for (i = 0x2B; i < 0x5E; i++) {
        checksum += p[i];
    }
    
    return (((checksum >> 16) + (checksum & 0xffff) + 1) * -1) & 0xffff;
}

int readBeacons(struct dsmbState *dsmb, char *filename) {
    FILE *file = fopen(filename, "rb");
    if (file == NULL) return 0;
    
    int i;
    for (i = 0; i < 10; i++) {
//        fread(&(dsmb->beaconPackets[i].length), 4, 1, file);
        dsmb->beaconPackets[i].length = 136;
        memset(dsmb->beaconPackets[i].packet, 0, sizeof(dsmb->beaconPackets[i].packet));
        fread(&(dsmb->beaconPackets[i].packet), dsmb->beaconPackets[i].length, 1, file);
    }
    
    for (i=0; i<10; i++) {
        struct mbPacketHeader *ph = (struct mbPacketHeader*)&(dsmb->beaconPackets[i].packet[0]);
                
        printf("final: %.2x, seqNum: %.2x, crc: %.4x, seqNum2: %.2x, num pkts: %.2x, length: %.4x --> ",
               ph->final, ph->seqNum, ph->checksum, ph->seqNum2, ph->numPackets, ph->length);        
               
        __u16 checksum = calcChecksum(&(dsmb->beaconPackets[i]));
        
        printf("Calced crc: %.4x - %s\n", checksum, checksum == ph->checksum ? "GOOD!" : "had to fix");
        
        ph->checksum = checksum;
    }      

    fclose(file);
    
    return 1;
}

int sendPacket(struct dsmbState *dsmb, void *data, int length) {
/*  struct sockaddr_ll etherSocket;

    memset(&etherSocket, 0, sizeof(etherSocket));
    etherSocket.sll_family = AF_PACKET;
    etherSocket.sll_ifindex = dsmb->ethIndex;
    memcpy(etherSocket.sll_addr, dsmb->MAC, 6);
    etherSocket.sll_halen = 6;
    
    if (sendto(dsmb->sd, data, length, 0, (struct sockaddr *)&etherSocket, sizeof(etherSocket)) == -1) {
        printf("Error in sending packet.\n");
        return 0;
    }*/
    
    if (send(dsmb->sd, data, length, 0) == -1) {
        printf("Error in sending packet.\n");
        return 0;
    }
    
    return 1;
}

int pollReadSocket(struct dsmbState *dsmb, int sec, int usec) {
    struct timespec timeout;
    
    fd_set socketsReady;
    memcpy(&socketsReady, &dsmb->socketsToCheck, sizeof(fd_set));

    timeout.tv_sec = 0;
    timeout.tv_nsec = 0;
    pselect(FOPEN_MAX, &socketsReady, NULL, NULL, &timeout, NULL);
    
    if (FD_ISSET(dsmb->sd, &socketsReady)) {
        // Data on this socket to be read
        return 1;
    }

    return 0;
}

void processRecievedPacket(struct dsmbState *dsmb, __u8 *packet, int length);

void processAuth(struct dsmbState *dsmb, __u8 *packet, int length) {
    struct a3_80211 *header = (struct a3_80211 *)packet;
    
    struct timespec delay;
    delay.tv_sec = 0;
    delay.tv_nsec = 15000;
    nanosleep(&delay, NULL);

    // Send AUTH response
    struct mgt_auth_one authSend;
    memcpy(&authSend, packet, sizeof(struct mgt_auth_one));

    authSend.at1_mac_hdr.mh_seq_num = dsmb->seqNum++;
    authSend.at1_seq++;
    authSend.at1_mac_hdr.mh_pwr_man = 0;
    // Set up our client MAC
    memcpy(dsmb->clientMAC, header->mh_mac2, 6);
    memcpy(authSend.at1_mac_hdr.mh_mac1, header->mh_mac2, 6);
    memcpy(authSend.at1_mac_hdr.mh_mac2, dsmb->MAC, 6);

    sendPacket(dsmb, &authSend, sizeof(authSend));        

    dsmb->curState = DSMB_WAITING_FOR_ASSOC;
}

void sendACK(struct dsmbState *dsmb) {
    __u8 sendRaw[sizeof(struct a1_80211)];
    memset(sendRaw, 0, sizeof(struct a1_80211));

    struct a1_80211 *ack = (struct a1_80211*)sendRaw;
    ack->mh_type = FC_TYPE_CTL;
    ack->mh_subtype = CTL_ACK;
    
    memcpy(ack->mh_mac1, dsmb->clientMAC, 6);

    sendPacket(dsmb, &sendRaw, sizeof(struct a1_80211)); 
}

void sendAssocReq(struct dsmbState *dsmb) {
    // Get the auth packet from the DS
    
    struct timespec delay;
    delay.tv_sec = 0;
    delay.tv_nsec = 1000;
    nanosleep(&delay, NULL);

    sendACK(dsmb);

//    struct timespec delay;
    delay.tv_sec = 0;
    delay.tv_nsec = 10000;
    nanosleep(&delay, NULL);

    {
        __u8 sendRaw[sizeof(struct mgt_association_resp_one) + 2];
        memset(sendRaw, 0, sizeof(struct a3_80211));

        struct mgt_association_resp_one *assocResp = (struct mgt_association_resp_one*)sendRaw;
    
        assocResp->as1_mac_hdr.mh_type = FC_TYPE_MGT;
        assocResp->as1_mac_hdr.mh_subtype = MGT_AS_RS;
        assocResp->as1_mac_hdr.mh_duration_id = 0xA2;
        memcpy(assocResp->as1_mac_hdr.mh_mac1, dsmb->clientMAC, 6);
        memcpy(assocResp->as1_mac_hdr.mh_mac2, dsmb->MAC, 6);
        memcpy(assocResp->as1_mac_hdr.mh_mac3, dsmb->MAC, 6);
        if (dsmb->sentAssocReq > 5) {
            dsmb->sentAssocReq = 0;
            dsmb->seqNum = 1;
        }
        if (dsmb->sentAssocReq == 0) {
            assocResp->as1_mac_hdr.mh_seq_num = dsmb->seqNum++;
        } else {
            assocResp->as1_mac_hdr.mh_seq_num = dsmb->seqNum;
            assocResp->as1_mac_hdr.mh_retry = 1;
        }
    
        assocResp->as1_cap.ci_ess = 1;
        assocResp->as1_cap.ci_shortPreamble = 1;
        assocResp->as1_status_code = 0;
        assocResp->as1_aid = 0xC001;
        assocResp->as1_rates_id.e_element_id = 1;
        assocResp->as1_rates_id.e_length = 2;
    
        sendRaw[sizeof(struct mgt_association_resp_one) + 0] = 0x82;
        sendRaw[sizeof(struct mgt_association_resp_one) + 1] = 0x84;

        sendPacket(dsmb, &sendRaw, sizeof(struct mgt_association_resp_one) + 2); 
        
        dsmb->sentAssocReq++;
    }
    
    dsmb->curState = DSMB_WAITING_FOR_ACK;
    dsmb->processState = processRecievedPacket;
}

void setup3Address80211Header(struct dsmbState *dsmb, struct a3_80211 *header, int type, int subType, __u8 destMAC[6]) {
    memset(header, 0, sizeof(struct a3_80211));
    
    header->mh_version = 0;
    header->mh_type = type;
    header->mh_subtype = subType;
    header->mh_from_ds = 1;
    header->mh_duration_id = 512;

    memcpy(header->mh_mac1, destMAC, 6);
    memcpy(header->mh_mac2, dsmb->MAC, 6);
    memcpy(header->mh_mac3, dsmb->MAC, 6);
    
    header->mh_seq_num = dsmb->seqNum++;
}

void processAcksSendingMBLoader(struct dsmbState *dsmb, __u8 *packet, int length) {
    struct a3_80211 *packetFrame = (struct a3_80211*)packet;
    __u8 rawPacket[65535];

    if (packetFrame->mh_type == FC_TYPE_DATA &&
        (packetFrame->mh_subtype == DATA_DACK || packetFrame->mh_subtype == DATA_NULL_ACK)) {
        printf("\n\n\n\n\nGot ACK packet: %x, %x, length %d\n\n\n\n\n", packetFrame->mh_type, packetFrame->mh_subtype, length);
    
        struct timespec delay;
        delay.tv_sec = 0;
        delay.tv_nsec = 15000;
        
        nanosleep(&delay, NULL);
        
        __u8 tmpMAC[6] = {0x03, 0x09, 0xBF, 0x00, 0x00, 0x03};
        setup3Address80211Header(dsmb, (struct a3_80211*)rawPacket, FC_TYPE_DATA, DATA_DACK, tmpMAC);

        __u8 ackData[4] = { 0x33, 0x00, 0x00, 0x00 };
        memcpy(rawPacket + sizeof(struct a3_80211), ackData, sizeof(ackData));
        
        sendPacket(dsmb, rawPacket, sizeof(struct a3_80211) + sizeof(ackData));
        
        dsmb->recvSendingACK = 0;
        dsmb->packetsSentWithoutACK = 0;
    } else {
        if (packetFrame->mh_type == FC_TYPE_CTL && packetFrame->mh_subtype == CTL_ACK) {
            printf(".");//Got small ACK, length %d, %x, ", length, packetFrame->mh_duration_id);
//            printMAC(packetFrame->mh_mac1);
//            printf("\n");
        } else {
            printf("Got unk packet: %x, %x, length %d\n", packetFrame->mh_type, packetFrame->mh_subtype, length);
            struct timespec delay;
            delay.tv_sec = 0;
            delay.tv_nsec = 15000;
            nanosleep(&delay, NULL);
            sendACK(dsmb);
        }
    }
}

struct multibootBeacon {
    struct a3_80211 macHeader;
    unsigned long long timestamp;
    __u16 interval;
    struct cap_info capInfo;
    __u8 rates[4];
    __u8 dsset[3];
    __u8 tim[7];
    __u8 payload[138];
};

void processRecievedPacket(struct dsmbState *dsmb, __u8 *packet, int length) {
    struct a3_80211 *packetFrame = (struct a3_80211*)packet;
    if (packetFrame->mh_type == FC_TYPE_MGT && packetFrame->mh_subtype == MGT_AUTH) {
        printf("Recieved auth packet, size %d\n", length);
        processAuth(dsmb, packet, length);
    } else if (packetFrame->mh_type == FC_TYPE_MGT && packetFrame->mh_subtype == MGT_AS_RQ) {
        printf("Got association request, size %d.\n", length);
        dsmb->sentAssocReq = 0;
        sendAssocReq(dsmb);
    } else if (dsmb->curState == DSMB_WAITING_FOR_ACK &&
               packetFrame->mh_type == FC_TYPE_CTL && packetFrame->mh_subtype == CTL_ACK) {
        printf("Got ACK for ASSOC. response, assuming DS connected.\n");
        printf("Beginning Multiboot handshaking...\n");
        dsmb->curState = DSMB_SENDING_MB_LOADER;
        dsmb->recvSendingACK = 0;
        dsmb->processState = processAcksSendingMBLoader;
    }
}

int runMessageLoop(struct dsmbState *dsmb) {
    __u8 rawPacket[65535];

    FD_ZERO(&dsmb->socketsToCheck);
    FD_SET(dsmb->sd, &dsmb->socketsToCheck);
    
    dsmb->curState = DSMB_WAITING_FOR_AUTH;
    dsmb->processState = processRecievedPacket;
    
    dsmb->blockTime = 10000;
    dsmb->curTime = 0;
    dsmb->packetsSentWithoutACK = 0;
    dsmb->sentAuth = 0;
    dsmb->sentAssocReq = 0;

    struct timeb baseTime, curTime;
    ftime(&baseTime);  
    
    int lastSent = 0;

    for (;;) {
        if (dsmb->advertisingOn) {
            if (lastSent + 200 < dsmb->curTime) {
                lastSent = dsmb->curTime;

                __u8 beaconPacket[188];
                memset(beaconPacket, 0, 188);
                struct multibootBeacon *beaconHeader = (struct multibootBeacon*)&beaconPacket[0];
                
                beaconHeader->macHeader.mh_type = FC_TYPE_MGT;
                beaconHeader->macHeader.mh_subtype = MGT_BEACON;
               __u8 broadcastMAC[6] = {0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
                memcpy(beaconHeader->macHeader.mh_mac1, broadcastMAC, 6);
                memcpy(beaconHeader->macHeader.mh_mac2, dsmb->MAC, 6);
                memcpy(beaconHeader->macHeader.mh_mac3, dsmb->MAC, 6);
                beaconHeader->macHeader.mh_seq_num = dsmb->seqNum++;
                
                beaconHeader->timestamp = dsmb->curTime * 100;
                beaconHeader->interval = 0xCB;
                beaconHeader->capInfo.ci_ess = 1;                
                beaconHeader->capInfo.ci_shortPreamble = 1;
                
                beaconHeader->rates[0] = 1;
                beaconHeader->rates[1] = 2;
                beaconHeader->rates[2] = 0x82;
                beaconHeader->rates[3] = 0x84;
                
                beaconHeader->dsset[0] = 3;
                beaconHeader->dsset[1] = 1;
                beaconHeader->dsset[2] = 1;
                
                beaconHeader->tim[0] = 5;
                beaconHeader->tim[1] = 5;
                beaconHeader->tim[2] = 1;
                beaconHeader->tim[3] = 2;
                beaconHeader->tim[4] = 0;
                beaconHeader->tim[5] = 0;
                beaconHeader->tim[6] = 0;
                
                beaconHeader->payload[0] = 0xDD;
                beaconHeader->payload[1] = dsmb->beaconPackets[dsmb->curBeacon].length;
                memcpy(&beaconHeader->payload[2], dsmb->beaconPackets[dsmb->curBeacon].packet, dsmb->beaconPackets[dsmb->curBeacon].length);
                                
                if (!sendPacket(dsmb, beaconPacket, 188)) {
                    return 0;
                }
            
                dsmb->curBeacon = (dsmb->curBeacon + 1) % 10;              

                if (dsmb->packetsSentWithoutACK > 30) {
                    struct timespec delay;
                    delay.tv_sec = 0;
                    delay.tv_nsec = 30000;

                    nanosleep(&delay, NULL);
                    
                    printf("Trying to associate again...\n");
                    dsmb->seqNum = 1;
                    sendAssocReq(dsmb);           
                    
                    dsmb->packetsSentWithoutACK = 0;
                }
            }
        }

        if (dsmb->curState == DSMB_SENDING_MB_LOADER) {
            struct timespec delay;
            delay.tv_sec = 0;
            delay.tv_nsec = 10000;

            nanosleep(&delay, NULL);
            
            // Send the keep-alive packets
            // We need to send a Data + CF-Poll
            __u8 tmpMAC[6] = {0x03, 0x09, 0xBF, 0x00, 0x00, 0x00};
            setup3Address80211Header(dsmb, (struct a3_80211*)rawPacket, FC_TYPE_DATA, DATA_POLL, tmpMAC);
            
            __u8 ackData[14] = { 0x06, 0x01, 0x02, 0x00, 0x03, 0x11, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };
            
            ackData[8] = 0x48;
            ackData[10] = 0X02; //(__u8)dsmb->recvSendingACK;
//            dsmb->recvSendingACK = 0x2;
            
            memcpy(rawPacket + sizeof(struct a3_80211), ackData, sizeof(ackData));
            
            sendPacket(dsmb, rawPacket, sizeof(struct a3_80211) + sizeof(ackData));

            dsmb->packetsSentWithoutACK++;
        }
        
        if (pollReadSocket(dsmb, 0, dsmb->blockTime)) {
            int bytesReceived = recv(dsmb->sd, rawPacket, sizeof(rawPacket), 0);
            if (bytesReceived < 0) {
                printf("Error in recieving packets.\n");
                return 0;
            }
            // Skip the prism2 capture header
            dsmb->processState(dsmb, &(rawPacket[0x90]), bytesReceived - 0x90);
        }
    
        ftime(&curTime);
        dsmb->curTime = ((curTime.time - baseTime.time)*1000) + curTime.millitm;
    }    
    
    return 1;
}

void processAuthRecieved(struct dsmbState *dsmb, __u8 *packet, int length);

__u8 beacon1Data, beacon2Data;
__u8 beaconData[6];

void processBeaconRecieved(struct dsmbState *dsmb, __u8 *packet, int length) {
    struct a3_80211 *header = (struct a3_80211 *)packet;
    
    if (header->mh_type == FC_TYPE_MGT && header->mh_subtype == MGT_BEACON) {
        memcpy(beaconData, packet + 52 + 12, 6);
//        beacon1Data = packet[52 + 16];
//        beacon2Data = packet[52 + 17];
        
        struct timespec delay;
        delay.tv_sec = 0;
        delay.tv_nsec = 15000;

        nanosleep(&delay, NULL);
    
        // Send Auth. request
        struct mgt_auth_one authSend;
        memset(&authSend, 0, sizeof(struct mgt_auth_one));
    
        authSend.at1_mac_hdr.mh_type = FC_TYPE_MGT;
        authSend.at1_mac_hdr.mh_subtype = MGT_AUTH;    
        authSend.at1_mac_hdr.mh_seq_num = dsmb->seqNum++;
        authSend.at1_mac_hdr.mh_pwr_man = 1;
        authSend.at1_mac_hdr.mh_duration_id = 0xa2;
        memcpy(dsmb->clientMAC, header->mh_mac2, 6);
        memcpy(authSend.at1_mac_hdr.mh_mac1, header->mh_mac2, 6);
        memcpy(authSend.at1_mac_hdr.mh_mac2, dsmb->MAC, 6);
        memcpy(authSend.at1_mac_hdr.mh_mac3, header->mh_mac2, 6);
    
        authSend.at1_algo = 0;
        authSend.at1_seq = 1;
        authSend.at1_challenge_id.e_element_id = 0;
        authSend.at1_challenge_id.e_length = 0;
        
        sendPacket(dsmb, &authSend, sizeof(authSend));
    } else if (header->mh_type == FC_TYPE_MGT && header->mh_subtype == MGT_AUTH) {
        dsmb->processState = processAuthRecieved;
        processAuthRecieved(dsmb, packet, length);
    }
}

__u8 gotAuth = 0;
void processAuthRecieved(struct dsmbState *dsmb, __u8 *packet, int length) {
    struct a3_80211 *header = (struct a3_80211 *)packet;
    
    if (header->mh_type == FC_TYPE_MGT && header->mh_subtype == MGT_AUTH) {
        struct timespec delay;
        delay.tv_sec = 0;
        delay.tv_nsec = 25000;

        nanosleep(&delay, NULL);

        if (gotAuth == 1) {    
            dsmb->seqNum = 1; gotAuth++;
        } else if (gotAuth == 0) gotAuth++;

        printf("Recieved Auth response\n");
        // Send Assoc. Request
        {
            __u8 sendRaw[sizeof(struct mgt_associate_req) - 6];
            memset(sendRaw, 0, sizeof(sendRaw));
       
            struct mgt_associate_req *assocReq = (struct mgt_association_req*)sendRaw;
        
            assocReq->ar_mac_hdr.mh_type = FC_TYPE_MGT;
            assocReq->ar_mac_hdr.mh_subtype = MGT_AS_RQ;
            assocReq->ar_mac_hdr.mh_pwr_man = 1;
            assocReq->ar_mac_hdr.mh_duration_id = 0xa2;
            assocReq->ar_mac_hdr.mh_seq_num = dsmb->seqNum++;
            memcpy(assocReq->ar_mac_hdr.mh_mac1, header->mh_mac2, 6);
            memcpy(assocReq->ar_mac_hdr.mh_mac2, dsmb->MAC, 6);
            memcpy(assocReq->ar_mac_hdr.mh_mac3, header->mh_mac2, 6);
        
            assocReq->ar_ssid_id.e_element_id = 0;
            assocReq->ar_ssid_id.e_length = 32;
            memcpy(&assocReq->ar_ssid[0], beaconData, 6);
/*            assocReq->ar_ssid[0] = 0x11;
            assocReq->ar_ssid[2] = 0x40;
            // This should be bytes 16&17 from the beacon packets
            assocReq->ar_ssid[4] = beacon1Data;
            assocReq->ar_ssid[5] = beacon2Data;*/

            assocReq->ar_cap.ci_ess = 1;
            assocReq->ar_cap.ci_shortPreamble = 1;
            assocReq->ar_listen_interval = 0x0001;
            assocReq->ar_rates_id.e_element_id = 1;
            assocReq->ar_rates_id.e_length = 2;       
            assocReq->ar_rates[0] = 0x82;
            assocReq->ar_rates[1] = 0x84;

            sendPacket(dsmb, &sendRaw, sizeof(sendRaw)); 
        }
    } else if (header->mh_type == FC_TYPE_DATA && header->mh_subtype == DATA_POLL) {
        printf("Got poll packet.\n");
/*        __u8 sendRaw[sizeof(struct a1_80211)];
        memset(sendRaw, 0, sizeof(struct a1_80211));

        struct a1_80211 *ack = (struct a1_80211*)sendRaw;
        ack->mh_type = FC_TYPE_CTL;
        ack->mh_subtype = CTL_ACK;
        
        memcpy(ack->mh_mac1, dsmb->clientMAC, 6);
    
        sendPacket(dsmb, &sendRaw, sizeof(struct a1_80211));*/
    } else {
//        printf("Got unk packet: %x, %x, length %d\n", header->mh_type, header->mh_subtype, length);
    }
}

int runMessageLoopConnectToDS(struct dsmbState *dsmb) {
    __u8 rawPacket[65535];

    FD_ZERO(&dsmb->socketsToCheck);
    FD_SET(dsmb->sd, &dsmb->socketsToCheck);
    
    dsmb->blockTime = 20000000;
    dsmb->processState = processBeaconRecieved;

    const __u8 MAC[8] = {0x00, 0x09, 0xBF, 0x0A, 0xF1, 0x86};
    memcpy(dsmb->MAC, MAC, 6);
    
    int advertisingCount = 0;    

    for (;;) {
       if (pollReadSocket(dsmb, 0, dsmb->blockTime)) {
            int bytesReceived = recv(dsmb->sd, rawPacket, sizeof(rawPacket), 0);

            if (bytesReceived < 0) {
                printf("Error in recieving packets.\n");
                return 0;
            }

            // Skip the prism2 capture header
            dsmb->processState(dsmb, &(rawPacket[0x90]), bytesReceived - 0x90);
        }
    }
}

int main(int argc, char **argv) {
    signal(SIGINT, shutdownDevices);
    signal(SIGQUIT, shutdownDevices);

    struct dsmbState *dsmb = (struct dsmbState*)malloc(sizeof(struct dsmbState));
    
    initDsMbState(dsmb);

    if (!setupNetwork(dsmb)) {
        return 1;
    }

//    if (!readBeacons(dsmb, "../../simplemariobeacons.orig")) {
//    if (!readBeacons(dsmb, "../../mrdriller.dat")) {
    if (!readBeacons(dsmb, "../../mario.dat")) {
        return 1;
    }       

//    int result = runMessageLoopConnectToDS(dsmb);
    int result = runMessageLoop(dsmb);
    
    free(dsmb);
    
    return result == 0 ? 1 : 0;
}