#!/usr/bin/env ruby

require 'uri'
require 'fileutils'
require 'net/http'
require 'openssl'
require 'yaml'
require 'stringio'
require 'zlib'

module RequireURI
  VERSION = '0.1.0'
  FORMAT_VERSION = 20050330

  # exceptions 
  class Error < Exception; end
  class NoSigningCertError < Error; end
  class ExpiredCertError < Error; end
  class BadCertSubjectError < Error; end
  class BadCachePermsError < Error; end
  class HTTPFetchError < Error; end
  class TooManyRedirectsError < Error; end
  class InvalidDigestAlgoError < Error; end
  class InvalidSignatureError < Error; end
  class InvalidURISchemaError < Error; end
  class InvalidFormatError < Error; end

  #
  # Get default options.
  #
  def RequireURI::get_opts(path = nil)
    ret = {
      # cache dir paths
      :cache_path         => File::join(ENV['HOME'], '.ruri', 'code-%s.rb'),
      :meta_path          => File::join(ENV['HOME'], '.ruri', 'meta-%s.rr'),
      :certs_path         => File::join(ENV['HOME'], '.ruri', 'cert-%s.pem'),

      # cache options
      :cache_perms        => [0700, 0500],
      :cache_perm_mask    => 0777,
      :cache_age          => 48 * 3600, # cache for 2 days
      :cache_digest_algo  => OpenSSL::Digest::SHA1,

      # warning proc (should be set by caller)
      :warn_proc          => Kernel::method(:warn).to_proc,

      # supported fetching schemes
      :allowed_uri_schemes  => %w{http https file},

      :allowed_encrypt_algos => { 
        'RSA'  => OpenSSL::PKey::RSA, 
        'DSA' => OpenSSL::PKey::DSA,
        'DH' => OpenSSL::PKey::DH,
      },

      :allowed_digest_algos => { 
        'MD5'  => OpenSSL::Digest::MD5, 
        'SHA1' => OpenSSL::Digest::SHA1,
      },

      ############################
      # build-specific options   #
      ############################
      :build_key_path   => File::join(ENV['HOME'], '.ruri', 'ruri-key.pem'),
      :build_cert_path  => File::join(ENV['HOME'], '.ruri', 'ruri-cert.pem'),
      :build_cert_age   => 365 * 24 * 3600, # 1 year
      :build_key_algo   => OpenSSL::PKey::RSA,
      :build_key_size   => 2048,
      :build_sign_algo  => OpenSSL::Digest::SHA1,

      # default new cert extensions
      :build_cert_exts  => {
        'basicConstraints'      => 'CA:FALSE',
        'subjectKeyIdentifier'  => 'hash',
        'keyUsage'              => 'keyEncipherment,dataEncipherment,digitalSignature',
      },
    }

    # TODO: support for user config file
