aes_kw.rb in jwe-0.3.1

- old
+ new

@@ -1,7 +1,10 @@
+require 'jwe/enc/cipher'
+
 module JWE
   module Alg
+    # Generic AES Key Wrapping algorithm for any key size.
     module AesKw
       attr_accessor :key
       attr_accessor :iv
 
       def initialize(key = nil, iv = "\xA6\xA6\xA6\xA6\xA6\xA6\xA6\xA6")
@@ -9,55 +12,63 @@
         self.key = key.force_encoding('ASCII-8BIT')
       end
 
       def encrypt(cek)
         a = iv
-        r = cek.scan(/.{8}/m)
+        r = cek.force_encoding('ASCII-8BIT').scan(/.{8}/m)
 
         6.times do |j|
-          r.length.times do |i|
-            b = encrypt_round(a + r[i])
-
-            a = b.chars.first(8).join
-            r[i] = b.chars.last(8).join
-
-            t = (r.length * j) + i + 1
-            a = xor(a, t)
-          end
+          a, r = kw_encrypt_round(j, a, r)
         end
 
         ([a] + r).join
       end
 
-      def decrypt(encrypted_cek)
-        c = encrypted_cek.scan(/.{8}/m)
-        a = c[0]
+      def kw_encrypt_round(j, a, r)
+        r.length.times do |i|
+          b = encrypt_round(a + r[i]).chars
 
-        r = c[1..c.length]
+          a, r[i] = a_ri(b)
 
-        5.downto(0) do |j|
-          r.length.downto(1) do |i|
-            t = (r.length * j) + i
-            a = xor(a, t)
+          a = xor(a, (r.length * j) + i + 1)
+        end
 
-            b = decrypt_round(a + r[i - 1])
+        [a, r]
+      end
 
-            a = b.chars.first(8).join
-            r[i - 1] = b.chars.last(8).join
-          end
+      def decrypt(encrypted_cek)
+        c = encrypted_cek.force_encoding('ASCII-8BIT').scan(/.{8}/m)
+        a, *r = c
+
+        5.downto(0) do |j|
+          a, r = kw_decrypt_round(j, a, r)
         end
 
         if a != iv
           raise StandardError.new('The encrypted key has been tampered. Do not use this key.')
         end
 
         r.join
       end
 
+      def kw_decrypt_round(j, a, r)
+        r.length.downto(1) do |i|
+          a = xor(a, (r.length * j) + i)
+
+          b = decrypt_round(a + r[i - 1]).chars
+
+          a, r[i - 1] = a_ri(b)
+        end
+
+        [a, r]
+      end
+
+      def a_ri(b)
+        [b.first(8).join, b.last(8).join]
+      end
+
       def cipher
-        @cipher ||= OpenSSL::Cipher.new(cipher_name)
-      rescue RuntimeError
-        raise JWE::NotImplementedError.new("The version of OpenSSL linked to your Ruby does not support the cipher #{cipher_name}.")
+        @cipher ||= Enc::Cipher.for(cipher_name)
       end
 
       def encrypt_round(data)
         cipher.encrypt
         cipher.key = key