/*
 *  spreadspace avr utils
 *
 *
 *  Copyright (C) 2013 Christian Pointner <equinox@spreadspace.org>
 *                     Othmar Gsenger <otti@wirdorange.org>
 *
 *  This file is part of spreadspace avr utils.
 *
 *  spreadspace avr utils 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 3 of the License, or
 *  any later version.
 *
 *  spreadspace avr utils 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 spreadspace avr utils. If not, see <http://www.gnu.org/licenses/>.
 */

#include <avr/io.h>
#include <avr/wdt.h>
#include <avr/interrupt.h>
#include <avr/power.h>

#include <stdio.h>
#include <string.h>

#include "util.h"
#include "led.h"

#include "LUFA/Drivers/Peripheral/TWI.h"
#include "LUFA/Drivers/Peripheral/Serial.h"
#include "LUFA/Drivers/Misc/RingBuffer.h"

#include <avr/eeprom.h>
#define EEPROM_SIZE 1024
typedef uint8_t keyslot_t[8];
keyslot_t EEMEM keystore[EEPROM_SIZE/sizeof(keyslot_t)];


const char* SL018_cmd_tostring(const uint8_t cmd) 
{
  switch(cmd) {
  case 0x01: return "Select Mifare card";
  case 0x02: return "Login to a sector";
  case 0x03: return "Read a data block";
  case 0x04: return "Write a data block";
  case 0x05: return "Read a value block";
  case 0x06: return "Initialize a value block";
  case 0x07: return "Write master key";
  case 0x08: return "Increment value";
  case 0x09: return "Decrement value";
  case 0x0A: return "Copy value";
  case 0x10: return "Read a data page";
  case 0x11: return "Write a data page";
  case 0x40: return "Control the red led";
  case 0xF0: return "Get firmware version";
  case 0xFF: return "Reset";
  default: return "unknown";
  }
}

const char* SL018_status_tostring(const uint8_t status) 
{
  switch(status) {
  case 0x0: return "Operation succeed";
  case 0x1: return "No tag";
  case 0x2: return "Login succeed";
  case 0x3: return "Login fail";
  case 0x4: return "Read fail";
  case 0x5: return "Write fail";
  case 0x6: return "Unable to read after write";
  case 0xA: return "Collision occur";
  case 0xC: return "Load key fail";
  case 0xD: return "Not authenticate";
  case 0xE: return "Not a value block";
  default: return "unknown";
  }
}

const char* SL018_tagtype_tostring(const uint8_t type) 
{
  switch(type) {
  case 0x1: return "Mifare 1k, 4 byte UID";
  case 0x2: return "Mifare 1k, 7 byte UID";
  case 0x3: return "Mifare Ultralight or NATG203, 7 byte UID";
  case 0x4: return "Mifare 4k, 4 byte UID";
  case 0x5: return "Mifare 4k, 7 byte UID";
  case 0x6: return "Mifare DesFire, 7 byte UID";
  default: return "unknown";
  }
}

uint8_t SL018_tagtype_to_uidlen(const uint8_t type) 
{
  switch(type) {
  case 0x1:
  case 0x4: return 4;
  case 0x2:
  case 0x3:
  case 0x5:
  case 0x6: return 7;
  default: return 0;
  }
}

const uint8_t SL018_CMD_ComSelectCard[]         = {1,0x01};
const uint8_t SL018_CMD_ComRedLedOn[]           = {2,0x40,1};
const uint8_t SL018_CMD_ComRedLedOff[]          = {2,0x40,0};
const uint8_t SL018_CMD_ComGetFirmwareVersion[] = {1,0xF0};
const uint8_t SL018_CMD_ComReset[]              = {1,0xFF};

uint8_t twi_recv_buf[256];
typedef struct __attribute__((__packed__))
{
  uint8_t len;
  uint8_t command;
  uint8_t status;
  uint8_t data[sizeof(twi_recv_buf)-3];
} sl018_message_t;

sl018_message_t* twi_recv_msg = (sl018_message_t *)&twi_recv_buf;

#define SL018_TWI_ADDR  0xA0
#define SL018_TAG_STA_PIN PINB
#define SL018_TAG_STA_BIT 7
#define CARD_PRESENT (!((SL018_TAG_STA_PIN >> SL018_TAG_STA_BIT) & 1))
uint8_t card_status = 0;



FILE usb_stream;
FILE serial_stream;
FILE * stdio = &serial_stream;

/*
             LUFA Library
     Copyright (C) Dean Camera, 2012.

  dean [at] fourwalledcubicle [dot] com
           www.lufa-lib.org
*/


#include <LUFA/Drivers/USB/USB.h>
#include "lufa-descriptor-usbserial.h"

USB_ClassInfo_CDC_Device_t VirtualSerial_CDC_Interface =
  {
    .Config =
      {
        .ControlInterfaceNumber         = 0,

        .DataINEndpointNumber           = CDC_TX_EPNUM,
        .DataINEndpointSize             = CDC_TXRX_EPSIZE,
        .DataINEndpointDoubleBank       = false,

        .DataOUTEndpointNumber          = CDC_RX_EPNUM,
        .DataOUTEndpointSize            = CDC_TXRX_EPSIZE,
        .DataOUTEndpointDoubleBank      = false,

        .NotificationEndpointNumber     = CDC_NOTIFICATION_EPNUM,
        .NotificationEndpointSize       = CDC_NOTIFICATION_EPSIZE,
        .NotificationEndpointDoubleBank = false,
      },
  };

