file_checksum.rb in vagrant-unbundled-2.2.6.0

- old
+ new

@@ -8,17 +8,31 @@
 end
 
 class FileChecksum
   BUFFER_SIZE = 1024 * 8
 
+  # Supported file checksum
+  CHECKSUM_MAP = {
+    :md5 => Digest::MD5,
+    :sha1 => Digest::SHA1,
+    :sha256 => Digest::SHA256,
+    :sha384 => Digest::SHA384,
+    :sha512 => Digest::SHA512
+  }.freeze
+
   # Initializes an object to calculate the checksum of a file. The given
   # ``digest_klass`` should implement the ``DigestClass`` interface. Note
   # that the built-in Ruby digest classes duck type this properly:
   # Digest::MD5, Digest::SHA1, etc.
   def initialize(path, digest_klass)
-    @digest_klass = digest_klass
-    @path         = path
+    if digest_klass.is_a?(Class)
+      @digest_klass = digest_klass
+    else
+      @digest_klass = load_digest(digest_klass)
+    end
+
+    @path = path
   end
 
   # This calculates the checksum of the file and returns it as a
   # string.
   #
@@ -38,8 +52,19 @@
           break
         end
       end
     end
 
-    return digest.hexdigest
+    digest.hexdigest
+  end
+
+  private
+
+  def load_digest(type)
+    digest = CHECKSUM_MAP[type.to_s.to_sym]
+    if digest.nil?
+      raise Errors::BoxChecksumInvalidType,
+        type: type.to_s
+    end
+    digest
   end
 end