from pwn import *
import json
from collections import Counter
from tqdm import tqdm
from termcolor import colored

#Connect to the server and parse until menu entry
def connect_and_menu():
	r = remote("anon.ctfcompetition.com",1337)
	r.recvline_contains("BACKUP")
	return r

#From the menu entry, download the JSON
def getBackUpJSON(r):
	r.send("BACKUP\n")
	return json.loads(r.recvline_endswith("}]"))

#Statistics on the number of cards per account
def getCardSizes(jsonBackup):
	cards = []
	for item in jsonBackup:
		cards.append(len(item["cards"]))
	print(Counter(cards))

#Statistics on the number of accounts with flagged cards, by number of cards in the account
def getFlaggedBySize(jsonBackup):
	sizeDict = {"0":0,"1":0,"2":0,"3":0}
	for acc in jsonBackup:
		numCards = len(acc["cards"])
		log.debug(numCards)
		for card in acc["cards"]:
			log.debug(card)
			if card.get("flagged") == "1":
				log.debug("FLAGGED!")
				sizeDict[str(numCards)] += 1
				break
	print(sizeDict)

#Create an account with a certain number of cards attached (max 3?)
def createAccountWithCards(r,n):
	r.send("NEWACC\n")
	accountName = r.recvline().split(" ")[2]
	log.debug("Created Account Name: " + accountName)
	for i in range(0,n):
		r.send("NEWCARD\n")
		cardName = r.recvline().split(" ")[2]
		log.debug("Card Name: " + cardName)
		r.send("ASSOC " + cardName + " " + accountName + "\n")
		response = r.recvline()
		if "OK" not in response:
			log.error(response)
			exit(-1)
	return accountName

def performASSOC(r,cardName,accountName):
	command = "ASSOC " + cardName + " " + accountName
	r.send(command + "\n")
	response = r.recvline()
	if "KO" in response:
		log.warn(command + " " + response)
		return -1
	else:
		log.debug(response)
		return 0

#Return a coloured string for an account. Red indicates it is flagged
def getAccountAsLine(acc, one, two, three):
	accName = acc["account"]
	cardNum = len(acc["cards"])
	output = accName + " "
	for i in range(0,cardNum):
		bg = "black"
		if acc["cards"][i]["card"] in one:
			bg = "on_yellow"
		elif acc["cards"][i]["card"] in two:
			bg = "on_blue"
		elif acc["cards"][i]["card"] in three:
			bg = "on_magenta"
		if acc["cards"][i].get("flagged") == "1":
			output += colored(acc["cards"][i]["card"] + " ", 'red', bg)
		else:
			output += colored(acc["cards"][i]["card"] + " ", 'green', bg)
	return output

def getAllCards(jsonBackup):
	cards = []
	for acc in jsonBackup:
		for car in acc["cards"]:
			cards.append(car["card"])
	return cards

def getNumFlagged(jsonBackup):
	cards = set()
	flagged = 0
	for acc in jsonBackup:
		for car in acc["cards"]:
			if car.get("flagged") == "1" and car["card"] not in cards:
				flagged += 1
				cards.add(car["card"])
	return str(flagged)


def getNumAccounts(jsonBackup):
	i = 0
	for acc in jsonBackup:
		i += 1
	return str(i)

def getUniqueCards(jsonBackup):
	return str(len(set(getAllCards(jsonBackup))))

def getNumCards(jsonBackup):
	return str(len(getAllCards(jsonBackup)))

def getCardsThatAppear(jsonBackup, n):
	cards = getAllCards(jsonBackup)
	appear = set()
	for car in cards:
		if cards.count(car) == n:
			appear.add(car)
	return appear

def makeGuess(r, guess):
	r.send(guess + "\n")
	return r.recvall()

context.log_level = 'warn' #change to debug to see raw sends/recieves

#getFlaggedBySize(getBackUpJSON(connect_and_menu()))

altered = connect_and_menu()
acc = createAccountWithCards(altered,0)
performASSOC(altered,"ccard"+hex(1),acc)
performASSOC(altered,"ccard"+hex(2),acc)
performASSOC(altered,"ccard"+hex(3),acc)
log.warn("Created ACC with [1,2,3]")
for i in tqdm(range(2,65,2)):
	acc = createAccountWithCards(altered,0)
	performASSOC(altered,"ccard"+hex(i),acc)
	performASSOC(altered,"ccard"+hex(i+1),acc)
	performASSOC(altered,"ccard"+hex(i+2),acc)
	log.warn("Created ACC with ["+str(i)+","+str(i+1)+","+str(i+2)+"]")

#Do some random associations between ccards and new account
# for i in tqdm(range(1,33)):
# 	hexAppend = hex(i)
# 	hexAppend2 = hex(i+32)
# 	acc = createAccountWithCards(altered,0)
# 	performASSOC(altered,"ccard"+hexAppend,acc)
# 	performASSOC(altered,"ccard"+hexAppend2,acc)
# for j in tqdm(range(1,65,3)):
# 	hexAppend = hex(j)
# 	hexAppend2 = hex(j+1)
# 	hexAppend3 = hex(j+2)
# 	acc2 = createAccountWithCards(altered,0)
# 	performASSOC(altered,"ccard"+hexAppend,acc2)
# 	performASSOC(altered,"ccard"+hexAppend2,acc2)
# 	performASSOC(altered,"ccard"+hexAppend3,acc2)
	#acc = createAccountWithCards(altered,0)
	#performASSOC(altered,"ccard"+hexAppend,acc)
	#acc = createAccountWithCards(altered,0)
	#performASSOC(altered,"ccard"+hexAppend,acc)


