require "base64"
require "openssl"

Puppet::Functions.create_function(:generate_password) do
  dispatch :generate_password do
    param 'Integer', :size
    param 'String',  :seed_file
    param 'String',  :password_key
    optional_param 'String', :method
    optional_param 'Boolean', :encode
    return_type 'String'
  end

  def generate_password(size, seed_file, password_key, method = nil, encode = false)
    key = get_key(seed_file, password_key)
    case method
    when nil
      pass = generate_string(size, key)
    when "curve25519"
      pass = generate_string(32, key, binary = true)
      pass[0] = (pass[0].ord & 248).chr
      pass[31] = ((pass[31].ord & 127) | 64).chr
    else
      raise "Unknown method"
    end

    if encode
      Base64.strict_encode64(pass).strip
    else
      pass
    end
  end

  def generate_string(size, key, binary = false)
    if binary
      set = (0 .. 255).map { |i| i.chr }
    else
      set = ('a' .. 'z').to_a + ('A' .. 'Z').to_a + ('0' .. '9').to_a
    end

    size.times.collect do |i|
      set[OpenSSL::HMAC.hexdigest(OpenSSL::Digest.new('sha256'), key, i.to_s).to_i(16) % set.size]
    end.join
  end

  def get_key(seed_file, password_key)
    "#{File.open(seed_file).read}:#{password_key}"
  end
end