refactoring AES part and add specs

lib/diaspora_federation/salmon/aes.rb

@@ -1,42 +1,55 @@
 module DiasporaFederation
   module Salmon
+    # class for AES encryption and decryption
     class AES
       # OpenSSL aes cipher definition
       CIPHER = "AES-256-CBC"
 
-      # encrypts the given data with a new, random AES cipher and returns the
-      # resulting ciphertext, the key and iv in a hash (each of the entries
-      # base64 strict_encoded).
+      # generates a random AES key and initialization vector
+      # @return [Hash] { key: "...", iv: "..." }
+      def self.generate_key_and_iv
+        cipher = OpenSSL::Cipher.new(CIPHER)
+        {key: cipher.random_key, iv: cipher.random_iv}
+      end
+
+      # encrypts the given data with an AES cipher defined by the given key
+      # and iv and returns the resulting ciphertext base64 strict_encoded.
       # @param [String] data plain input
-      # @return [Hash] { key: "...", iv: "...", ciphertext: "..." }
-      def self.encrypt(data)
+      # @param [String] key AES key
+      # @param [String] iv AES initialization vector
+      # @param [Hash] key_and_iv { key: "...", iv: "..." }
+      # @return [String] base64 encoded ciphertext
+      def self.encrypt(data, key, iv)
+        raise ArgumentError unless data.instance_of?(String)
+
         cipher = OpenSSL::Cipher.new(CIPHER)
         cipher.encrypt
-        key = cipher.random_key
-        iv = cipher.random_iv
+        cipher.key = key
+        cipher.iv = iv
+
         ciphertext = cipher.update(data) + cipher.final
 
-        enc = [key, iv, ciphertext].map {|i| Base64.strict_encode64(i) }
-
-        {key: enc[0], iv: enc[1], ciphertext: enc[2]}
+        Base64.strict_encode64(ciphertext)
       end
 
       # decrypts the given ciphertext with an AES cipher defined by the given key
-      # and iv. parameters are expected to be base64 encoded
+      # and iv. +ciphertext+ is expected to be base64 encoded
       # @param [String] ciphertext input data
       # @param [String] key AES key
       # @param [String] iv AES initialization vector
       # @return [String] decrypted plain message
+      # @raise [ArgumentError] if any of the arguments is missing or not the correct type
       def self.decrypt(ciphertext, key, iv)
-        dec = [ciphertext, key, iv].map {|i| Base64.decode64(i) }
+        raise ArgumentError unless ciphertext.instance_of?(String) &&
+                                   key.instance_of?(String) &&
+                                   iv.instance_of?(String)
 
         decipher = OpenSSL::Cipher.new(CIPHER)
         decipher.decrypt
-        decipher.key = dec[1]
-        decipher.iv = dec[2]
+        decipher.key = key
+        decipher.iv = iv
 
-        plain = decipher.update(dec[0]) + decipher.final
-        plain
+        decipher.update(Base64.decode64(ciphertext)) + decipher.final
       end
     end
   end

lib/diaspora_federation/salmon/encrypted_slap.rb

@@ -154,9 +154,9 @@ module DiasporaFederation
       # @return [Nokogiri::XML::Element] header xml document
       def self.decrypt_header(data, pkey)
         cipher_header = JSON.parse(Base64.decode64(data))
-        header_key = JSON.parse(pkey.private_decrypt(Base64.decode64(cipher_header["aes_key"])))
+        key = JSON.parse(pkey.private_decrypt(Base64.decode64(cipher_header["aes_key"])))
 
-        xml = AES.decrypt(cipher_header["ciphertext"], header_key["key"], header_key["iv"])
+        xml = AES.decrypt(cipher_header["ciphertext"], Base64.decode64(key["key"]), Base64.decode64(key["iv"]))
         Nokogiri::XML::Document.parse(xml).root
       end
       private_class_method :decrypt_header
@@ -169,12 +169,13 @@ module DiasporaFederation
       # @param parent_node [Nokogiri::XML::Element] parent element for insering in XML document
       def self.encrypted_header(author_id, envelope_key, pubkey, parent_node)
         data = header_xml(author_id, envelope_key)
-        encryption_data = AES.encrypt(data)
+        key = AES.generate_key_and_iv
+        ciphertext = AES.encrypt(data, key[:key], key[:iv])
 
-        json_key = JSON.generate(key: encryption_data[:key], iv: encryption_data[:iv])
+        json_key = JSON.generate(key: Base64.strict_encode64(key[:key]), iv: Base64.strict_encode64(key[:iv]))
         encrypted_key = Base64.strict_encode64(pubkey.public_encrypt(json_key))
 
-        json_header = JSON.generate(aes_key: encrypted_key, ciphertext: encryption_data[:ciphertext])
+        json_header = JSON.generate(aes_key: encrypted_key, ciphertext: ciphertext)
 
         header = Nokogiri::XML::Element.new("encrypted_header", parent_node.document)
         header.content = Base64.strict_encode64(json_header)

lib/diaspora_federation/salmon/magic_envelope.rb

@@ -88,10 +88,9 @@ module DiasporaFederation
       #
       # @return [Hash] AES key and iv. E.g.: { key: "...", iv: "..." }
       def encrypt!
