<commit_before>""" AES-128 """

from __future__ import division, absolute_import
import pyrtl
# from pyrtl import rtllib
# from pyrtl.rtllib import libutils
# import sys
# sys.path.append('../..')
from pyrtl.rtllib import libutils

# TODO:
# 1) Right now all ROMs are global -- need to be moved into function so each
#    instance of the AES block creates it's own memories
# 2) All ROMs should be syncronous.  This should be easy once (3) is completed
# 3) Right now aes_decryption generates one GIANT combinational block. Instead
#    it should generate one of 2 options -- Either an interative design or a
#    pipelined design.  Both will add registers between each round of AES
# 4) aes_encryption should be added to this file as well so that an
#    aes encrypter similar to (3) above is generated
# 5) a single "aes-unit" combining encryption and decryption (without making
#    full independent hardware units) would be a plus as well
# 6) The constants right now are a big ugly list -- it might be better to
#    store them as a multi-line string like: '0f bf 53 2c ...' and then
#    provide a small helper funciton to parse them out.  It might be a nice
#    utility for pyrtl itself to include.


class AES(object):
    def __init__(self):
        self.memories_built = False

    def _g(self, word, key_expand_round):
        # change key_expand_round into a wirevector
        self.build_memories_if_not_exists()
        # One-byte left circular rotation, substitution of each byte
        a = libutils.partition_wire(word, 8)
        sub = pyrtl.concat(self.sbox[a[2]], self.sbox[a[1]], self.sbox[a[0]], self.sbox[a[3]])
        # round_const = pyrtl.Register(32, 'round_const')
        # xor substituted bytes with round constant.
        round_const = pyrtl.concat(self.rcon[key_expand_round + 1], pyrtl.Const(0, bitwidth=24))
        # round_const.next <<= round_const ^ sub
        return round_const 

    def key_expansion(self, old_key, key_expand_round):
        self.build_memories_if_not_exists()
        w = list(libutils.partition_wire(old_key, 32))
        # for key_expand_round in range(10):
        last = key_expand_round * 4
        w.append(w[0] ^ self._g(w[3], key_expand_round))
        w.append(w[-1] ^ w[1])
        w.append(w[-1] ^ w[2])
        w.append(w[-1] ^ w[3])
        return pyrtl.concat(*w[4:8])

    def inv_sub_bytes(self, in_vector):
        self.build_memories_if_not_exists()
        subbed = [self.inv_sbox[byte] for byte in libutils.partition_wire(in_vector, 8)]
        return pyrtl.concat(*reversed(subbed))

    def inv_shift_rows(self, in_vector):
        # a = libutils.partition_wire(in_vector, 8)
        a = [in_vector[offset - 8:offset] for offset in range(128, 0, -8)]
        out_vector = pyrtl.concat(a[0], a[13], a[10], a[7],
                                  a[4], a[1], a[14], a[11],
                                  a[8], a[5], a[2], a[15],
                                  a[12], a[9], a[6], a[3])
        return out_vector

    def inv_galois_mult(self, c, d):
        return self._inv_gal_mult_dict[d][c]

    @staticmethod
    def _mod_add(base, add, mod):
        base_mod_floor = base // mod
        return (base + add) % mod + base_mod_floor * mod

    _igm_divisor = [14, 11, 13, 9]

    def inv_mix_columns(self, in_vector):
        self.build_memories_if_not_exists()
        def _inv_mix_single(index):
            mult_items = [self.inv_galois_mult(a[self._mod_add(index, loc, 4)], mult_table)
                          for loc, mult_table in enumerate(self._igm_divisor)]
            return mult_items[0] ^ mult_items[1] ^ mult_items[2] ^ mult_items[3]

        # a = libutils.partition_wire(in_vector, 8)
        a = [in_vector[offset - 8:offset] for offset in range(128, 0, -8)]
        inverted = [_inv_mix_single(index) for index in range(len(a))]
        return pyrtl.concat(*inverted)

    def addroundkey(self, t, key):
        return t ^ key

    def aes__decryption_states(self, ciphertext_in, key, reset):
        """
        return ready, decryption_result: ready is a one bit signal showing that the answer decryption 
        result has been calculated.
        """
        # substates of the calcul
        # create the counter here
        counter = pyrtl.Register(4, 'counter')
        cipher_text = pyrtl.Register(128)
        current_key = pyrtl.Register(128)
        key_expansion_in = WireVector()
        round = WireVector()
        add_round_in = WireVector()
        # the hardware for a given round
        counter.next <<= round
        inv_shift = self.inv_shift_rows(cipher_text)
        inv_sub = self.inv_sub_bytes(inv_shift)
        key_exp_out = key_expansion(self, key_expansion_in, round)
        add_round_out = addroundkey(inv_sub, add_round_in)
        inv_mix_out = self.inv_mix_columns(add_round_out)
        current_key.next <<= key_exp_out

        key_exp_out = key_expansion(self, key_expansion_in, round)
        with conditional_assignment:
            with reset == 1:
                round |= 0
                key_expansion_in |= key
                cipher_text.next |= add_round_out
                add_round_in  |= ciphertext_in
                # reset everything to initial values
                    
            with counter == 11:
                # keep everything the same
                round |= counter
                cipher_text.next |= cipher_text
            # might be some special thing to do durnign round 1 or 10

            with pyrtl.otherwise:
                round |= counter + 1
                key_expansion_in |= current_key
                add_round_in |= key_exp_out 
                with counter == 10:
                    cipher_text.next |=add_round_out 
                with pyrtl.otherwise:
                    cipher_text.next |= inv_mix_out

        ready = (counter == 11)
        return ready, cipher_text
        # for round in range(1,11):
        #     expanded_key = pyrtl.WireVector(128, 'expanded_key')
        #     key_state = pyrtl.Register(128, 'key_state')
        #     key_state <<= expanded_key
        #     key_state.next <<= key_expansion(self, expanded_key, round)

    def aes_decryption(self, ciphertext, key):
        expanded_key = self.key_expansion(key)  # Expanding the key (key expansion).
        t = self.addroundkey(ciphertext, expanded_key, 0)  # Initial AddRoundKey.

        for round in range(1, 11):
            t = self.inv_shift_rows(t)
            t = self.inv_sub_bytes(t)
            t = self.addroundkey(t, expanded_key)
            if round != 10:
                t = self.inv_mix_columns(t)
