A handwaving implementation of AES Polynomials in Ruby. Don't use this in production. SRSLY, don't!

aes.rb

# frozen_string_literal: true

class Integer
  def to_aes
    Aes::Polynomial.from(self)
  end

  def to_hex_s
    hex = to_s(16).upcase

    hex.prepend('0') until hex.length.even?

    hex.scan(/.{2}/).join(' ')
  end
end

module Aes
  class Polynomial
    # We use a separate class for the coefficients in order to avoid
    # an infinite loop problem:
    #
    # We want to reduce every polynomial down to its representative of minimal degree,
    # BUT we must not do that with the AES modulus itself, because this would require
    # us to instantiate the modulus first, which would then trigger a reduction,
    # which would then again instantiate the modulus, etc.
    #
    # Hence we implement the coefficients as unpruned arrays, i.e.
    #
    #    1 + x^2 + x^5 == [1,0,1,0,0,1] == [1,0,1,0,0,1,0,0,0]
    #
    # and wrap them in a Aes::Polynomial class which handles reduction.
    class Coefficients
      class << self
        def from(integer)
          raise ArgumentError if integer.negative?

          coefficients = []

          while integer.positive?
            coefficients.push(integer % 2)
            integer /= 2
          end

          new(coefficients)
        end

        def zero
          from(0x00)
        end
      end

      attr_reader :coefficients

      def initialize(array)
        @coefficients = array.map(&:abs)
      end

      def [](index)
        @coefficients[index] || 0
      end

      def +(other)
        n = [0, degree, other.degree].max

        coefficients = Array.new(n + 1) do |i|
          (self[i] + other[i]) % 2
        end

        self.class.new(coefficients)
      end
      alias - +

      def *(other)
        n = [degree + other.degree, 0].max

        coefficients = Array.new(n + 1) do |i|
          (i + 1).times.sum do |j|
            self[j] * other[i - j]
          end % 2
        end

        self.class.new(coefficients)
      end

      def mod(other)
        quotient = self.class.zero
        remainder = dup

        while remainder.degree >= other.degree
          factor = self.class.from(2**(remainder.degree - other.degree))

          quotient += factor
          remainder -= factor * other
        end

        [quotient, remainder]
      end

      def /(other)
        mod(other).first
      end

      def %(other)
        mod(other).last
      end

      def ==(other)
        return false unless other.is_a?(self.class)

        to_i == other.to_i
      end

      def degree
        return -Float::INFINITY if zero?

        @coefficients.length - 1 - @coefficients.reverse.find_index { |x| !x.zero? }
      end

      def inspect
        @coefficients.reverse.inspect
      end

      def reduce!
        @coefficients = (self % MODULUS).coefficients
      end

      def to_i
        @coefficients.map.with_index { |c, i| c * 2**i }.sum
      end

      def to_s
        return '0' if zero?

        @coefficients.map.with_index do |c, i|
          next if c.zero?

          case i
          when 0 then '1'
          when 1 then 'x'
          else
            "x**#{i}"
          end
        end.compact.reverse.join(' + ')
      end

      def zero?
        @coefficients.empty? || @coefficients.all?(&:zero?)
      end
    end

    MODULUS = Coefficients.from(0x11B).freeze
  end
end

module Aes
  class Polynomial
    class << self
      def from(value)
        value.is_a?(Polynomial) ? value : Polynomial.new(value)
      end

      def zero
        new(0x00)
      end
    end

    attr_reader :coefficients

    def initialize(value, reduce: true)
      @coefficients = case value
                      when Coefficients
                        value
                      when Integer
                        Coefficients.from(value)
                      else
                        raise ArgumentError
                      end

      self.reduce if reduce
    end

    def +(other)
      self.class.new(@coefficients + other.coefficients)
    end

    alias - +

    def *(other)
      self.class.new(@coefficients * other.coefficients)
    end

    def mod(other)
      @coefficients.mod(other.coefficients).map { |i| self.class.new(i) }
    end

    def /(other)
      mod(other).first
    end

    def %(other)
      mod(other).last
    end

    def ==(other)
      return false unless other.is_a?(self.class)

      @coefficients == other.coefficients
    end

    def degree
      @coefficients.degree
    end

    def inspect
      "#{to_hex_s} (#{self})"
    end

    def to_i
      @coefficients.to_i
    end

    def to_hex_s
      to_i.to_hex_s
    end

    def to_s
      @coefficients.to_s
    end

    def zero?
      @coefficients.zero?
    end

    private

    def reduce
      @coefficients.reduce!
    end
  end
end

# rubocop:disable Lint/BinaryOperatorWithIdenticalOperands

