=begin
Copyright (c) 2008 GOTOU Yuuzou

The MIT License

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
=end

require "openssl"
require "digest/sha1"
require "digest/md5"
require "enumerator"
require "stringio"
require "find"
require "logger"

class JarSigner
  Error = Class.new(StandardError)
  InvalidFormatError = Class.new(Error)
  NoManifestFile = Class.new(Error)
  NoSignatureFile = Class.new(Error)
  NoPublicKeySignatureFile = Class.new(Error)
  NoSignature = Class.new(Error)

  Logger = ::Logger.new($stderr, ::Logger::ERROR)
  Logger.progname = File.basename($0, ".rb")
  Logger.formatter = Proc.new{|level, datetime, progname, message|
    "%s: %s\n" % [progname, message]
  }

  def logger=(logger)
    @logger = logger
  end

  def logger
    @logger || Logger
  end

  class ManifestEntry
    attr_accessor :params, :digest
  end

  class Manifest
    attr_accessor :entries, :digest, :header, :header_digest

    def initialize
      @header = nil
      @entries = Hash.new
      @digest = @header_digest = nil
    end
  end

  class Signature
    attr_accessor :header, :public_key_signature, :content
    attr_reader :entries

    def initialize
      @header = nil
      @entries = Hash.new
      @public_key_signature = nil
      @content = nil
    end
  end

  class DigestCollection
    attr_reader :md5, :sha1

    def initialize(str="")
      @digests = [
        @md5 = Digest::MD5.new,
        @sha1 = Digest::SHA1.new,
      ]
      update(str)
    end

    def update(str)
      @digests.each{|digest| digest.update(str) }
    end
  end

  module PublicKeySignatureUtilities
    module_function

    DefaultStore = OpenSSL::X509::Store.new
    DefaultStore.set_default_paths
    DefaultStore.verify_callback = Proc.new do |preverify_ok, store_ctx|
      PublicKeySignatureUtilities.print_verify_info(preverify_ok, store_ctx) ||
      PublicKeySignatureUtilities.check_purpose(preverify_ok, store_ctx) ||
      false
    end

    def check_purpose(prevefiy_ok, store_ctx)
      cert = store_ctx.current_cert
      if store_ctx.error != OpenSSL::X509::V_ERR_INVALID_PURPOSE
        return false
      elsif ext = cert.extensions.find{|ext| ext.oid == "extendedKeyUsage" }
        if ext.value.split(/,/).find{|val| val.strip == "Code Signing" }
          return true
        end
      end
      return false
    end

    def print_verify_info(preverify_ok, store_ctx)
      cert = store_ctx.current_cert