-        encryption_data = AES.encrypt(@payload)
-        @payload = encryption_data[:ciphertext]
-
-        {key: encryption_data[:key], iv: encryption_data[:iv]}
+        key = AES.generate_key_and_iv
+        @payload = AES.encrypt(@payload, key[:key], key[:iv])
+        strict_base64_encode(key)
       end
 
       # Extracts the entity encoded in the magic envelope data, if the signature
@@ -129,7 +128,7 @@ module DiasporaFederation
 
         data = Base64.urlsafe_decode64(magic_env.at_xpath("me:data").content)
         unless cipher_params.nil?
-          data = AES.decrypt(data, cipher_params[:key], cipher_params[:iv])
+          data = AES.decrypt(data, Base64.decode64(cipher_params[:key]), Base64.decode64(cipher_params[:iv]))
         end
 
         XmlPayload.unpack(Nokogiri::XML::Document.parse(data).root)
@@ -177,6 +176,12 @@ module DiasporaFederation
       def self.sig_subject(data_arr)
         data_arr.map {|i| Base64.urlsafe_encode64(i) }.join(".")
       end
+
+      # @param [Hash] hash { key: "...", iv: "..." }
+      # @return [Hash] encoded hash: { key: "...", iv: "..." }
+      def strict_base64_encode(hash)
+        Hash[hash.map {|k, v| [k, Base64.strict_encode64(v)] }]
+      end
     end
   end
 end

spec/lib/diaspora_federation/salmon/aes_spec.rb

@@ -0,0 +1,59 @@
+module DiasporaFederation
+  describe Salmon::AES do
+    let(:data) { "test data string" }
+
+    describe ".generate_key_and_iv" do
+      it "generates a random key and iv" do
+        key_and_iv = Salmon::AES.generate_key_and_iv
+
+        expect(key_and_iv[:key]).not_to be_empty
+        expect(key_and_iv[:iv]).not_to be_empty
+      end
+
+      it "generates a different key and iv every time" do
+        key_and_iv = Salmon::AES.generate_key_and_iv
+        key_and_iv_2 = Salmon::AES.generate_key_and_iv
+
+        expect(key_and_iv[:key]).not_to eq(key_and_iv_2[:key])
+        expect(key_and_iv[:iv]).not_to eq(key_and_iv_2[:iv])
+      end
+    end
+
+    describe ".encrypt" do
+      let(:key_and_iv) { Salmon::AES.generate_key_and_iv }
+
+      it "encrypts the data" do
+        ciphertext = Salmon::AES.encrypt(data, key_and_iv[:key], key_and_iv[:iv])
+
+        expect(Base64.decode64(ciphertext)).not_to eq(data)
+      end
+
+      it "raises an error when the data is missing or the wrong type" do
+        [nil, 1234, true, :symbol].each do |val|
+          expect {
+            Salmon::AES.encrypt(val, key_and_iv[:key], key_and_iv[:iv])
+          }.to raise_error ArgumentError
+        end
+      end
+    end
+
+    describe ".decrypt" do
+      it "decrypts what it has encrypted" do
+        key = Salmon::AES.generate_key_and_iv
+        ciphertext = Salmon::AES.encrypt(data, key[:key], key[:iv])
+
+        decrypted_data = Salmon::AES.decrypt(ciphertext, key[:key], key[:iv])
+
+        expect(decrypted_data).to eq(data)
+      end
+
+      it "raises an error when the params are missing or the wrong type" do
+        [nil, 1234, true, :symbol].each do |val|
+          expect {
+            Salmon::AES.decrypt(val, val, val)
+          }.to raise_error ArgumentError
+        end
+      end
+    end
+  end
+end

spec/lib/diaspora_federation/salmon/encrypted_slap_spec.rb

@@ -49,7 +49,7 @@ module DiasporaFederation
           expect(json_header).to include("aes_key", "ciphertext")
         end
 
-        it "encrypted the public_key encrypted header correctly" do
+        it "encrypts the public_key encrypted header correctly" do
           key = {}
           expect {
             key = JSON.parse(okey.private_decrypt(Base64.decode64(cipher_header["aes_key"])))
@@ -57,12 +57,12 @@ module DiasporaFederation
           expect(key).to include("key", "iv")
         end
 
-        it "encrypted the aes encrypted header correctly" do
+        it "encrypts the aes encrypted header correctly" do
           header = ""
           expect {
             header = Salmon::AES.decrypt(cipher_header["ciphertext"],
-                                         header_key["key"],
-                                         header_key["iv"])
+                                         Base64.decode64(header_key["key"]),
+                                         Base64.decode64(header_key["iv"]))
           }.not_to raise_error
           header_doc = Nokogiri::XML::Document.parse(header)
           expect(header_doc.root.name).to eq("decrypted_header")

spec/support/shared_slap_specs.rb

@@ -5,7 +5,7 @@ shared_examples "a Slap instance" do
 
   context "#entity" do
     it "requires the pubkey for the first time (to verify the signature)" do
-      expect { subject.entity }.to raise_error
+      expect { subject.entity }.to raise_error ArgumentError
     end
 
     it "works when the pubkey is given" do