# 
#     path ||= ENV['RUBY_RR_PATH']
#     if path && File::exists?(path)
#       if st = File::stat(path) && st.mode == 
#         ret.update(YAML::load(File.read(path))) if path && File::exists?(path)
#       end
#     end
# 
  end

  #
  # create default cache directory
  #
  def RequireURI::create_cache(path, warn_proc, be_quiet)
    warn_proc.call("Creating cache path: #{path}")
    FileUtils::mkdir_p(path)
    File::chmod(0700, path)
  end

  #
  # fetch contents of given URI
  #
  def RequireURI::fetch_uri(uri, redirect_limit = 10)
    # check number of redirects
    raise TooManyRedirectsError, 'too many redirects' if redirect_limit == 0

    # parse URI
    uri = URI::parse(uri)

    # check URI scheme
    case uri.scheme
    when /^http/
      # fetch URI (TODO: proxy support)
      http = Net::HTTP.new(uri.host, uri.port)
      http.use_ssl if uri.scheme == 'https'
      
      resp = http.get(uri.path)

      # check HTTP response
      # TODO: conditional HTTP get
      case resp
      when Net::HTTPOK
        resp.body
      when Net::HTTPRedirection
        fetch_uri(resp['location'], redirect_limit - 1)
      else
        raise HTTPFetchError, resp.to_s
      end
    when 'file'
      File.read(uri.path)
    else
      raise InvalidURISchemaError, "invalid URI schema: '#{uri.schema}'"
    end
  end

  #
  # encrypt/decrypt data.
  #
  def RequireURI::crypt_data(data, key, bits, encrypt = true)
    size = bits / 8 - 12
    meth = (encrypt) ? :private_encrypt : :public_decrypt
    buf = ''

    # create input and output buffer streams
    i_io, o_io = StringIO.new(data), StringIO.new

    # iterate over input data and {en/de}crypt it
    o_io.write(key.send(meth, buf)) while buf = i_io.read(size)

    # set result, clear buffers
    ret = o_io.string
    i_io, o_io = nil, nil

    # return result
    ret
  end

  #
  # convert an email address to an X509-style DN
  # 
  def RequireURI::email_to_dn(email_addr)
    # TODO: this could probably be improved to munge chars, etc
    user, host = email_addr.split(/@/)
    host = host.split(/\./)
    "/CN=#{user}/" << host.map { |i| "DC=#{i}" }.join('/')
  end

  #
  # given a signer DN, get cert associated with it
  #
  def RequireURI::get_cert(signer, opts_hash = {})
    opt = get_opts.update(opts_hash)

    # build certificate path
    hash = opt[:cache_digest_algo].hexdigest(signer)
    cert_path = opt[:certs_path] % [hash]

    # check to make sure signing cert exists, try and read it
    unless File::exists?(cert_path) && cert_buf = File.read(cert_path)
      raise NoSigningCertError, "missing signing cert '#{cert_path}'"
    end

    # load certificate from buffer
    ret = OpenSSL::X509::Certificate.new(cert_buf)

    # check certificate validity
    # TODO: possible spot for OCSP support, chaining, etc
    now = Time.now
    if ret.not_before && ret.not_before > now
      raise ExpiredCertError, "cert not valid before '#{ret.not_before}'"
    elsif ret.not_after && ret.not_after < now
      raise ExpiredCertError, "cert not valid after '#{ret.not_after}'"
    elsif ret.subject.to_s != signer
      raise BadCertSubjectError, "cert subject '#{ret.subject}' doesn't match signer '#{signer}'"
    end

    # return certificate
    ret
  end


  #
  # build meta hash of data
  #
  def RequireURI::build_meta(input_data, opts_hash = {})
    opt = get_opts.update(opts_hash)
    
    # load key and certificate
    priv_path, cert_path = opt[:build_key_path], opt[:build_cert_path]
    key = opt[:build_key_algo].new(File.read(priv_path))
    cert = OpenSSL::X509::Certificate.new(File.read(cert_path))

    # check certificate validity
    # TODO: possible spot for OCSP support, chaining, etc
    now = Time.now
    if cert.not_before && cert.not_before > now
      raise ExpiredCertError, "cert not valid before '#{cert.not_before}'"
    elsif cert.not_after && cert.not_after < now
      raise ExpiredCertError, "cert not valid after '#{cert.not_after}'"
    end

    # create metadata hash
    meta = {
      'format'      => FORMAT_VERSION,
      'data'        => input_data, 

      # when was I created
      'create_time' => Time.now.to_s,

      # is the data encrypted?
      'encrypted'   => opt[:encrypted],

      # public key (but only if we're not encrypting, it seems kind of
      # stupid to pass the public key along with an encrypted data
      # stream)
      'public_key'  => opt[:encrypted] ? nil : cert.public_key.to_s,
      'algo'        => {
        'sign'          => opt[:build_sign_algo].name.split(/::/)[-1],
        'encrypt'       => opt[:build_key_algo].name.split(/::/)[-1],

        # i _really_ hate saving this, but i need it to determine the
        # block size.  i think myabe i'll fix this when (if?) i switch
        # to a symmetric cipher for encrypted data
        'encrypt_size'  => opt[:build_key_size],
      },
    }

    # compress input data
    meta['data'] = Zlib::Deflate::deflate(meta['data'])

    # encrypt data (if requested)
    if opt[:encrypted]
      meta['data'] = crypt_data(meta['data'], key, opt[:build_key_size], true)
    end

    # sign data
    dgst = opt[:build_sign_algo].new
    meta['signature'] = [key.sign(dgst, meta['data'])].pack('m')

    # packing this kind of defeats the purpose of encryption, but YAML
    # doesn't seem to preserve the string contents exactly otherwise :/
    meta['data'] = [meta['data']].pack('m')

    # return meta hash string
    Zlib::Deflate::deflate(meta.to_yaml)
  end

  #
  # simple wrapper for creating an .rr file
  #
  def RequireURI::pack(src_path, dst_path, opts_hash = {})
    opt = get_opts.update(opts_hash)

    unless File::exists?(opt[:build_key_path])
      puts "Looks like you haven't set up your build environment.",
           "Would you like to do that now? [Y/n]"
      exit(-1) unless $stdin.gets =~ /y/

      puts "Please enter your email address:"
      email = $stdin.gets
      unless email =~ /@/
        warn "Invalid email address."
        exit(-1) 
      end

      puts "Generating build environment"
      RequireURI::gen_build_env(email, opt)
      puts "Finished generating build environment."
    end

    data = build_meta(src_path ? File.read(src_path) : $stdin.read, opt)
    if dst_path
      File::open(dst_path, 'wb') { |file| file.write(data) }
    else
      $stdout.write(data)
    end

    true
  end

  #
  # simple wrapper to unpack the source of a .rr file
  #
  # Note: at the moment you can only decrypt rr files if you're the
  # creator.  This will be fixed eventually (assuming you have the
  # appropriate public key, otherwise you'll always be SOL :D).
  #
  def RequireURI::unpack(src_path, dst_path, opts_hash = {})
    opt = get_opts.update(opts_hash)

    fh = dst_path ? File::open(dst_path, 'wb') : $stdout
    meta = YAML::load(Zlib::Inflate::inflate(File.read(src_path)))
    data = meta['data'].unpack('m')[0]

    # if this is an encrypted rr, then load the appropriate cert (at the
    # moment, it always loads the build cert, this will be fixed
    # eventually), then decrypt the data
    if meta['encrypted']
      # don't attempt to validate certificate (this is deliberate)
      cert = OpenSSL::X509::Certificate.new(File::read(opt[:build_cert_path]))
      size = meta['algo']['encrypt_size']

      # decrypt data
      data = crypt_data(data, cert.public_key, size, false)
    end

    # decompress the input data and spit it out on stdout
    fh.puts(Zlib::Inflate::inflate(data))

    fh.close if dst_path
  end

  #
  # simple wrapper to add a trusted certificate (associated with the
  # given email address)
  #
  def RequireURI::add_trusted_cert(email_addr, cert_file, opts_hash = {})
    opt = get_opts.update(opts_hash)
    dn = RequireURI::email_to_dn(email_addr)

    hash = opt[:cache_digest_algo].hexdigest(dn)
    cert_path = opt[:certs_path] % [hash]

    puts "DN: \"#{dn}\"", "adding as: \"#{cert_path}\""
    FileUtils::cp(cert_file, cert_path)
  end

  # 
  # generate a build environment for the given email address
  #
  def RequireURI::gen_build_env(email_addr, opts_hash = {})
    # get options
    opt = get_opts.update(opts_hash)

    # generate certificate DN
    dn = RequireURI::email_to_dn(email_addr)
    puts "DN: \"#{dn}\""

    # generate private key
    size, algo = opt[:build_key_size], opt[:build_key_algo]
    algo_name = algo.name.split(/::/)[-1]
    puts 'Generating %d-bit %s private key.  (Note: This might take a while)...' % [size, algo_name]
    priv_key = algo.new(size)

    # save private key
    path = opt[:build_key_path]
    dir_path = File::dirname(path)
    begin 
      unless File::directory?(dir_path)
        create_cache(dir_path, opt[:warn_proc], opt[:no_warnings]) 
      end
      File::open(path, 'wb') { |file| file.chmod(0600); file.write(priv_key) }
    rescue Exception => err
      raise "Couldn't save private key to '#{path}': #{err}"
    end

    # build self-signed certificate
    # TODO: should do issuer, etc
    puts 'Generating self-signed SSL certificate...'
    cert = OpenSSL::X509::Certificate.new
    cert.version = 2
    cert.serial = 0
    cert.public_key = priv_key.public_key
    cert.not_before = Time.now
    cert.not_after = Time.now + opt[:build_cert_age]
    ef = OpenSSL::X509::ExtensionFactory.new(nil, cert)
    cert.extensions = opt[:build_cert_exts].map { |key, val|
      ef.create_extension(key, val)
    }
    cert.sign(priv_key, opt[:build_sign_algo].new)
    name = OpenSSL::X509::Name::parse(dn)
    cert.subject = name
    cert.issuer = name # NOTE: this would change if we issued certs

    # save cert
    path = opt[:build_cert_path]
    dir_path = File::dirname(path)
    begin 
      FileUtils::mkdir_p(dir_path) unless File::directory?(dir_path)
      File::open(path, 'wb') { |file| file.chmod(0600); file.write(cert) }
    rescue Exception => err
      raise "Couldn't save public certificate to '#{path}': #{err}"
    end

    # adding self as trusted cert
    # TODO: make this an option?
    add_trusted_cert(email_addr, path, opt)

    true
  end

  #
  # actual require_uri implementation
  #
  def RequireURI::require_uri(uri, signer, opts_hash = {})
    # options that can be overridden by the caller
    opt = get_opts.update(opts_hash)

    # check to see if there was a signing cert, if there is, 
    # convert it to DN notation
    raise NoSigningCertError, 'missing signing cert' unless signer
    signer = email_to_dn(signer)

    # create the code cache if it doesn't exist
    cache_dir_path = File::dirname(opt[:cache_path])
    unless File.directory?(cache_dir_path)
      create_cache(cache_dir_path, opt[:warn_proc], opt[:no_warnings]) 
    end

    # check cache directory permissions
    st = File.stat(cache_dir_path)
    unless opt[:cache_perms].any? { |perm|
      st.mode & opt[:cache_perm_mask] == perm & opt[:cache_perm_mask]
    }
      raise BadCachePermsError, 'bad cache permissions'
    end

    # build cached code/metadata paths
    hash = opt[:cache_digest_algo].hexdigest(uri)
    cache_path = File::join(opt[:cache_path] % [hash])
    meta_path = File::join(opt[:meta_path] % [hash])

    # check to see if the cache exists, and if it's valid
    cache_limit = Time.now - opt[:cache_age]
    unless File.exists?(cache_path) && File.mtime(cache_path) >= cache_limit
      # load certificate
      # (could raise exception here)
      cert = get_cert(signer, opt)

      # code isn't cached, fetch from source
      # (could raise exception here)
      content = fetch_uri(uri)
      
      # write content to the meta file
      File::open(meta_path, 'wb') { |file| file.write(content) }

      # deserialize the meta file
      # (YAML could raise an exception here)
      meta = YAML::load(Zlib::Inflate::inflate(content))

      # grab data from meta info, and un-base64 it
      data = meta['data'].unpack('m')[0]

      # verify format version 
      unless meta['format'] == FORMAT_VERSION
        raise InvalidFormatError, "invalid format: #{meta['format']}"
      end

      # check cert digest algorithm
      algo = meta['algo']['sign']
      unless opt[:allowed_digest_algos].key?(algo)
        raise InvalidDigestAlgoError, "invalid digest algorithm '#{algo}'"
      end

      # create an instance of the given digest algo
      dgst = opt[:allowed_digest_algos][algo].new

      #########################
      # VERIFY SIGNATURE HERE #
      #########################
      sig = meta['signature'].unpack('m')[0]
      unless cert.public_key.verify(dgst, sig, data)
        raise InvalidSignatureError, 'invalid signature'
      end

      # if it's encrypted, then decrypt it
      # (could raise exception here)
      if meta['encrypted']
        size = meta['algo']['encrypt_size'] # really shouldn't save this :/
        data = crypt_data(data, cert.public_key, size, false)
      end

      # decompress data
      data = Zlib::Inflate::inflate(data)

      # if we're okay, then write the data to the cache file
      File::open(cache_path, 'wb') { |file| file.write(data) }
    end

    # if we've gotten this far, then load the cache file, and return the
    # result of require()
    require cache_path
  end
