import nfc
import sys
from nfc.clf import RemoteTarget
from nfc.tag.tt2 import Type2TagCommandError
import logging
from base64 import b64encode, b16encode
from os.path import join
from time import time
from hashlib import sha3_256
import ndef
import coloredlogs
import binascii
import os
import json
from colorama import init, Fore, Back, Style

coloredlogs.install(fmt='%(levelname)s %(message)s')
init()

uids = json.load(open('uids.json'))

def get_hash(payload, count=10000, salt=b"x!"):
	data = sha3_256(payload + (salt or bytes(0)))
	for _ in range(count):
		data = sha3_256(data.digest() + (salt or bytes(0)))
	return data.digest()

def activate_and_unlock(*args, **kwargs):
	tag = nfc.tag.activate(*args, **kwargs)
	tag.has_password = False
	assert tag.signature
	if not tag.ndef.is_writeable:
		tag.has_password = True
		unlock_result = tag.authenticate(password=get_hash(tag.identifier+tag.signature))
		if not unlock_result:
			log.error("This tag is locked and password is incorrect!")
			return
	return tag

if __name__ == '__main__':
	log = logging.getLogger('nfc')
	clf = nfc.ContactlessFrontend()
	
	try:
		assert clf.open('tty:USB0:pn532')
	except TimeoutError:
		log.error('Got TimeoutError on reader connection attempt.')
	except AssertionError:
		log.error('There is no reader connected.')
		exit()

	clf.device.chipset.set_serial_baudrate(115200)
	clf.device.chipset.transport.baudrate = 115200
	log.info(f"Setting device baud rate to {clf.device.chipset.transport.baudrate}.")
	
	print(Style.BRIGHT + "What mode would you like to work in?\ns = single enroll\nm = mass enroll\n>>> ", end='')
	mode = input()[0]
	
	if mode == 'm':
		print(Style.BRIGHT + "Choose the number from where to begin enrolling. To stop enrolling, just use CTRL+C to exit the program.\n>>> ", end='')

		try:
			uid = int(input())
		except KeyboardInterrupt: exit()
		except:
			log.error("The supplied number is not valid.")
			exit()
		print(Fore.GREEN + "✔ Mass enroll mode" + Style.RESET_ALL)
	else:
		print(Fore.GREEN + "✔ Single enroll mode" + Style.RESET_ALL)

	try:
		while 1:

			# Ask for a number if in single enroll mode.
			if mode == 's':
				while 1:
					print(Style.BRIGHT + "Input the attendee number\n>>> ", end='')			
					try:
						uid = int(input())
					except KeyboardInterrupt: exit()
					except:
						log.error("The supplied number is not valid.")
						continue
					else:
						break
			
			
			user_id = uids[str(uid)][0]
			url = "http://go.foxo.me/" + user_id
			
			print(Fore.GREEN + "Attendee info:")
			print(f"NAME: {uids[str(uid)][1]}")
			print(f" INT: {uid}")
			print(f"CODE: {user_id}\n")
			
			print(Fore.YELLOW + "⌛ Waiting for nfc tag", end='')
			
			# Wait for a successful reading
			while 1:
				target = clf.sense(RemoteTarget('106A'), iterations=10)
				if not target:
					print('.', end='')
					sys.stdout.flush()
					continue
				break
				
			print(Style.RESET_ALL)
			
			# Tag was found, write the details
			start = time()
			tag = activate_and_unlock(clf, target)

			if not tag:
				log.error("Tag was gone mid-write.")
				continue

			if tag.product != 'NXP NTAG215':
				log.error("Incorrect tag type detected! Tag must be a NXP 215!")
				while clf.sense(RemoteTarget('106A')): pass
				continue
	
			if not tag.signature:
				log.error("Tag signature was not read correctly. Try again.")
				while clf.sense(RemoteTarget('106A')): pass
				continue
				
			password = get_hash(tag.identifier+tag.signature, 10000)[:6]

			# If tag count is disabled, enable it
			try:
				count = tag.transceive(binascii.unhexlify('3902'), timeout=0.05) # READ_CNT command
			except Type2TagCommandError:
				page = tag.read(0x84)[:4]
				if not page[0] & (1<<4):
					log.info("Enabling tag counter")
					page[0] |= 1 << 4 # Set NFC_CNT_EN = 1
					self.tag.write(0x84, page)

			tag = activate_and_unlock(clf, target)
			count = tag.transceive(binascii.unhexlify('3902'), timeout=0.05) # READ_CNT command

			print(Fore.CYAN+Style.BRIGHT)
			print(f"UUID: {b16encode(tag.identifier).decode()}")
			print(f"SIGN: {b16encode(tag.signature).decode()}")
			print(f"PASS: {b16encode(password[:4]).decode()}")
			print(f"PWCK: {b16encode(password[4:]).decode()}")
			print(" CNT:", int.from_bytes(count, byteorder='little'))
			print(Style.RESET_ALL)
			
			tag.ndef.records = [ndef.UriRecord(url), ]
			try:
				assert tag.ndef.has_changed is False
			except AssertionError:
				log.error("NDEF data was written incorrectly!")
				while clf.sense(RemoteTarget('106A')): pass
				continue
			else:
				log.info("Writing URL")

			if not tag.has_password:
				log.info("Setting password")
				tag.protect(password=password, read_protect=False, protect_from=0)

			print(Back.GREEN + Fore.BLACK + f"\n✔ Badge enrolled with success in {time()-start:.2f}s!" + Style.RESET_ALL)
			uid += 1
			
			while clf.sense(RemoteTarget('106A')): pass
			
			os.system('clear')

	except KeyboardInterrupt:
		pass
		
	clf.close()