// GenerateEncryptionKeys creates a white-boxed version of AES with given key for encryption, with any non-determinism
// generated by seed. Opts specifies what type of input and output masks we put on the construction and should be in
// common.{IndependentMasks, SameMasks, MatchingMasks}.
func GenerateEncryptionKeys(key, seed []byte, opts common.KeyGenerationOpts) (out Construction, inputMask, outputMask matrix.Matrix) {
	rs := random.NewSource("Chow Encryption", seed)

	constr := saes.Construction{key}
	roundKeys := constr.StretchedKey()

	// Apply ShiftRows to round keys 0 to 9.
	for k := 0; k < 10; k++ {
		constr.ShiftRows(roundKeys[k])
	}

	skinny := func(pos int) table.Byte {
		return common.TBox{constr, roundKeys[9][pos], roundKeys[10][pos]}
	}

	wide := func(round, pos int) table.Word {
		return table.ComposedToWord{
			common.TBox{Constr: constr, KeyByte1: roundKeys[round][pos]},
			common.TyiTable(pos % 4),
		}
	}

	generateKeys(&rs, opts, &out, &inputMask, &outputMask, common.ShiftRows, skinny, wide)

	return
}

func TestRecoverEncodings(t *testing.T) {
	constr, key := testConstruction()
	fastConstr := fastTestConstruction()

	baseConstr := saes.Construction{key}
	roundKeys := baseConstr.StretchedKey()

	outAff := getOutputAffineEncoding(constr, fastConstr, 1, 0)

	// Manually recover the output encoding.
	Q, Ps := RecoverEncodings(fastConstr, 1, 0)

	if fmt.Sprintf("%x %v", outAff.Linear, outAff.Constant) != fmt.Sprintf("%x %v", Q.Linear, Q.Constant) {
		t.Fatalf("RecoverEncodings recovered the wrong output encoding!")
	}

	// Verify that all Ps composed with their corresponding output encoding equals XOR by a key byte.
	id := matrix.GenerateIdentity(8)
	for pos, P := range Ps {
		outAff := getOutputAffineEncoding(constr, fastConstr, 0, unshiftRows(pos))
		A, b := DecomposeAffineEncoding(encoding.ComposedBytes{outAff, P})

		if fmt.Sprintf("%x", id) != fmt.Sprintf("%x", A) {
			t.Fatalf("Linear part of encoding was not identity!")
		}

		if roundKeys[1][unshiftRows(pos)] != b {
			t.Fatalf("Constant part of encoding was not key byte!")
		}
	}
}

// GenerateDecryptionKeys creates a white-boxed version of AES with given key for decryption, with any non-determinism
// generated by seed. Opts specifies what type of input and output masks we put on the construction and should be in
// common.{IndependentMasks, SameMasks, MatchingMasks}.
func GenerateDecryptionKeys(key, seed []byte, opts common.KeyGenerationOpts) (out Construction, inputMask, outputMask matrix.Matrix) {
	rs := random.NewSource("Chow Decryption", seed)

	constr := saes.Construction{key}
	roundKeys := constr.StretchedKey()

	// Last key needs to be unshifted for decryption to work right.
	constr.UnShiftRows(roundKeys[10])

	skinny := func(pos int) table.Byte {
		return common.InvTBox{constr, 0x00, roundKeys[0][pos]}
	}

	wide := func(round, pos int) table.Word {
		if round == 0 {
			return table.ComposedToWord{
				common.InvTBox{Constr: constr, KeyByte1: roundKeys[10][pos], KeyByte2: roundKeys[9][pos]},
				common.InvTyiTable(pos % 4),
			}
		} else {
			return table.ComposedToWord{
				common.InvTBox{Constr: constr, KeyByte2: roundKeys[9-round][pos]},
				common.InvTyiTable(pos % 4),
			}
		}
	}

	generateKeys(&rs, opts, &out, &inputMask, &outputMask, common.UnShiftRows, skinny, wide)

	return
}

// GenerateDecryptionKeys creates a white-boxed version of the AES key `key` for decryption, with any non-determinism
// generated by `seed`.  The `opts` argument works the same as above.
func GenerateDecryptionKeys(key, seed []byte, opts KeyGenerationOpts) (out Construction, inputMask, outputMask matrix.Matrix) {
	constr := saes.Construction{key}
	roundKeys := constr.StretchedKey()

	// Last key needs to be unshifted for decryption to work right.
	constr.UnShiftRows(roundKeys[10])

	skinny := func(pos int) table.Byte {
		return InvTBox{constr, 0x00, roundKeys[0][pos]}
	}

	wide := func(round, pos int) table.Word {
		if round == 0 {
			return table.ComposedToWord{
				InvTBox{constr, roundKeys[10][pos], roundKeys[9][pos]},
				InvTyiTable(pos % 4),
			}
		} else {
			return table.ComposedToWord{
				InvTBox{constr, 0x00, roundKeys[9-round][pos]},
				InvTyiTable(pos % 4),
			}
		}
	}

	generateKeys(seed, opts, &out, &inputMask, &outputMask, unshiftRows, skinny, wide)

	return
}

