Create C header file describing an AES substitution box given an arbitrary irreducible polynomial which defines the finite field GF(2^8)

gen_from_poly.py

import sys

from gf2 import GF2, GF2Element

def append_c_header(path, s):
    with open(path, 'a') as o:
        o.write(s)

def clear_c_header(path):
    open(path, 'w').close()


pad_hex = lambda x: "0x" + hex(x)[2:].rjust(2, '0')

def to_c_arr(values, prefix=""):

    result = prefix + "{\n"
    for i in range(0, len(values), 16):
        slice_ = values[i:(i+16)]
        result += prefix +  '    ' + ', '.join(pad_hex(x) for x in slice_) + ",\n"

    result += prefix + "}"
    return result



def create_gf(polynomial):
    """
    Create inverse and multiplication table of GF(2^8) with defining polynomial `polynomial`
    """

    field = GF2(8, polynomial)

    mult = [None] * (0x100 * 0x100)

    for i in range(0, 0x100):
        for j in range(0, 0x100):
            mult[i*0x100+j] = (field.element(i) * field.element(j)).value

    return {"mult": mult}


def create_sboxes(polynomial):

    field = GF2(8, polynomial)

    rol8 = lambda x, n: ((x << n) | ( x >> (8 - n)))

    def s_transform(e):
            if e != 0:
                e = field.element(e).inverse.value

            return field.element(e ^ rol8(e, 1) ^ rol8(e, 2) ^ rol8(e, 3) ^ rol8(e, 4) ^ 0x63)

    sbox = [s_transform(i).value for i in range(0x100)]

    return {"fwd_sbox": sbox}


if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: ./gen_from_poly.py <polynomial> <sbox out> <gf out>")
        exit()

    polynomial = int(sys.argv[1])

    sbox_out = sys.argv[2]
    gf_out = sys.argv[3]

    clear_c_header(gf_out)
    append_c_header(gf_out, "#pragma once\n")
    append_c_header(gf_out, "#include <stdint.h>\n")
    append_c_header(gf_out, "typedef uint8_t gf_t;\n\n")

    for k, v in create_gf(polynomial).items():
        append_c_header(gf_out, f"gf_t gf_28_{k}[] = " + to_c_arr(v) + ";\n\n")


    clear_c_header(sbox_out)
    append_c_header(sbox_out, "#pragma once\n")
    append_c_header(sbox_out, "#include <gf/gf.h>\n\n")

    for k, v in create_sboxes(polynomial).items():
        append_c_header(sbox_out, f"gf_t aes_{k}[] = " + to_c_arr(v) + ";\n\n")