'''
	Finding lower bounds on the number of active sboxes up to 10 rounds.
	It is used to prove a security bound.
'''

# global variables
bN = 4


def shiftRows(state):
	'''
		Doing shift rows.
	'''

	# row 1 goes one to left
	tmp = state[1]
	for j in range(4):
		state[(1 + j * 4) % 16] = state[(1 + (j + 1) * 4) % 16]
	state[13] = tmp

	# row 2 goes two to left
	for i in range(2):
		tmp = state[2]
		for j in range(4):
			state[(2 + j * 4) % 16] = state[(2 + (j + 1) * 4) % 16]
		state[14] = tmp;

    # row 3 goes three to left (but easier to go one right here)
	tmp = state[15]
	for j in range(4):
		state[(15 - j * 4) % 16] = state[(15 - (j + 1) * 4) % 16];
	state[3] = tmp

	return state


def mixColumn(state, nextD, fd):
	'''
		Doing MixColumn
	'''

	newState = [i for i in range(max(state)+1, max(state)+17)]

	for i in range(4):
		fd.write("p.add_constraint(x[{0}] + x[{1}] + x[{2}] + x[{3}] + x[{4}] + x[{5}] + x[{6}]+ x[{7}] >= {8}*d[{9}])\n".format(state[4*i], state[4*i+1], state[4*i+2], state[4*i+3], newState[4*i], newState[4*i+1], newState[4*i+2], newState[4*i+3], bN, nextD))
		for j in range(4):
			fd.write("p.add_constraint(x[{0}] <= d[{1}])\n".format(state[4*i+j], nextD))
			fd.write("p.add_constraint(x[{0}] <= d[{1}])\n".format(newState[4*i+j], nextD))
		fd.write("p.add_constraint(x[{0}] + x[{1}] + x[{2}] + x[{3}] >= d[{4}])\n".format(state[4*i], state[4*i+1], state[4*i+2], state[4*i+3], nextD))
		fd.write("p.add_constraint(x[{0}] + x[{1}] + x[{2}] + x[{3}] >= d[{4}])\n\n".format(newState[4*i], newState[4*i+1], newState[4*i+2], newState[4*i+3], nextD))
		nextD+=1
	return newState, nextD


def main():
	'''
		Here we make/open a file and write the required equations in it.
		Then we solve the equations in sage.
	'''

	double_shift = [True, False]
	num_rounds = 7

	with open("equations.sage",'w') as fd:

		# the objective is to minimizing (maximization=False)
		fd.write("# set the objective to minimizing ...\n")
		fd.write('p = MixedIntegerLinearProgram(maximization=False, solver = "GLPK")\n')
		fd.write("x = p.new_variable(binary=True)\n")
		fd.write("d = p.new_variable(binary=True)\n")

		for ds in double_shift:
			for rounds in range(1, num_rounds):
				nextD = 0

				# initialize the state with 0 to 15
				state = [i for i in range(16)]

				# set the objective: x[0] + x[1] + ... + x[#rounds*16 - 1]
				fd.write("p.set_objective(")
				for i in range(rounds * 16):
					fd.write("x[{0}]".format(i))
					if(i < ((rounds * 16)-1)):
						fd.write(" + ")
					else: # last element
						fd.write(")\n\n")

				for i in range(rounds):
					state = shiftRows(state) # do shiftrows
					if(double_shift): # what is the double shift???
						state = shiftRows(state)
					state, nextD = mixColumn(state, nextD, fd) # do mixcolumn

				# Constrains: make sure that there is one Aktive S-Box
				fd.write("p.add_constraint(")
				for i in range(16):
					fd.write("x[{0}] ".format(i))
					if(i < 15):
						fd.write(" + ")
					else:
						fd.write(" >= 1)\n\n")

				fd.write("print 'Rounds: {0} -> Active SBoxes: ' + str(p.solve()) + 'Double Shift: {1}'\n\n".format(rounds, ds))


if __name__ == '__main__':
	main()