#for i in tqdm(range(0, 0)):
#	createAccountWithCards(altered,0) #Max 3 cards per account

alteredJSON = getBackUpJSON(altered)
log.warn("Initial Size: " + str(len(alteredJSON)))
prunedJSON = []
for acc in alteredJSON:
	if len(acc["cards"]) == 3:
		prunedJSON.append(acc)

alteredJSON = prunedJSON
log.warn("Pruned Size: " + str(len(alteredJSON)))
#print("Accounts: "+ getNumAccounts(alteredJSON))
#print("Cards (total): " + getNumCards(alteredJSON))
#print("Cards (unique): " + getUniqueCards(alteredJSON))
#print("Flagged Cards (unique): " +getNumFlagged(alteredJSON))

#getCardSizes(alteredJSON)
#getFlaggedBySize(alteredJSON)
#print(Counter(getAllCards(alteredJSON)))

one = getCardsThatAppear(alteredJSON, 1)
two = getCardsThatAppear(alteredJSON, 2)
three = getCardsThatAppear(alteredJSON, 3)

#target = 0
#guess = ""
#for acc in alteredJSON:
#	if len(acc["cards"]) == 3:
#		if acc["cards"][0]["card"] in two and acc["cards"][1]["card"] in two and acc["cards"][2]["card"] in two:
#			print(getAccountAsLine(acc, one, two, three))
#			target += 3
#			if acc["cards"][0].get("flagged") == "1":
#				guess += "1"
#			else:
#				guess += "0"
#			if acc["cards"][1].get("flagged") == "1":
#				guess += "1"
#			else:
#				guess += "0"

#print("Target=" + str(target))
#print("Guess=" + guess)

def getAccountCards(acc):
	cards = set()
	for card in acc["cards"]:
		cards.add(card["card"])
	log.info(str(acc) +" has cards " + str(cards))
	return cards

def getNextAccount(alteredJSON, targetCards, noMatch):
	for acc in alteredJSON:
		if acc["account"] in noMatch:
			continue
		log.info("TESTING " + str(targetCards) + " in " + str(getAccountCards(acc)))
		if targetCards in getAccountCards(acc):
			log.debug("Found " + str(targetCards) + " in " + str(getAccountCards(acc)))
			return acc["account"]
	log.info("DID NOT FIND SUCCESSOR FOR " + str(targetCards))
	return "NONE"

def getNextAccountTWO(alteredJSON, targetCardA, targetCardB, noMatch):
	for acc in alteredJSON:
		if acc["account"] in noMatch:
			continue
		if targetCardA in getAccountCards(acc) and targetCardB in getAccountCards(acc):
			log.debug("Found " + str(targetCardA) + " " + str(targetCardB) + " in " + str(getAccountCards(acc)))
			return acc["account"]
	log.info("DID NOT FIND SUCCESSOR FOR " + str(targetCardA) + " " + str(targetCardB))
	return "NONE"

def getPossibleNextAccounts(alteredJSON, currentAcc, targetSize):
	possible = []
	if targetSize == 1:
		for card in currentAcc["cards"]:
			log.debug("TARGET CARD: " + str(card["card"]))
			noMatch = [currentAcc["account"]]
			candidate = getNextAccount(alteredJSON, card["card"], noMatch)
			while candidate != "NONE":
				possible.append(candidate)
				noMatch.append(candidate)
				candidate = getNextAccount(alteredJSON, card["card"], noMatch)
		log.info("Found "+str(len(possible)))
		return possible
	else:
		for card in currentAcc["cards"]:
			cand = list(getAccountCards(currentAcc))
			cand.remove(card["card"])
			log.info("TARGET CARDS: " + str(cand))
			noMatch = [currentAcc["account"]]
			candidate = getNextAccountTWO(alteredJSON, cand[0], cand[1], noMatch)
			while candidate != "NONE":
				possible.append(candidate)
				noMatch.append(candidate)
				candidate = getNextAccountTWO(alteredJSON, cand[0], cand[1], noMatch)
		log.info("Found "+str(len(possible)))
		return possible

def getAccounts(alteredJSON):
	accounts = []
	for acc in alteredJSON:
		accounts.append(acc["account"])
	return accounts

def getAccountByName(alteredJSON, name):
	for acc in alteredJSON:
		if acc["account"] == name:
			return acc
	log.error("NO SUCH NAME")
	return "NONE"

