require 'dionysus'
require 'digest'

##
# Convenience methods for the Digest module.
#
#     require 'dionysus/digest'
# 
# The <tt>Digest::DEFAULT_DIGESTS</tt> are automatically registered, if they
# exist.  You can register additional digests with Digest.register_digest.
# The given class must give the digest with the <tt>digest(string)</tt>.
#
# TODO add digest detection -- by length and by proc on the digests hash
module Digest
  DEFAULT_DIGESTS = [:md5, :sha1, :sha2, :sha256, :sha384, :sha512]
  @digests = {}
  
  ##
  # Register a digest.  Raises an error if the interpreted class doesn't exist.
  # It will interpret the klass as <tt>Digest::SYM</tt> if it's <tt>nil</tt>, 
  # and it will run the <tt>digest</tt> method on the klass to determine the 
  # digests bit length if bits is <tt>nil</tt>.
  # 
  # This will register <tt>:my_digest</tt> and automatically determine the bit 
  # length by executing the class's <tt>digest</tt> method on the string 
  # <tt>'1'</tt>:
  #
  #     Digest.register_digest!( :my_digest, :klass => MyDigestClass )
  #
  # Options:
  # [klass]       The digest class (also can be an arbitrary object).  Default:
  #               <tt>Digest::#{sym.to_s.upcase}</tt>
  # [bit_length]  The bit length of the digest.  Default: calculated by 
  #               running the digest on the string <tt>'1'</tt>.
  # [method]      The calculation method for the digest.  Default: 
  #               <tt>:digest</tt>
  def self.register_digest!( sym, options = {} )
    options = options.with_indifferent_access
    options[:method]      ||= :digest
    options[:klass]       ||= "Digest::#{sym.to_s.upcase}".constantize
    options[:bit_length]  ||= options[:klass].send(options[:method], '1').length * 8
    @digests[sym.to_sym] = options
  end
  
  ##
  # Register a digest.  Returns nil if an error occurs.
  def self.register_digest( sym, options = {} )
    self.register_digest!(sym, options)
  rescue LoadError
    nil
  end
  
  ##
  # The hash of registered digests.
  def self.digests
    @digests
  end
  
  ##
  # The available digests.
  def self.available_digests
    self.digests.keys
  end
  
  ##
  # The lengths of the registered digests in the given encoding.
  def self.digest_lengths( encoding = :binary )
    if encoding.is_a?(Symbol)
      bits_per_char = case encoding
        when :binary then 8
        when :base64 then 6
        when :hex, :hexidecimal then 4
        when :bit then 1
        else raise ArgumentError, "Invalid encoding"
      end
    elsif encoding.is_a?(Integer) and encoding > 0
      bits_per_char = encoding
    else
      raise ArgumentError, "Invalid encoding"
    end
    
    Hash[ self.digests.collect { |dig, info| [dig, info[:bit_length] / bits_per_char] } ]
  end
  
  ##
  # Calculate the given digest of the given string.
  #
  # Examples:
  #
  #     Digest.digest(:sha512, 'foobar') #=> binary digest
  #     Digest.digest(Digest::SHA512, 'foobar') #=> binary digest
  def self.digest( sym_or_klass, str )
    if sym_or_klass.is_a?(Class)
      sym_or_klass
    else
      Digest.const_get(sym_or_klass.to_s.upcase)
    end.digest(str)
  end
end

# Register some default digests
Digest::DEFAULT_DIGESTS.each do |dig|
  Digest.register_digest(dig)
end