end

module Kernel
  def require_uri(*args)
    RequireURI::require_uri(*args)
  end
end

if __FILE__ == $0
  test_dir = File::join(File::dirname(__FILE__), '..', 'test')
  in_path = File::join(test_dir, 'test.rb')
  rr_path = File::join(test_dir, 'test.rr')
  opt = { 
    :build_key_size   => 1024,
    :build_cert_path  => File::join(test_dir, 'dot_ruri', 'ruri-cert.pem'),
    :build_key_path   => File::join(test_dir, 'dot_ruri', 'ruri-key.pem'),
    :cache_path       => File::join(test_dir, 'dot_ruri', 'code-%s.rb'),
    :meta_path        => File::join(test_dir, 'dot_ruri', 'meta-%s.yaml'),
    :certs_path       => File::join(test_dir, 'dot_ruri', 'cert-%s.pem'),
  }
  email = 'test@example.com'

  unless File::exists?(opt[:build_cert_path])
    $stderr.puts "creating build environment for #{email}"
    RequireURI::gen_build_env(email, opt)
  end

  $stderr.puts "creating test.rr"
  RequireURI::pack(in_path, rr_path, opt)

  uri = 'file://' << rr_path
  $stderr.puts "require_uri '#{uri}', '#{email}'"

  # initialize test class
  require_uri uri, email, opt
  TestClass.new
end