//-----------------------------------------------------------------------------
// This code is licensed to you under the terms of the GNU GPL, version 2 or,
// at your option, any later version. See the LICENSE.txt file for the text of
// the license.
//-----------------------------------------------------------------------------
// main code for standalone HF Sniff (and ULC/NTAG/ULEV1 pwd storing)
//-----------------------------------------------------------------------------

/*
This can actually be used in two separate ways.
It can either be used to just HF 14a sniff on the go and/or grab the 
authentication attempts for ULC/NTAG/ULEV1 into the flash mem (RDV4).

The retrieved sniffing session can be acquired by connecting the device
to a client that supports the reconnect capability and issue 'hf 14a list'.

In order to view the grabbed authentication attempts in the flash mem,
you can simply run 'script run read_pwd_mem' or just 'mem read l 256' 
from the client to view the stored quadlets.
*/

#include "hf_bog.h"

#define DELAY_READER_AIR2ARM_AS_SNIFFER (2 + 3 + 8) 
#define DELAY_TAG_AIR2ARM_AS_SNIFFER (3 + 14 + 8)

// Maximum number of auth attempts per standalone session
#define MAX_PWDS_PER_SESSION 64

uint8_t FindOffsetInFlash() {
	uint8_t mem[4] = { 0x00, 0x00, 0x00, 0x00 };
	uint8_t eom[4] = { 0xFF, 0xFF, 0xFF, 0xFF };
	uint8_t memcnt = 0;
	
	while (memcnt < 0xFF)
	{
		Flash_ReadData(memcnt, mem, 4);
		if (memcmp(mem, eom, 4) == 0) {
			return memcnt;
		}
		memcnt += 4;
	}
	
	return 0; // wrap-around
}

void EraseMemory() {
    if (!FlashInit()){
        return;
    }
    
    Flash_CheckBusy(BUSY_TIMEOUT);
    Flash_WriteEnable();
    Flash_Erase4k(0,0);

    if (MF_DBGLEVEL > 1) Dbprintf("[!] Erased flash!");
    FlashStop();
	SpinDelay(100);
}