<commit_msg>Lots of updates to AES (fixing things)<commit_after>""" AES-128 """

from __future__ import division, absolute_import
import pyrtl
# from pyrtl import rtllib
# from pyrtl.rtllib import libutils
# import sys
# sys.path.append('../..')
from pyrtl.rtllib import libutils

# TODO:
# 1) Right now all ROMs are global -- need to be moved into function so each
#    instance of the AES block creates it's own memories
# 2) All ROMs should be syncronous.  This should be easy once (3) is completed
# 3) Right now aes_decryption generates one GIANT combinational block. Instead
#    it should generate one of 2 options -- Either an interative design or a
#    pipelined design.  Both will add registers between each round of AES
# 4) aes_encryption should be added to this file as well so that an
#    aes encrypter similar to (3) above is generated
# 5) a single "aes-unit" combining encryption and decryption (without making
#    full independent hardware units) would be a plus as well
# 6) The constants right now are a big ugly list -- it might be better to
#    store them as a multi-line string like: '0f bf 53 2c ...' and then
#    provide a small helper funciton to parse them out.  It might be a nice
#    utility for pyrtl itself to include.


class AES(object):
    def __init__(self):
        self.memories_built = False

    def _g(self, word, key_expand_round):
        # change key_expand_round into a wirevector
        self.build_memories_if_not_exists()
        # One-byte left circular rotation, substitution of each byte
        a = libutils.partition_wire(word, 8)
        sub = pyrtl.concat(self.sbox[a[2]], self.sbox[a[1]], self.sbox[a[0]], self.sbox[a[3]])
        round_const = pyrtl.concat(self.rcon[key_expand_round + 1], pyrtl.Const(0, bitwidth=24))
        return round_const

    def key_expansion(self, old_key, key_expand_round):
        self.build_memories_if_not_exists()
        w = list(libutils.partition_wire(old_key, 32))
        w.append(w[0] ^ self._g(w[3], key_expand_round))
        w.append(w[-1] ^ w[1])
        w.append(w[-1] ^ w[2])
        w.append(w[-1] ^ w[3])
        return pyrtl.concat(*w[4:8])

    def inv_sub_bytes(self, in_vector):
        self.build_memories_if_not_exists()
        subbed = [self.inv_sbox[byte] for byte in libutils.partition_wire(in_vector, 8)]
        return pyrtl.concat(*reversed(subbed))

    def inv_shift_rows(self, in_vector):
        # a = libutils.partition_wire(in_vector, 8)
        a = [in_vector[offset - 8:offset] for offset in range(128, 0, -8)]
        out_vector = pyrtl.concat(a[0], a[13], a[10], a[7],
                                  a[4], a[1], a[14], a[11],
                                  a[8], a[5], a[2], a[15],
                                  a[12], a[9], a[6], a[3])
        return out_vector

    def inv_galois_mult(self, c, d):
        return self._inv_gal_mult_dict[d][c]

    @staticmethod
    def _mod_add(base, add, mod):
        base_mod_floor = base // mod
        return (base + add) % mod + base_mod_floor * mod

    _igm_divisor = [14, 11, 13, 9]

    def inv_mix_columns(self, in_vector):
        self.build_memories_if_not_exists()
        def _inv_mix_single(index):
            mult_items = [self.inv_galois_mult(a[self._mod_add(index, loc, 4)], mult_table)
                          for loc, mult_table in enumerate(self._igm_divisor)]
            return mult_items[0] ^ mult_items[1] ^ mult_items[2] ^ mult_items[3]

        # a = libutils.partition_wire(in_vector, 8)
        a = [in_vector[offset - 8:offset] for offset in range(128, 0, -8)]
        inverted = [_inv_mix_single(index) for index in range(len(a))]
        return pyrtl.concat(*inverted)

    def addroundkey(self, t, key):
        return t ^ key

    def aes_decryption_statem(self, ciphertext_in, key, reset):
        """
        return ready, decryption_result: ready is a one bit signal showing that the answer decryption 
        result has been calculated.
        """
        if len(key) != len(ciphertext_in):
            raise pyrtl.PyrtlError("AES key and ciphertext should be the same length")

        cipher_text, current_key = (pyrtl.Register(len(key)) for i in range(2))
        key_expansion_in, add_round_in = (pyrtl.WireVector(len(key)) for i in range(2))

        counter = pyrtl.Register(4, 'counter')
        round = pyrtl.WireVector(4)
        counter.next <<= round

        inv_shift = self.inv_shift_rows(cipher_text)
        inv_sub = self.inv_sub_bytes(inv_shift)
        key_exp_out = self.key_expansion(key_expansion_in, round)
        add_round_out = self.addroundkey(inv_sub, add_round_in)
        inv_mix_out = self.inv_mix_columns(add_round_out)

        current_key.next <<= key_exp_out

        with pyrtl.conditional_assignment:
            with reset == 1:
                round |= 0
                key_expansion_in |= key
                cipher_text.next |= add_round_out
                add_round_in |= ciphertext_in
                # reset everything to initial values
                    
            with counter == 11:  # keep everything the same
                round |= counter
                cipher_text.next |= cipher_text
                # might be some special thing to do durnign round 1 or 10

            with pyrtl.otherwise:  # running through AES
                round |= counter + 1
                key_expansion_in |= current_key
                add_round_in |= key_exp_out 
                with counter == 10:
                    cipher_text.next |= add_round_out
                with pyrtl.otherwise:
                    cipher_text.next |= inv_mix_out

        ready = (counter == 11)
        return ready, cipher_text
        # for round in range(1,11):
        #     expanded_key = pyrtl.WireVector(128, 'expanded_key')
        #     key_state = pyrtl.Register(128, 'key_state')
        #     key_state <<= expanded_key
        #     key_state.next <<= key_expansion(self, expanded_key, round)

    def aes_decryption(self, ciphertext, key):
        expanded_key = self.key_expansion(key)  # Expanding the key (key expansion).
        t = self.addroundkey(ciphertext, expanded_key, 0)  # Initial AddRoundKey.

        for round in range(1, 11):
            t = self.inv_shift_rows(t)
            t = self.inv_sub_bytes(t)
            t = self.addroundkey(t, expanded_key)
            if round != 10:
                t = self.inv_mix_columns(t)

{'fraction_non_alphanumeric': 0.0973097617217525, 'fraction_numerical': 0.018601076095311298, 'mean_word_length': 2.3378655720882504, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6844817', 'n_tokens_mistral': 4346, 'n_tokens_neox': 4012, 'n_words': 1128}