void EVENT_USB_Device_ConfigurationChanged(void)
{
  CDC_Device_ConfigureEndpoints(&VirtualSerial_CDC_Interface);
  CDC_Device_CreateStream(&VirtualSerial_CDC_Interface,&usb_stream);
  stdio = &usb_stream;
}

void EVENT_USB_Device_ControlRequest(void)
{
  CDC_Device_ProcessControlRequest(&VirtualSerial_CDC_Interface);
}

void EVENT_USB_Device_Disconnect(void)
{
  stdio = &serial_stream;
}
/* end LUFA CDC-ACM specific definitions*/



int16_t stdio_bytes_received(void)
{
  if(stdio == &usb_stream)
    return CDC_Device_BytesReceived(&VirtualSerial_CDC_Interface);
  else
    return (int16_t)Serial_IsCharReceived();
}

uint8_t sl018_cmd_raw(const uint8_t* twi_send_buf, bool wait_for_answer)
{
  uint8_t pos = 0;

  if (TWI_StartTransmission(SL018_TWI_ADDR | TWI_ADDRESS_WRITE,10) == TWI_ERROR_NoError) {
    for(pos=0; pos<=twi_send_buf[0]; pos++) {
      if( ! TWI_SendByte(twi_send_buf[pos])) {
        TWI_StopTransmission();
        return 1;
      }
    }
    TWI_StopTransmission();
  } else 
    return 1;

  if(!wait_for_answer) return 0;

  memset(twi_recv_buf, 0, sizeof(twi_recv_buf));
  _delay_ms(50);

  if (TWI_StartTransmission(SL018_TWI_ADDR | TWI_ADDRESS_READ,10) == TWI_ERROR_NoError) {
    TWI_ReceiveByte(twi_recv_buf, 0);
    for(pos=1; pos<=twi_recv_buf[0]; pos++) {
      if (! TWI_ReceiveByte(&twi_recv_buf[pos], (pos == twi_recv_buf[0]) ? 1:0 ) ) {
        TWI_StopTransmission();
        return 1;
      }
    }
    TWI_StopTransmission();
  } else
    return 1;

  return 0;
}

uint8_t sl018_reset(void)
{
  if(sl018_cmd_raw(SL018_CMD_ComReset, 0)) {
    fprintf(stdio, "I2C error\n\r");
    return 1;
  }
  return 0;
}

uint8_t sl018_cmd(const uint8_t* twi_send_buf)
{
  if(sl018_cmd_raw(twi_send_buf, 1)) {
    fprintf(stdio, "I2C error\n\r");
    return 1;
  } else {
    if(twi_recv_msg->len < 2) {
      fprintf(stdio,"SL018 Cmd,Error: short message received\n\r");
      return 1;
    }
    if(twi_recv_msg->status) {
      fprintf(stdio,"SL018 Cmd,Error: '%s','%s'\n\r",SL018_cmd_tostring(twi_recv_msg->command),SL018_status_tostring(twi_recv_msg->status));
      return 1;
    }
    sl018_message_t * twi_send_msg = (sl018_message_t *)twi_send_buf;
    if(twi_send_msg->command != twi_recv_msg->command) {
      fprintf(stdio,"SL018 Cmd,Error: mismatch of sent and received command code: %02X,%02X\n\r",twi_send_msg->command,twi_recv_msg->command);
    }
  }
  return 0;
}



void flash_keystore_from_stdio(void)
{
  keyslot_t ks;
  uint8_t byte_pos=0;
  fprintf(stdio,"Flashing keys:\n\r");
  fflush(stdio);
  for(uint8_t ks_pos=0;ks_pos<EEPROM_SIZE/sizeof(ks);) {
    CDC_Device_USBTask(&VirtualSerial_CDC_Interface);
    USB_USBTask();

    int16_t bytes_received = stdio_bytes_received();
    while(bytes_received > 0) {
      ks[byte_pos++]=fgetc(stdio);
      bytes_received--;
      if (byte_pos == sizeof(ks)) {
        byte_pos=0;
        eeprom_update_block(&ks,&keystore[ks_pos],sizeof(ks));
        ks_pos++;
        fputc('.', stdio);
        fflush(stdio);
        led_toggle();
      }
    }
  }
  fprintf(stdio,"\n");
  fputc(0, stdio);
  led_off();
}

void dump_keystore_to_stdio(void)
{
  keyslot_t ks;
  for(uint8_t ks_pos=0;ks_pos<EEPROM_SIZE/sizeof(ks);ks_pos++) {
    eeprom_read_block(&ks,&keystore[ks_pos],sizeof(ks));
    for (uint8_t i=0; i< sizeof(ks); i++)
      fprintf(stdio,"%02X",ks[i]);
    fprintf(stdio,"\n\r");
  }
}