#       Logger.info do
#         msg = [
#           "",
#           "------------------------------------------",
#           "depth:    %d"      % [store_ctx.error_depth,],
#           "error:    %d: %s"  % [store_ctx.error, store_ctx.error_string,],
#           "subject:  %s"      % [certificate_identitity(cert.subject),],
#           "issuer:   %s"      % [certificate_identitity(cert.issuer),],
#           "validity: %s - %s" % [cert.not_before, cert.not_after,],
#           "------------------------------------------",
#         ]
#         msg.join("\n")
#       end
      return preverify_ok
    end

    def certificate_identitity(dn)
      dn = dn.to_a
      rdn = nil
      rdn ||= dn.find{|n| n.first=="CN"}
      rdn ||= dn.find{|n| n.first=="O"}
      rdn ||= dn.last
      return "%1$s=%2$s" % rdn
    end
  end

  #############################################

  attr_reader :path, :manifest, :signatures
  attr_accessor :compatible_mode

  def load_manifest
    File.open(find_mf_file) do |stream|
      text, header = read_entry(stream)
      assume_compatible_mode(header)
      @manifest.header = header
      @manifest.digest = DigestCollection.new(text)
      @manifest.header_digest = entry_digest(text)
      while entry = read_entry(stream)
        text, hash = *entry
        unless name = hash["Name"]
          raise InvalidFormatError,
            "'Name:' must be in manifest entry: #{hash.inspect}"
        end
        manifest_entry = ManifestEntry.new
        manifest_entry.params = hash
        manifest_entry.digest = entry_digest(text)
        @manifest.entries[name] = manifest_entry
        @manifest.digest.update(text)
      end
    end
  end

  def available_signatures
    list = []
    entries = Dir.entries("#{@path}/META-INF")
    entries.each do |entry|
      if /(.+).sf/i =~ entry
        list << $1
      end
    end
    return list.sort
  end

  def load_signature(sign_name=nil)
    sign_name ||= available_signatures.first
    sign = Signature.new
    sign.content = ""

    # load META-INF/XXXXXXX.SF
    File.open(find_sf_file(sign_name), "rb") do |stream|
      text, header = read_entry(stream)
      sign.header = header
      sign.content << text
      while entry = read_entry(stream)
        text, hash = *entry
        sign.entries[hash["Name"]] = hash
        sign.content << text
      end
    end

    # load META-INF/XXXXXXX.RSA
    File.open(find_public_key_signature_file(sign_name), "rb") do |stream|
      sign.public_key_signature = OpenSSL::PKCS7::PKCS7.new(stream.read)
    end

    @signatures[sign_name] = sign
    return true
  end

  #############################################

  def verify(sign_name=nil, store=nil)
    sign_name ||= available_signatures.first
    return verify_manifest(sign_name) &&
      verify_digest_manifest(sign_name) &&
      verify_publick_key_signature(sign_name, store)
  end

  def verify_manifest(sign_name=nil)
    unsigned = []
    result = true
    verified = Hash.new(false)
    sign_name ||= available_signatures.first

    entries = Dir.entries(@path)
    entries -= ["META-INF", ".", ".."]
    entries.collect!{|entry| File.join(@path, entry) }
    Find.find(*entries) do |path|
      next if File.directory?(path)
      digest = DigestCollection.new(File.open(path, "rb"){|io| io.read })
      name = path.gsub(@path+"/", "")
      m = verify_manifest_entry(digest, name)
      s = verify_digest_entry(sign_name, name)
      logger.debug("%s%s    %s" % [m ? "m" : "-", s ? "s" : "-", name])
      result &&= m
      unsigned << name unless s
      verified[name] = true
    end

    @manifest.entries.each do |name, entry|
      if !verified.key?(name)
        logger.info "unverified file: #{name}"
        result = false
      elsif !verified[name]
        logger.info "signature verification failed for a file: #{name}"
        result = false
      end
    end

    unless unsigned.empty?
      logger.warn "This jar contains unsigned entries which have not been integrity-checked."
      unsigned.each do |name|
        logger.warn "  unsigned: #{name}"
      end
    end

    return result
  end

  def verify_digest_manifest(sign_name)
    assert_signature_loaded(sign_name)
    sign = @signatures[sign_name]
    unless result =
      verify_digest(@manifest.digest.sha1.digest, sign.header["SHA1-Digest-Manifest"]) ||
      verify_digest(@manifest.header_digest.sha1.digest, sign.header["SHA1-Digest"])
      logger.info "digest in manifest signature not match: #{sign_name}"
    end
    return result
  end

  def verify_publick_key_signature(sign_name, store=nil)
    assert_signature_loaded(sign_name)
    sign = @signatures[sign_name]
    pkcs7 = sign.public_key_signature
    flags = OpenSSL::PKCS7::DETACHED|OpenSSL::PKCS7::BINARY
    certs = pkcs7.certificates
    store ||= PublicKeySignatureUtilities::DefaultStore
    unless pkcs7.verify(certs, store, sign.content, flags)
      logger.info "failed to verify PKCS #7: #{sign_name}"
      return false
    end
    return true
  end

  #############################################

  def sign(name, pkey, certs)
    # not yet implemented
  end

  #############################################

  def jar?
    return @compatible_mode == :jar
  end

  def xpi?
    return @compatible_mode == :xpi
  end

  #############################################

  private

  def initialize(path)
    @compatible_mode = nil
    @path = File.expand_path(path)
    @manifest = Manifest.new
    @signatures = Hash.new
  end

  def verify_manifest_entry(digest, name)
    if entry = @manifest.entries[name]
      if verify_digest(digest.sha1.digest, entry.params["SHA1-Digest"]) ||
         verify_digest(digest.md5.digest, entry.params["MD5-Digest"])
        return true
      end
    end
    return false
  end

  def verify_digest_entry(sign_name, name)
    return false unless sign = @signatures[sign_name]
    if entry = @manifest.entries[name]
      if sigentry = sign.entries[name]
        if verify_digest(entry.digest.sha1.digest, sigentry["SHA1-Digest"]) ||
           verify_digest(entry.digest.md5.digest, sigentry["MD5-Digest"])
          return true
        end
      end
    end
    return false
  end

  def assume_compatible_mode(hash)
    case hash["Created-By"]
    when /Sun Microsystems Inc/
      @compatible_mode = :jar
    when /Signtool/
      @compatible_mode = :xpi
    else
      @compatible_mode = :unknown
    end
  end

  def assert_signature_loaded(sign_name)
    unless sign = @signatures[sign_name]
      raise NoSignature, "missing signature entry: #{sign_name}"
    end
    unless sign.public_key_signature
      raise NoSignature, "missing public key cryptgraphic signature: #{sign_name}"
    end
  end

  def entry_digest(text)
    digest = DigestCollection.new
    text.each_line do |line|
      next if line.chomp.empty? && xpi?
      digest.update(line)
    end
    return digest
  end

  def find_inf_file(basenames, exception)
    basenames.each do |basename|
      basename = basename.downcase
      entries = Dir.entries("#{@path}/META-INF")
      entries.each do |entry|
        if basename == entry.downcase
          return File.expand_path(entry, "#{@path}/META-INF")
        end
      end
    end
    raise exception, "could not find: #{basenames}"
  end

  def find_mf_file
    return find_inf_file("manifest.mf", NoManifestFile)
  end

  def find_sf_file(name)
    return find_inf_file("#{name}.sf", NoSignatureFile)
  end

  def find_public_key_signature_file(name)
    return find_inf_file(["#{name}.rsa", "#{name}.dsa"], NoPublicKeySignatureFile)
  end

  def verify_digest(digest, expected)
    if expected
      return digest == decode64(expected)
    end
    return false
  end

  def encode64(str)
    return str ? [str].pack("m").delete("\r\n") : nil
  end

  def decode64(str)
    return str ?  str.unpack("m").first : nil
  end

  def read_entry(lines)
    text = ""
    hash = Hash.new
    value = nil
    lines.each_line do |line|
      text << line
      case line.chomp
      when /^$/
        break
      when /^([\w-]+):\s+(.*)$/n
        name = $1
        value = $2
        hash[name] = value
      when /^\s+(.*)$/n
        value << $1
      else
        raise InvalidFormatError, "invalid format: #{line}"
      end
    end
    return hash.empty? ? nil : [text, hash]
  end
end

if __FILE__ == $0
  js = JarSigner.new(ARGV[0] || ".")
  js.logger.level = ::Logger::DEBUG
  js.load_manifest
  js.load_signature
  puts "verify -> %p" % js.verify

  js = JarSigner.new(ARGV[0] || ".")
  js.logger.level = ::Logger::ERROR
  js.load_manifest
  js.available_signatures.each do |sign_name|
    puts "-----  %s -----" % sign_name
    js.load_signature(sign_name)
    puts "verify_manifest -> %p" % js.verify_manifest(sign_name)
    puts "verify_digest_manifest-> %p" % js.verify_digest_manifest(sign_name)
    puts "verify_publick_key_signature -> %p" % js.verify_publick_key_signature(sign_name)
  end
end