// RecoverEncodings returns the full affine output encoding of affine-encoded f at the given position, as well as the
// input affine encodings for all neighboring bytes up to a key byte. Returns (out, []in)
func RecoverEncodings(constr *chow.Construction, round, pos int) (encoding.ByteAffine, []encoding.ByteAffine) {
Ps := make([]encoding.ByteAffine, 4) // Approximate input encodings.
Ds := make([]number.ByteFieldElem, 4) // Array of gamma * MC coefficient
q := byte(0x00) // The constant part of the output encoding.
L := RecoverL(constr, round, pos)
Atilde := FindAtilde(constr, L)
AtildeInv, _ := Atilde.Invert()
for i := 0; i < 4; i++ {
j := pos/4*4 + i
inEnc, _ := RecoverAffineEncoded(
constr, encoding.IdentityByte{}, round-1, common.UnShiftRows(j), common.UnShiftRows(j),
)
_, f := RecoverAffineEncoded(constr, inEnc, round, j, pos)
var c byte
Ds[i], c, Ps[i] = FindPartialEncoding(constr, f, L, AtildeInv)
q ^= c
if i == 0 {
q ^= f.Get(0x00)
}
}
DInv, _ := DecomposeAffineEncoding(encoding.ByteMultiplication(FindDuplicate(Ds).Invert()))
A := Atilde.Compose(DInv)
AInv, _ := A.Invert()
return encoding.ByteAffine{encoding.ByteLinear{A, AInv}, q}, Ps
}
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, common.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][common.UnShiftRows(pos)] != b {
t.Fatalf("Constant part of encoding was not key byte!")
}
}
}
// RecoverKey returns the AES key used to generate the given white-box construction.
func RecoverKey(constr *chow.Construction) []byte {
temp := make([]encoding.Byte, 16)
// Recover all output affine encodings for round 1.
for pos := 0; pos < 16; pos++ {
temp[pos], _ = RecoverEncodings(constr, 1, pos)
}
// Recover all input affine encodings up to a key byte for round 2.
// Compose it with the above's output.
for col := 0; col < 4; col++ {
_, in := RecoverEncodings(constr, 2, 4*col)
for row := 0; row < 4; row++ {
backPos := common.UnShiftRows(4*col + row)
temp[backPos] = encoding.ComposedBytes{temp[backPos], in[row]}
}
}
// Recover round key for round 2.
// The output encoding of round 1 composed with the approximate input encoding of round 2 should be an affine
// transformation with the identity matrix as the linear part and a key byte as the constant part.
roundKey := make([]byte, 16)
for pos := 0; pos < 16; pos++ {
roundKey[pos] = temp[pos].Encode(0)
}
// Recover the master key from the round key and return.
return BackOneRound(BackOneRound(roundKey, 2), 1)
}