Possible bug in the shift_rows_output leak model

Using the shift_rows_output leak model, I noticed that Pearson correlation is exactly the same for all subkey candidates at all master key positions except 0, 4, 8, 12.

Looking an implementation of the ShiftColumns_output class details I found the reason of such behaviour.

ShiftColumns_output.leakage() doesn’t properly manage

The math of working positions:

output[4] = value from input[4] = sbox(pt[4] ^ key[4]) - correct

The math of non-working positions:

Mapping for output position 5: mapping[5] = 9

This means: output[5] = value from input[9] = sbox(pt[9] ^ key[9])

Visualization for better understanding:

State before Shift Raw

Row 0: [00] [04] [08] [12]
Row 1: [01] [05] [09] [13]
Row 2: [02] [06] [10] [14]
Row 3: [03] [07] [11] [15]

State after Shift Raw

Row 0: [00] [04] [08] [12]  <- No shift
Row 1: [05] [09] [13] [01]  <- Shift left 1
Row 2: [10] [14] [02] [06]  <- Shift left 2
Row 3: [15] [03] [07] [11]  <- Shift left 3

So, when we try to find HW of the key index 5, we get static value as a result of 09 ^ [some fixed subkey] (position 09 moved to position 05) which produces the same HW for each power trace. This will produce the same leakage vector for each subkey hypothesis (as a plaintexts and a subkey hypothesis are the same).

Since ALL hypotheses have the same leakage vector, they all get the same correlation value

I changed

#class ShiftColumns_output(AESLeakageHelper):
#    name = 'HW: AES ShiftColumns Output'
#    def leakage(self, pt, ct, key, bnum):
#        state = [pt[i] ^ key[i] for i in range(0, 16)]
#        state = self.subbytes(state)
#        state = self.shiftrows(state)
#        return state[bnum]

on

class ShiftColumns_output(AESLeakageHelper):
    name = 'HW: AES ShiftColumns Output'

    KEY_TO_OUTPUT = [0, 13, 10, 7, 4, 1, 14, 11, 8, 5, 2, 15, 12, 9, 6, 3]

    def leakage(self, pt, ct, key, bnum):
        """
        key: Full 16-byte key (with guess at position bnum)
        bnum: KEY BYTE INDEX (0-15) being guessed

        Returns: Byte value at the OUTPUT position affected by key[bnum]
        """
        output_pos = self.KEY_TO_OUTPUT[bnum]

        state = [pt[i] ^ key[i] for i in range(0, 16)]
        state = self.subbytes(state)
        state = self.shiftrows(state)

        return state[output_pos]

and got the correct correlation values.

How it was before the fix

and how it looks with the fixed code

This also affects round_1_2_state* family leakage models.

Here are my fixes:

class Round1Round2StateDiff_Text(AESLeakageHelper):
    name = 'HD: AES Round1/Round2 State diff for text'
    KEY_TO_OUTPUT = [0, 13, 10, 7, 4, 1, 14, 11, 8, 5, 2, 15, 12, 9, 6, 3]

    def leakage(self, pt, ct, key, bnum):
        state1 = pt[:]
        state = [pt[i] ^ key[i] for i in range(0, 16)]
        state = self.subbytes(state)
        state = self.shiftrows(state)
        state = self.mixcolumns(state)
        output_pos = self.KEY_TO_OUTPUT[bnum]  # FIX
        return state[output_pos] ^ state1[output_pos]

class Round1Round2StateDiff_KeyMix(AESLeakageHelper):
    name = 'HD: AES Round1/Round2 State diff for key addition'
    KEY_TO_OUTPUT = [0, 13, 10, 7, 4, 1, 14, 11, 8, 5, 2, 15, 12, 9, 6, 3]

    def leakage(self, pt, ct, key, bnum):
        state = [pt[i] ^ key[i] for i in range(0, 16)]
        state1 = state[:]
        state = self.subbytes(state)
        state = self.shiftrows(state)
        state = self.mixcolumns(state)
        key2 = self.key_schedule_rounds(key, 0, 1)
        state = [state[i] ^ key2[i] for i in range(0, 16)]
        output_pos = self.KEY_TO_OUTPUT[bnum]  # FIX
        return state[output_pos] ^ state1[output_pos]

Any thoughts about this issue?