AES128 in Ruby 1.9.2

aes128.rb

# coding: utf-8
# Short code to explain the AES128 (Rijndael) cipher
# Some array operations require Ruby 1.9.2+
# Steps defined by http://en.wikipedia.org/wiki/Advanced_Encryption_Standard

class AES128

  # Mix Rows coefficients
  MIX = [
    [2, 3, 1, 1],
    [1, 2, 3, 1],
    [1, 1, 2, 3],
    [3, 1, 1, 2]
  ]

  # Multiplicative inverse in GF(2)[x]/(x**8 + x**4 + x**3 + x + 1)
  S_BOX = %w[
    63 7c 77 7b f2 6b 6f c5 30 01 67 2b fe d7 ab 76
    ca 82 c9 7d fa 59 47 f0 ad d4 a2 af 9c a4 72 c0
    b7 fd 93 26 36 3f f7 cc 34 a5 e5 f1 71 d8 31 15
    04 c7 23 c3 18 96 05 9a 07 12 80 e2 eb 27 b2 75
    09 83 2c 1a 1b 6e 5a a0 52 3b d6 b3 29 e3 2f 84
    53 d1 00 ed 20 fc b1 5b 6a cb be 39 4a 4c 58 cf
    d0 ef aa fb 43 4d 33 85 45 f9 02 7f 50 3c 9f a8
    51 a3 40 8f 92 9d 38 f5 bc b6 da 21 10 ff f3 d2
    cd 0c 13 ec 5f 97 44 17 c4 a7 7e 3d 64 5d 19 73
    60 81 4f dc 22 2a 90 88 46 ee b8 14 de 5e 0b db
    e0 32 3a 0a 49 06 24 5c c2 d3 ac 62 91 95 e4 79
    e7 c8 37 6d 8d d5 4e a9 6c 56 f4 ea 65 7a ae 08
    ba 78 25 2e 1c a6 b4 c6 e8 dd 74 1f 4b bd 8b 8a
    70 3e b5 66 48 03 f6 0e 61 35 57 b9 86 c1 1d 9e
    e1 f8 98 11 69 d9 8e 94 9b 1e 87 e9 ce 55 28 df
    8c a1 89 0d bf e6 42 68 41 99 2d 0f b0 54 bb 16
  ].map {|x| x.to_i 16 }

  # rcon(i) = 2 ** (i - 1)
  # Note the exponenential operation is defined by galois multiplication
  RCON = %w[
    8d 01 02 04 08 10 20 40 80 1b 36 6c d8 ab 4d 9a
  ].map {|x| x.to_i 16 }

  def xor v1, v2
    v1.zip(v2).map {|(a, b)| a ^ b }
  end

  def shift_rows bytes
    i = -1
    bytes = bytes.each_slice(4).to_a.transpose.map do |v|
      i += 1
      v.rotate i
    end
    bytes.transpose.flatten
  end

  # galois multiplication
  def gmul byte, n
    b =\
      case n
      when 1; return byte
      when 2; byte << 1
      when 3; (byte << 1) ^ byte
      end
    b > 255 ? ((b & 0xFF) ^ 0x1B) : b
  end

  def mix_columns bytes
    # raise "size not right: #{bytes}" if bytes.size != 16
    bytes.each_slice(4).flat_map do |slice|
      MIX.map do |coef|
        r = 0
        slice.zip(coef) {|b, n| r ^= gmul b, n}
        r
      end
    end
  end

  def sub_bytes bytes
    bytes.map {|b| S_BOX[b] }
  end

  # [key, round_key1, ..., round_key10]
  def key_schedule key
    cols = key.each_slice(4).to_a
    4.upto 43 do |i|
      col = cols[i - 1]
      if i % 4 == 0
        col = col.rotate
        col = sub_bytes col
        col[0] = col[0] ^ RCON[i / 4]
      end
      cols << xor(col, cols[i - 4])
    end
    cols.each_slice(4).map(&:flatten)
  end

  def encrypt128 bytes
    bytes = xor bytes, @rkeys[0] # add round key
    1.upto 9 do |i|
      bytes = sub_bytes bytes
      bytes = shift_rows bytes
      bytes = mix_columns bytes
      bytes = xor bytes, @rkeys[i]
    end
    bytes = sub_bytes bytes
    bytes = shift_rows bytes
    xor bytes, @rkeys[10]
  end

  def encrypt s
    raise 'state size should be 16 bytes' if s.bytesize != 16
    encrypt128(s.bytes).pack 'C*'
  end

  def initialize key
    raise 'key size should be 16 bytes' if key.bytesize != 16
    @rkeys = key_schedule key.bytes
  end

  attr_reader :rkeys
end

if $0 == __FILE__
  aes = AES128.new 'b' * 16
  res = aes.encrypt 'h' * 16
  p res.bytes.to_a
end