from time import time
from base64 import b16encode
from hashlib import sha3_512
import os.path
from security_const import *
import logging
import binascii
import ndef
from io import BytesIO
from nfc.tag.tt2 import Type2TagCommandError

logging.basicConfig(level='INFO')
log = logging.getLogger('badge')

START_PAGE = {
	"NXP NTAG213": 0x11,
	"NXP NTAG215": 0x6B,
	"NXP NTAG216": 0xCB
}

USER_PAGE_BOUNDS = {
	"NXP NTAG213": (0x04, 0x27),
	"NXP NTAG215": (0x04, 0x81),
	"NXP NTAG216": (0x04, 0xE1),
}

CONFIG_PAGES = {
	"NXP NTAG213": (0x29, 0x2A, 0x2B, 0x2C),
	"NXP NTAG215": (0x83, 0x84, 0x85, 0x86),
	"NXP NTAG216": (0xE3, 0xE4, 0xE5, 0xE6),	
}

class Badge:
	def __init__(self, tag, uuid=None):
		self.tag = tag
		self.uuid = uuid
		
	def read_transaction(self):
		bounds = USER_PAGE_BOUNDS[self.tag.product]
		start_page = START_PAGE[self.tag.product] # We assume transactions are 22 pages wide

		data = b""
		offset = 0

		while 1:
			print("Reading", hex(start_page+offset), end=' ')
		
			chunk = self.tag.read(start_page+offset)
			print(b16encode(chunk))
			if len(data) == 0:
				if bytes(chunk[:4]) != HEADER:
					print('header is', HEADER)
					print('chunk 4', bytes(chunk[:4]))
					print('Does not start with header :(')
			
			data += chunk
			offset += 4
			
			if FOOTER in chunk:
				data = data[:data.find(FOOTER)+4]
				break

		return data

	def write_transaction(self, tx, verify=True):
		bounds = USER_PAGE_BOUNDS[self.tag.product]
		start = START_PAGE[self.tag.product]

		with BytesIO(tx.get_bytes()) as data:
		
			offset = 0
			
			# Write the transaction
			while 1:
				chunk = data.read(4)
				if not chunk: break
				
				print("Writing", hex(start+offset), b16encode(chunk))
				assert bounds[0] <= start+offset <= bounds[1]
				
				self.tag.write(start+offset, chunk)
				
				offset += 1
				
		# Read the transaction back
		if verify:
			print('expected:::', tx.get_bytes())
			print('read:::::', self.read_transaction())
			assert tx.get_bytes() == self.read_transaction()

	def enable_count(self):
	
		page_addr = CONFIG_PAGES[self.tag.product][1]
	
		page = self.tag.read(page_addr)[:4]
		if not page[0] & (1<<4):
			log.info("Enabling tag counter")
			page[0] |= 1 << 4 # Set NFC_CNT_EN = 1
			self.tag.write(page_addr, page)
		else:
			log.info("Tag counter is already enabled!")

		try:
			count = self.tag.transceive(binascii.unhexlify('3902'), timeout=0.05) # READ_CNT command
		except Type2TagCommandError:
			log.error("Tag count is enabled, but didn't get any response anyway.")
		else:
			log.info(f"Tag count is {int.from_bytes(count, byteorder='little')}")

	def protect(self):
		self.tag.protect(self.password())

	def password(self, with_uuid=True):
		x = self.tag.identifier
		for i in range(ITERATIONS):
			x = sha3_512(x + (self.uuid if with_uuid and self.uuid else bytes(0)) + SALT).digest()
		return x[:6]

if __name__ == '__main__':

	from types import SimpleNamespace
	
	tag = SimpleNamespace()
	tag.identifier = bytes(8)

	pt = Protect(tag, uuid=b"bye")
	start = time()
	print('with uuid', b16encode(pt.password()))
	print('without uuid', b16encode(pt.password(with_uuid=False)))
	print("took", (time()-start)*1000, 'ms')