lib/lockbox/box.rb in lockbox-0.2.0 vs lib/lockbox/box.rb in lockbox-0.2.1

@@ -1,10 +1,10 @@
 require "securerandom"
 
 class Lockbox
   class Box
-    def initialize(key: nil, algorithm: nil, encryption_key: nil, decryption_key: nil)
+    def initialize(key: nil, algorithm: nil, encryption_key: nil, decryption_key: nil, padding: false)
       raise ArgumentError, "Cannot pass both key and public/private key" if key && (encryption_key || decryption_key)
 
       key = Lockbox::Utils.decode_key(key) if key
       encryption_key = Lockbox::Utils.decode_key(encryption_key) if encryption_key
       decryption_key = Lockbox::Utils.decode_key(decryption_key) if decryption_key
@@ -32,13 +32,15 @@
       else
         raise ArgumentError, "Unknown algorithm: #{algorithm}"
       end
 
       @algorithm = algorithm
+      @padding = padding == true ? 16 : padding
     end
 
     def encrypt(message, associated_data: nil)
+      message = Lockbox.pad(message, size: @padding) if @padding
       case @algorithm
       when "hybrid"
         raise ArgumentError, "No public key set" unless @encryption_box
         raise ArgumentError, "Associated data not supported with this algorithm" if associated_data
         nonce = generate_nonce(@encryption_box)
@@ -52,22 +54,25 @@
       end
       nonce + ciphertext
     end
 
     def decrypt(ciphertext, associated_data: nil)
-      case @algorithm
-      when "hybrid"
-        raise ArgumentError, "No private key set" unless @decryption_box
-        raise ArgumentError, "Associated data not supported with this algorithm" if associated_data
-        nonce, ciphertext = extract_nonce(@decryption_box, ciphertext)
-        @decryption_box.decrypt(nonce, ciphertext)
-      when "xsalsa20"
-        nonce, ciphertext = extract_nonce(@box, ciphertext)
-        @box.decrypt(nonce, ciphertext)
-      else
-        nonce, ciphertext = extract_nonce(@box, ciphertext)
-        @box.decrypt(nonce, ciphertext, associated_data)
-      end
+      message =
+        case @algorithm
+        when "hybrid"
+          raise ArgumentError, "No private key set" unless @decryption_box
+          raise ArgumentError, "Associated data not supported with this algorithm" if associated_data
+          nonce, ciphertext = extract_nonce(@decryption_box, ciphertext)
+          @decryption_box.decrypt(nonce, ciphertext)
+        when "xsalsa20"
+          nonce, ciphertext = extract_nonce(@box, ciphertext)
+          @box.decrypt(nonce, ciphertext)
+        else
+          nonce, ciphertext = extract_nonce(@box, ciphertext)
+          @box.decrypt(nonce, ciphertext, associated_data)
+        end
+      message = Lockbox.unpad(message, size: @padding) if @padding
+      message
     end
 
     # protect key for xchacha20 and hybrid
     def inspect
       to_s