SelectionVectors.py

import os
import itertools

from sage.all import VectorSpace, Matrix, GF
from tqdm import tqdm

AESSBox = [
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16
          ]

def selectionFunction(plaintextByte, keyByteGuess, mask):
    return LINEAR_MAPS[mask][AESSBox[plaintextByte ^ keyByteGuess]]

def getPlaintexts(pathToTraces):
    SELorBU=0 #SEL traces = 1, BU traces = 2, not found = 0
    try:
        f = open(pathToTraces / "plaintext.txt", "rt").readlines()
        SELorBU=1
    except IOError:
        pass
    try:
        numOfBytes = os.path.getsize(pathToTraces / "0000.pt")
        if SELorBU==1:
            print("/!\\ Traces of SEL masking scheme and BU masking scheme are located in the same folder /!\\")
            print("Please separate \"traces.txt\" into a different one and relaunch the program.")
            exit()
        else:
            SELorBU=2
    except IOError:
        pass
    if not SELorBU:
        print("/!\\ The given path to the traces directory does not contain the file \"0000.bin\" or \"plaintext.txt\" /!\\")
        print("  Either the path is not correct, either the traces are not in the correct format")
        exit()

    PLAINTEXTS = []

    if SELorBU==1:
        T = len(f)
        for traceNumber in range(T):
            t=0
            i=0
            plaintext=[]
            while f[traceNumber][i]!='\n':
                t=0
                while f[traceNumber][i]!=' ':
                    t*=10
                    t+=int(f[traceNumber][i])
                    i+=1
                plaintext.append(t)
                i+=1
            PLAINTEXTS.append(plaintext)
    else:
        T = 0
        while True:
            fpt = pathToTraces / ("%04d.pt" % T)
            try:
                with open(fpt, "rb") as file:
                    try:
                        PLAINTEXTS.append(file.read(16))
                    except IOError as err:
                        print("/!\\ The file %04d.pt contains less than 16 bytes /!\\" % T)
                        print("Exiting the program.")
                        exit()
                T+=1
                file.close()
            except IOError as err:
                break

    return(PLAINTEXTS)

def getSelectionVectors(PLAINTEXTS, T, mode='list', bytePositions=(0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15), masks=(1, 2, 4, 8, 16, 32, 64, 128)):
    SELECTIONVECTORS=[]
    INFO = []

    if T>len(PLAINTEXTS):
        print("Impossible to read Plaintexts:")
        print("Only %d traces available, cannot get T=%d traces" % (len(PLAINTEXTS), T))
        print("Exiting the program.")
        exit()

    print("generating selection vectors:", mode)

    if mode=='list':
        for keyByteGuess in range(256):
            vectPos=[]
            for bytePosition in range(16):
                vectTrace=[]
                for traceNumber in range(T):
                    vectTrace.append(selectionFunction(PLAINTEXTS[traceNumber][bytePosition], keyByteGuess, mask=masks[-1]))
                vectPos.append(vectTrace)
            SELECTIONVECTORS.append(vectPos)
        return(SELECTIONVECTORS)

    elif mode=='vect':
        VS=VectorSpace(GF(2), T)
        for keyByteGuess in range(256):
            vectPos=[]
            for bytePosition in range(16):
                vectTrace=[]
                for traceNumber in range(T):
                    vectTrace.append(selectionFunction(PLAINTEXTS[traceNumber][bytePosition], keyByteGuess, mask=masks[-1]))
                vectPos.append(VS(vectTrace))
            SELECTIONVECTORS.append(vectPos)
        return(SELECTIONVECTORS)

    elif mode=='mat':
        bytes_keys = list(itertools.product(bytePositions,list(range(256))))
        for bytePosition, keyByteGuess in tqdm(bytes_keys):
            for mask in masks:
                selVect=[]
                for traceNumber in range(T):
                    selVect.append(selectionFunction(PLAINTEXTS[traceNumber][bytePosition], keyByteGuess, mask=mask))
                SELECTIONVECTORS.append(Matrix(GF(2), 1, len(selVect), selVect))
                INFO.append((bytePosition, keyByteGuess, mask))
        return SELECTIONVECTORS, INFO

    elif mode=='fulmat':
        bytes_keys = list(itertools.product(bytePositions,list(range(256))))
        for bytePosition, keyByteGuess in tqdm(bytes_keys):
            for mask in masks:
                selVect=[]
                for traceNumber in range(T):
                    selVect.append(selectionFunction(PLAINTEXTS[traceNumber][bytePosition], keyByteGuess, mask=mask))
                SELECTIONVECTORS.append(selVect)
                INFO.append((bytePosition, keyByteGuess, mask))
        return Matrix(GF(2),SELECTIONVECTORS), INFO

    print("Unrecognized mode for the getSelectionVectors function")
    print("Existing modes are 'list', 'vect', 'mat' and 'fulmat'")
    print("Exiting the program.")
    exit()


def scalar_bin(a, b):
    v = a & b
    res = 0
    while v:
        res ^= v & 1
        v >>= 1
    return res


LINEAR_MAPS = [
    [scalar_bin(x, mask) for x in range(256)]
    for mask in range(256)
]

def reportKeyMatch(info, **extra):
    bytePosition, keyByteGuess, mask = info
    extra = " ".join(f"{k}={v}" for k, v in extra.items())
    print(f"Match Byte#{bytePosition:02d} Key {keyByteGuess:02x} ({keyByteGuess}) mask {mask:08b} ({extra})")