func (sr shiftrows) Encode(in [16]byte) (out [16]byte) {
	constr := saes.Construction{}

	copy(out[:], in[:])
	constr.ShiftRows(out[:])

	return
}

func TestBackOneRound(t *testing.T) {
	_, key := testConstruction()
	baseConstr := saes.Construction{key}
	roundKeys := baseConstr.StretchedKey()

	for round := 1; round < 11; round++ {
		a, b := roundKeys[round-1], BackOneRound(roundKeys[round], round)
		if bytes.Compare(a, b) != 0 {
			t.Fatalf("Failed to move back one round on round %v!\nReal: %x\nCand: %x\n", round, a, b)
		}
	}
}

func TestDecrypt(t *testing.T) {
	key := make([]byte, 16)
	rand.Read(key)

	constr1 := saes.Construction{Key: key}
	constr2 := Construction{Key: Expand(key)}

	in, out := make([]byte, 16), make([]byte, 16)
	rand.Read(in)
	constr1.Decrypt(out, in)

	if !constr2.decrypt(Expand(in)).Equals(Expand(out)) {
		t.Fatal("BES Decrypt didn't agree with AES Decrypt!")
	}
}

// backOneRound takes round key i and returns round key i-1.
func backOneRound(roundKey [16]byte, round int) (out [16]byte) {
	constr := saes.Construction{}

	// Recover everything except the first word by XORing consecutive blocks.
	for pos := 4; pos < 16; pos++ {
		out[pos] = roundKey[pos] ^ roundKey[pos-4]
	}

	// Recover the first word by XORing the first block of the roundKey with f(last block of roundKey), where f is a
	// subroutine of AES' key scheduling algorithm.
	for pos := 0; pos < 4; pos++ {
		out[pos] = roundKey[pos] ^ constr.SubByte(out[12+(pos+1)%4])
	}
	out[0] ^= powx[round-1]

	return
}

// GenerateKeys creates a white-boxed version of the AES key `key`, with any non-determinism generated by `seed`.
func GenerateKeys(key, seed []byte) (out Construction, inputMask, outputMask encoding.BlockAffine) {
	rs := random.NewSource("Ful Construction", seed)

	// Generate two completely random affine transformations, to be put on input and output of SPN.
	input, output := generateAffineMasks(&rs)

	// Steal key schedule logic from the standard AES construction.
	contr := saes.Construction{key}
	roundKeys := contr.StretchedKey()

	// Generate an SPN which has the input and output masks, but is otherwise un-obfuscated.
	out[0] = decomposition[0].compose(&blockAffine{
		linear:   matrix.GenerateIdentity(128),
		constant: matrix.Row(roundKeys[0]),
	}).compose(input)
	copy(out[1:5], decomposition[1:5])

	for i := 1; i < 10; i++ {
		out[4*i+0] = decomposition[0].compose(&blockAffine{
			linear:   round,
			constant: matrix.Row(roundKeys[i]).Add(subBytesConst),
		}).compose(out[4*i+0])
		copy(out[4*i+1:4*i+5], decomposition[1:5])
	}

	out[40] = output.compose(&blockAffine{
		linear:   lastRound,
		constant: matrix.Row(roundKeys[10]).Add(subBytesConst),
	}).compose(out[40])

	// Sample self-equivalences of the S-box layer and mix them into adjacent affine layers.
	label := make([]byte, 16)
	copy(label, []byte("Self-Eq"))
	r := rs.Stream(label)

	for i := 0; i < 40; i++ {
		a, bInv := generateSelfEquivalence(r, stateSize[i%4], compressSize[i%4])
		out[i] = a.compose(out[i])
		out[i+1] = out[i+1].compose(bInv)
	}

	return out, input.BlockAffine(), output.BlockAffine()
}

// GenerateEncryptionKeys creates a white-boxed version of the AES key `key` for encryption, with any non-determinism
// generated by `seed`.  The `opts` specifies what type of input and output masks we put on the construction and should
// be either IndependentMasks, SameMasks, or MatchingMasks.
func GenerateEncryptionKeys(key, seed []byte, opts KeyGenerationOpts) (out Construction, inputMask, outputMask matrix.Matrix) {
	constr := saes.Construction{key}
	roundKeys := constr.StretchedKey()

	// Apply ShiftRows to round keys 0 to 9.
	for k := 0; k < 10; k++ {
		constr.ShiftRows(roundKeys[k])
	}

	skinny := func(pos int) table.Byte {
		return TBox{constr, roundKeys[9][pos], roundKeys[10][pos]}
	}

	wide := func(round, pos int) table.Word {
		return table.ComposedToWord{
			TBox{constr, roundKeys[round][pos], 0x00},
			TyiTable(pos % 4),
		}
	}

	generateKeys(seed, opts, &out, &inputMask, &outputMask, shiftRows, skinny, wide)

	return
}

func (sbox sbox) Decode(in byte) byte {
	constr := saes.Construction{}
	return constr.UnSubByte(in)
}