void handle_stdio(uint8_t cmd)
{  
  switch(cmd) {
    case 'r': 
             reset2bootloader(); 
             break;
    case '0': 
             if(!sl018_reset())
               fprintf(stdio, "ok\n\r");
             break;
    case 'f': //get cardreader firmware version
             if(!sl018_cmd(SL018_CMD_ComGetFirmwareVersion)) {
               twi_recv_msg->data[sizeof(twi_recv_msg->data) - 1] = 0;
               fprintf(stdio, "%s\n\r",twi_recv_msg->data);
             }
             break;
    case '4': //turn cardreader led off
             if(!sl018_cmd(SL018_CMD_ComRedLedOff))
               fprintf(stdio, "ok\n\r");
             break;
    case '5': //turn cardreader led on
             if(!sl018_cmd(SL018_CMD_ComRedLedOn))
               fprintf(stdio, "ok\n\r");
             break;
    case 'e': //flash eeprom
             flash_keystore_from_stdio();
             break;
    case 'd': //dump eeprom - this breaks security!
             dump_keystore_to_stdio();
             break;
    default: fprintf(stdio, "error, unknown command %02X '%c'\n\r",cmd, cmd); return;
  }
}



/* this generates a Fletcher8 checksum  */
/* code from: http://stackoverflow.com/questions/13491700/8-bit-fletcher-checksum-of-16-byte-data */
uint8_t generate_csum(uint8_t* data)
{
  uint16_t sum1 = 0xf, sum2 = 0xf, len = sizeof(keyslot_t) - 1;
  do { sum2 += ( sum1 += *data++ ); } while (--len);
  return sum2<<4 | sum1;
}

uint8_t compare_keyslots(const keyslot_t a, const keyslot_t b)
{
  uint8_t tmp=0;
      // constant time compare
  for(uint8_t i=0; i<sizeof(keyslot_t); ++i)
    tmp |= a[i] ^ b[i];
  return tmp;
}

bool check_card(const uint8_t * uid, uint8_t uid_len)
{
  keyslot_t card, ks;
  memset(card, 0, sizeof(card));
  for (uint8_t pos=0; pos<uid_len; pos++)
    card[pos]=uid[uid_len-pos-1];
  card[sizeof(keyslot_t)-1]=generate_csum(card);
  bool valid=0;
  for(uint8_t ks_pos=0;ks_pos<(EEPROM_SIZE/sizeof(ks));ks_pos++) {
    eeprom_read_block(&ks,&keystore[ks_pos],sizeof(ks));
    if(!compare_keyslots(card, ks)) {
      valid=1;
      // break;  // this would break security (not constant time)
    }
  }
  return valid;
}

void handle_card(void)
{
  fprintf(stdio, "CARD IN: ");
  if(!sl018_cmd(SL018_CMD_ComSelectCard))
  {
    uint8_t uid_len = twi_recv_msg->len - sizeof(twi_recv_msg->command) - sizeof(twi_recv_msg->status) - 1;
    if(uid_len == 255 || uid_len > sizeof(keyslot_t) - 1) {
      fprintf(stdio," received UID length (%d) is to big for keystore \n\r", uid_len);
      return;
    }
    uint8_t type = twi_recv_msg->data[uid_len];
    uint8_t expected_uid_len = SL018_tagtype_to_uidlen(type);
    if(expected_uid_len != uid_len) {
      fprintf(stdio," Invalid uid length (%d) for tag type: %s\n\r", uid_len, SL018_tagtype_tostring(type));
      return;
    }    

    for (uint8_t pos=0; pos<uid_len; pos++)
      fprintf(stdio,"%02X",twi_recv_msg->data[uid_len-pos-1]);
    fprintf(stdio, ", %s\n\r", SL018_tagtype_tostring(type));

    if (0 < type && type < 7) {
      if(check_card(twi_recv_msg->data,uid_len)) {
        sl018_cmd(SL018_CMD_ComRedLedOn);
        fprintf(stdio,"Card allowed - opening/closing door\n\r"); // TODO: open/close door!
        _delay_ms(255);
        sl018_cmd(SL018_CMD_ComRedLedOff);
      } else {
        fprintf(stdio,"Card not found in Database\n\r");
      }
    } else {
      fprintf(stdio,"Ignoring unknown card type %02x\n\r",type);
    }
  }  
}



int main(void)
{
  MCUSR &= ~(1 << WDRF);
  wdt_disable();

  cpu_init();
  led_init();
  USB_Init();
  TWI_Init(TWI_BIT_PRESCALE_1, TWI_BITLENGTH_FROM_FREQ(1, 200000));
  sei();
  Serial_Init(115200,false);
  Serial_CreateStream(&serial_stream);
  
  sl018_reset();
  for(;;) {
    CDC_Device_USBTask(&VirtualSerial_CDC_Interface);
    USB_USBTask();

    int16_t bytes_received = stdio_bytes_received();
    if(bytes_received > 0)
      handle_stdio(fgetc(stdio));

    if(CARD_PRESENT != card_status) {
      card_status = CARD_PRESENT;
      if(card_status)
        handle_card();
    }
  }
}