def getPath(alteredJSON, sequences):
	newSequences = []
	for seq in sequences:
		head = seq[-1]
		nextSteps = []
		if len(seq) > 1:
			log.debug("Looking for one match")
			nextSteps = getPossibleNextAccounts(alteredJSON, getAccountByName(alteredJSON,head), 1)
		else:
			log.info("Looking for two matches")
			nextSteps = getPossibleNextAccounts(alteredJSON, getAccountByName(alteredJSON,head), 2)
		log.debug("Current head: " + str(head))
		if len(nextSteps) > 0:
			log.info("FOUND NEXT STEP")
			for step in nextSteps:
				if step not in seq:
					newSeq = list(seq)
					newSeq.append(step)
					log.info("OLD: " + str(seq) + " NEW: " + str(newSeq))
					newSequences.append(newSeq)
				else:
					log.info("REPEAT! SEQ:"+str(seq)+" "+"CAND:"+str(step))
					log.info("PRUNED a candidate sequence of size " + str(len(seq)))
					if len(seq) == 32:
						return [seq]
		else:
			log.info("PRUNED a candidate sequence of size " + str(len(seq)))
			if len(seq) == 32:
				return [seq]
	return newSequences

def isCardFlagged(alteredJSON, target):
	for acc in alteredJSON:
		for card in acc["cards"]:
			if card["card"] == target:
				if card.get("flagged") == "1":
					return 1
				else:
					return 0
	log.error("CARD NOT FOUND")

accounts = getAccounts(alteredJSON)
current = []
for acc in accounts:
	l = []
	l.append(acc)
	current.append(l)
while len(current) > 0:
	log.warn("Current Accounts:" + str(current))
	current = getPath(alteredJSON, current)
	if len(current) == 1 and len(current[0]) == 32:
		print("Result: "+str(current[0]))
		for a in current[0]:
			print(getAccountAsLine(getAccountByName(alteredJSON,a),one,two,three))
		break
result = []
guess = ""
for a in current[0]:
	result.append((getAccountByName(alteredJSON,a)))
overlapNext = getAccountCards(result[0]).intersection(getAccountCards(result[1]))
noOverlap = getAccountCards(result[0]).difference(overlapNext)
elem = list(noOverlap)[0]
r = isCardFlagged(alteredJSON, elem)
guess += str(r)
for a in list(overlapNext):
	r = isCardFlagged(alteredJSON, a)
	guess += str(r)
noOverlap = getAccountCards(result[1]).difference(overlapNext)
elem = list(noOverlap)[0]
r = isCardFlagged(alteredJSON, elem)
guess += str(r)
for i in range(2, 31):
	overlapLast = getAccountCards(result[i]).intersection(getAccountCards(result[i-1]))
	overlapNext = getAccountCards(result[i]).intersection(getAccountCards(result[i+1]))
	bothOverlap = overlapLast.union(overlapNext)
	noOverlap = getAccountCards(result[i]).difference(bothOverlap)
	elem = list(noOverlap)[0]
	r = isCardFlagged(alteredJSON, elem)
	guess += str(r)
	elem = list(overlapNext)[0]
	r = isCardFlagged(alteredJSON, elem)
	guess += str(r)
overlapLast = getAccountCards(result[30]).intersection(getAccountCards(result[31]))
noOverlap = getAccountCards(result[31]).difference(overlapLast)
for a in list(noOverlap):
	r = isCardFlagged(alteredJSON, a)
	guess += str(r)
print(guess)
print(str(len(guess)))
print(makeGuess(altered, guess))

#print("Cards by Appearance: Once=" + str(len(one))," Twice=" + str(len(two))+ " Thrice=" + str(len(three)))

exit(0)

for acc in sorted(alteredJSON,key=lambda acc: acc["account"]):
	if len(acc["cards"]) == 3:
		print(getAccountAsLine(acc, one, two, three))

print()
print()
print()

#for acc in sorted(alteredJSON,key=lambda acc: acc["account"]):
#	if len(acc["cards"]) == 1:
#		print(getAccountAsLine(acc, one, two, three))


# Parameters
# We have 40 ccards we are intertested in: 	ccard0x1 through ccard0x40 (64 cards)
# When we make a test card the name format is:  ucard0x1 and so on
# When we make a test account the name format is: uaccount0x1
# We can create cards and associate these cards to accounts
# "their encryption is deterministic"
# "generate a anonymized encrypted jsonified maybe-emojified backup."
# WORKING ASSUMPTION: New cards are NEVER flagged.
# WORKING ASSUMPTION: Existing cards are not named caccount0x1 etc
# DEDUCED: An account can be assigned at most 3 cards
# DEDUCED An card can be assigned to at most 3 accounts
# LOGICALLY: Deterministic Encryption so we know we can associate the same card
#  attached to different accounts
# DEDUCED: There are always 300 accounts + amount generated by user

# I strongly suspect the answer will be in the association. Need a clever mapping
#  such that repeated runs are reasonably obvious.

#I can't tell if the server is throttling making solution impossible OR if
#there is a fixed time out after 55 seconds.

#Wait I just realised, id doesn't suffice to identify the set ccards
#and the set of their correspdoning values. We do actually need a one to
#one mapping between the two sets. In particular it seems unlikely such a
#mapping can be produced in the time alotted for the connection. So now I
#feel like a break in the encryption will be needed. Strange for a Misc challenge though...