// This is actually copied from SniffIso14443a
void RAMFUNC SniffAndStore(uint8_t param) {
	
	iso14443a_setup(FPGA_HF_ISO14443A_SNIFFER);
	
	// Allocate memory from BigBuf for some buffers
	// free all previous allocations first
	BigBuf_free(); BigBuf_Clear_ext(false);
	clear_trace();
	set_tracing(true);
	
	// Array to store the authpwds
	uint8_t *capturedPwds = BigBuf_malloc(4 * MAX_PWDS_PER_SESSION);
	
	// The command (reader -> tag) that we're receiving.
	uint8_t *receivedCmd = BigBuf_malloc(MAX_FRAME_SIZE);
	uint8_t *receivedCmdPar = BigBuf_malloc(MAX_PARITY_SIZE);
	
	// The response (tag -> reader) that we're receiving.
	uint8_t *receivedResp = BigBuf_malloc(MAX_FRAME_SIZE);
	uint8_t *receivedRespPar = BigBuf_malloc(MAX_PARITY_SIZE);
	
	// The DMA buffer, used to stream samples from the FPGA
	uint8_t *dmaBuf = BigBuf_malloc(DMA_BUFFER_SIZE);
	uint8_t *data = dmaBuf;

	uint8_t previous_data = 0;
	int dataLen = 0;
	bool TagIsActive = false;
	bool ReaderIsActive = false;
	
	// Set up the demodulator for tag -> reader responses.
	DemodInit(receivedResp, receivedRespPar);
	
	// Set up the demodulator for the reader -> tag commands
	UartInit(receivedCmd, receivedCmdPar);
	
	// Setup and start DMA.
	if ( !FpgaSetupSscDma((uint8_t*) dmaBuf, DMA_BUFFER_SIZE) ){
		if (MF_DBGLEVEL > 1) Dbprintf("FpgaSetupSscDma failed. Exiting"); 
		return;
	}
	
	tUart* uart = GetUart();
	tDemod* demod = GetDemod();
	
	// We won't start recording the frames that we acquire until we trigger;
	// a good trigger condition to get started is probably when we see a
	// response from the tag.
	// triggered == false -- to wait first for card
	bool triggered = !(param & 0x03); 
	
	uint32_t rsamples = 0;
	
	// Current captured passwords counter
	uint8_t auth_attempts = 0;

	SpinDelay(50);
	
	// loop and listen
	while (!BUTTON_PRESS()) {
        WDT_HIT();
        LED_A_ON();
		
		int register readBufDataP = data - dmaBuf;
		int register dmaBufDataP = DMA_BUFFER_SIZE - AT91C_BASE_PDC_SSC->PDC_RCR;
		if (readBufDataP <= dmaBufDataP)
			dataLen = dmaBufDataP - readBufDataP;
		else
			dataLen = DMA_BUFFER_SIZE - readBufDataP + dmaBufDataP;
		
		// test for length of buffer
		if (dataLen > DMA_BUFFER_SIZE) { // TODO: Check if this works properly
			Dbprintf("[!] blew circular buffer! | datalen %u", dataLen);
			break;
		}
		if (dataLen < 1) continue;

		// primary buffer was stopped( <-- we lost data!
		if (!AT91C_BASE_PDC_SSC->PDC_RCR) {
			AT91C_BASE_PDC_SSC->PDC_RPR = (uint32_t) dmaBuf;
			AT91C_BASE_PDC_SSC->PDC_RCR = DMA_BUFFER_SIZE;
			//Dbprintf("[-] RxEmpty ERROR | data length %d", dataLen); // temporary
		}
		// secondary buffer sets as primary, secondary buffer was stopped
		if (!AT91C_BASE_PDC_SSC->PDC_RNCR) {
			AT91C_BASE_PDC_SSC->PDC_RNPR = (uint32_t) dmaBuf;
			AT91C_BASE_PDC_SSC->PDC_RNCR = DMA_BUFFER_SIZE;
		}

		LED_A_OFF();
		
		// Need two samples to feed Miller and Manchester-Decoder
		if (rsamples & 0x01) {

			if (!TagIsActive) {		// no need to try decoding reader data if the tag is sending
				uint8_t readerdata = (previous_data & 0xF0) | (*data >> 4);
				if (MillerDecoding(readerdata, (rsamples-1)*4)) {
					LED_C_ON();

					// check - if there is a short 7bit request from reader
					if ((!triggered) && (param & 0x02) && (uart->len == 1) && (uart->bitCount == 7)) triggered = true;

					if (triggered) {
						if ((receivedCmd) && ((receivedCmd[0] == MIFARE_ULEV1_AUTH) || (receivedCmd[0] == MIFARE_ULC_AUTH_1))) {
							if (MF_DBGLEVEL > 1) Dbprintf("PWD-AUTH KEY: 0x%02x%02x%02x%02x", receivedCmd[1], receivedCmd[2], receivedCmd[3], receivedCmd[4]);
							
							// temporarily save the captured pwd in our array
							memcpy(&capturedPwds[4 * auth_attempts], receivedCmd+1, 4);
							auth_attempts++;
						}
						
						if (!LogTrace(receivedCmd, 
										uart->len, 
										uart->startTime*16 - DELAY_READER_AIR2ARM_AS_SNIFFER,
										uart->endTime*16 - DELAY_READER_AIR2ARM_AS_SNIFFER,
										uart->parity, 
										true)) break;
					}
					/* ready to receive another command. */
					UartReset();
					/* reset the demod code, which might have been */
					/* false-triggered by the commands from the reader. */
					DemodReset();
					LED_B_OFF();
				}
				ReaderIsActive = (uart->state != STATE_UNSYNCD);
			}

			// no need to try decoding tag data if the reader is sending - and we cannot afford the time
			if (!ReaderIsActive) {		
				uint8_t tagdata = (previous_data << 4) | (*data & 0x0F);
				if (ManchesterDecoding(tagdata, 0, (rsamples-1)*4)) {
					LED_B_ON();

					if (!LogTrace(receivedResp, 
									demod->len, 
									demod->startTime*16 - DELAY_TAG_AIR2ARM_AS_SNIFFER, 
									demod->endTime*16 - DELAY_TAG_AIR2ARM_AS_SNIFFER,
									demod->parity,
									false)) break;

					if ((!triggered) && (param & 0x01)) triggered = true;

					// ready to receive another response.
					DemodReset();
					// reset the Miller decoder including its (now outdated) input buffer
					UartReset();
					//UartInit(receivedCmd, receivedCmdPar);
					LED_C_OFF();
				} 
				TagIsActive = (demod->state != DEMOD_UNSYNCD);
			}
		}

		previous_data = *data;
		rsamples++;
		data++;
		if (data == dmaBuf + DMA_BUFFER_SIZE) {
			data = dmaBuf;
		}
	} // end main loop

	FpgaDisableSscDma();
	set_tracing(false);
	
	Dbprintf("Stopped sniffing");
	
	SpinDelay(200);
	
	// Write stuff to flash
	if (auth_attempts > 0) {
		if (MF_DBGLEVEL > 1) Dbprintf("[!] Authentication attempts = %u", auth_attempts);
		
		// Setting the SPI Baudrate to 48MHz to avoid the bit-flip issue (https://github.com/RfidResearchGroup/proxmark3/issues/34)
		FlashmemSetSpiBaudrate(48000000);
		
		// Find the offset in flash mem to continue writing the auth attempts
		uint8_t memoffset = FindOffsetInFlash();
		if (MF_DBGLEVEL > 1) Dbprintf("[!] Memory offset = %u", memoffset);
		
		if ((memoffset + 4 * auth_attempts) > 0xFF)
		{
			// We opt to keep the new data only
			memoffset = 0;
			if (MF_DBGLEVEL > 1) Dbprintf("[!] Size of total data > 256 bytes. Discarding the old data.");
		}
		
		// Get previous data from flash mem
		uint8_t *previousdata = BigBuf_malloc(memoffset);
		if (memoffset > 0)
		{
			uint16_t readlen = Flash_ReadData(0, previousdata, memoffset);
			if (MF_DBGLEVEL > 1) Dbprintf("[!] Read %u bytes from flash mem", readlen);
		}
				
		// create new bigbuf to hold all data
		size_t total_size = memoffset + 4 * auth_attempts;
		uint8_t *total_data = BigBuf_malloc(total_size);
		
		// Add the previousdata array into total_data array
		memcpy(total_data, previousdata, memoffset);
		
		// Copy bytes of capturedPwds immediately following bytes of previousdata
		memcpy(total_data + memoffset, capturedPwds, 4 * auth_attempts);
		
		// Erase first page of flash mem
		EraseMemory();
				
		// Write total data to flash mem
		uint16_t writelen = Flash_WriteData(0, total_data, memoffset + 4 * auth_attempts);
		if (MF_DBGLEVEL > 1) Dbprintf("[!] Wrote %u bytes into flash mem", writelen);
		
                // If pwd saved successfully, blink led A three times
		if (writelen > 0) {
			SpinErr(0, 200, 5); // blink led A
		}
		
		SpinDelay(100);
		
		// Reset the SPI Baudrate to the default value (24MHz)
		FlashmemSetSpiBaudrate(24000000);
	}
}

void RunMod() {
	
	StandAloneMode();
	
	Dbprintf(">>  Bogiton 14a Sniff UL/UL-EV1/NTAG a.k.a BogitoRun Started  <<");
	Dbprintf("Starting to sniff");
	
	// param:
	// bit 0 - trigger from first card answer
	// bit 1 - trigger from first reader 7-bit request	
	SniffAndStore(0);	
	LEDsoff();	
	SpinDelay(300);
	Dbprintf("- [ End ] -> You can take shell back ...");
	Dbprintf("- [  !  ] -> use 'script run read_pwd_mem' to print passwords");	
}