RSpec.describe Aes::Polynomial do # rubocop:disable Metrics/BlockLength
  describe '#+' do
    it 'adds correctly' do
      expect(0b00.to_aes + 0b00.to_aes).to eq 0b00.to_aes
      expect(0b00.to_aes + 0b01.to_aes).to eq 0b01.to_aes
      expect(0b01.to_aes + 0b01.to_aes).to eq 0b00.to_aes
      expect(0b10.to_aes + 0b01.to_aes).to eq 0b11.to_aes
    end
  end

  describe '#-' do
    it 'subtracts correctly' do
      expect(0b00.to_aes - 0b00.to_aes).to eq 0b00.to_aes
      expect(0b00.to_aes - 0b01.to_aes).to eq 0b01.to_aes
      expect(0b01.to_aes - 0b01.to_aes).to eq 0b00.to_aes
      expect(0b10.to_aes - 0b01.to_aes).to eq 0b11.to_aes
    end
  end

  describe '#*' do
    it 'multiplies correctly' do
      expect(0b00.to_aes * 0b00.to_aes).to eq 0b00.to_aes
      expect(0b00.to_aes * 0b01.to_aes).to eq 0b00.to_aes
      expect(0b01.to_aes * 0b01.to_aes).to eq 0b01.to_aes
      expect(0b10.to_aes * 0b01.to_aes).to eq 0b10.to_aes
      expect(0b10.to_aes * 0b10.to_aes).to eq 0b0100.to_aes
    end
  end

  describe '#/' do
    it 'divides correctly' do
      expect(0b1110.to_aes / 0b0010.to_aes).to eq(0b0111.to_aes)
      expect(0b1111.to_aes / 0b0010.to_aes).to eq(0b0111.to_aes)
      expect(0b0001_0110.to_aes / 0b101.to_aes).to eq(0b100.to_aes)
    end
  end

  describe '#%' do
    it 'calculates the remainder correctly' do
      expect(0b1110.to_aes % 0b0010.to_aes).to eq(0b0000.to_aes)
      expect(0b1111.to_aes % 0b0010.to_aes).to eq(0b0001.to_aes)
      expect(0b0001_0110.to_aes % 0b101.to_aes).to eq(0b0010.to_aes)
    end
  end

  describe '#==' do
    it 'compares correctly' do
      expect(0b00.to_aes).to eq(0b00.to_aes)
      expect(0b1_0001_1011.to_aes).to eq(0b00.to_aes)
      expect(0b1_0000_0000.to_aes).to eq(0b0_0001_1011.to_aes)
    end
  end

  describe '#degree' do
    it 'returns the degree of the polynomial' do
      expect(0b0000.to_aes.degree).to eq(-Float::INFINITY)
      expect(0b0001.to_aes.degree).to eq(0)
      expect(0b0010.to_aes.degree).to eq(1)
      expect(0b0101.to_aes.degree).to eq(2)
      expect(0b1010.to_aes.degree).to eq(3)
      expect(0b1111.to_aes.degree).to eq(3)
      expect(0b1111_1111.to_aes.degree).to eq(7)
      expect(0b1_0000_0000.to_aes.degree).to eq(4)
    end
  end

  describe '#inspect' do
    it 'returns the polynomial in mathematical and hex notation' do
      expect(0b0000.to_aes.inspect).to eq('00 (0)')
      expect(0b0001.to_aes.inspect).to eq('01 (1)')
      expect(0b0010.to_aes.inspect).to eq('02 (x)')
      expect(0b0011.to_aes.inspect).to eq('03 (x + 1)')
      expect(0b0111.to_aes.inspect).to eq('07 (x**2 + x + 1)')
      expect(0b0001_0110.to_aes.inspect).to eq('16 (x**4 + x**2 + x)')
      expect(0b1_0001_1011.to_aes.inspect).to eq('00 (0)')
      expect(0b1_0000_0000.to_aes.inspect).to eq('1B (x**4 + x**3 + x + 1)')
    end
  end

  describe '#to_s' do
    it 'returns the polynomial in mathematical notation' do
      expect(0b0000.to_aes.to_s).to eq('0')
      expect(0b0001.to_aes.to_s).to eq('1')
      expect(0b0010.to_aes.to_s).to eq('x')
      expect(0b0011.to_aes.to_s).to eq('x + 1')
      expect(0b0111.to_aes.to_s).to eq('x**2 + x + 1')
      expect(0b0001_0110.to_aes.to_s).to eq('x**4 + x**2 + x')
      expect(0b1_0001_1011.to_aes.to_s).to eq('0')
      expect(0b1_0000_0000.to_aes.to_s).to eq('x**4 + x**3 + x + 1')
    end
  end
end

# rubocop:enable Lint/BinaryOperatorWithIdenticalOperands