#
# Copyright (C) 2004 Mauricio Julio Fernández Pradier
# See LICENSE.txt for additional licensing information.
#

require 'yaml'
require 'yaml/syck'
require 'fileutils'
require 'zlib'
require 'digest/md5'
require 'fileutils'
require 'find'
require 'stringio'

require 'rubygems/specification'
require 'rubygems/security'

module Gem

# Wrapper for FileUtils meant to provide logging and additional operations if
# needed.
class FileOperations
    extend FileUtils
    class << self
            # additional methods not implemented in FileUtils
    end
    def initialize(logger = nil)
        @logger = logger
    end

    def method_missing(meth, *args, &block)
        case
        when FileUtils.respond_to?(meth)
            @logger.log "#{meth}: #{args}" if @logger
            FileUtils.send meth, *args, &block
        when FileOperations.respond_to?(meth)
            @logger.log "#{meth}: #{args}" if @logger
            FileOperations.send meth, *args, &block
        else
            super
        end
    end
end


module Package

class NonSeekableIO < StandardError; end
class ArgumentError < ::ArgumentError; end
class ClosedIO < StandardError; end
class BadCheckSum < StandardError; end
class TooLongFileName < StandardError; end

module FSyncDir
    private
    def fsync_dir(dirname)
            # make sure this hits the disc
        begin
            dir = open(dirname, "r")
            dir.fsync
        rescue # ignore IOError if it's an unpatched (old) Ruby
        ensure
            dir.close if dir rescue nil 
        end
    end
end

class TarHeader
    FIELDS = [:name, :mode, :uid, :gid, :size, :mtime, :checksum, :typeflag,
            :linkname, :magic, :version, :uname, :gname, :devmajor, 
            :devminor, :prefix]
    FIELDS.each {|x| attr_reader x}

    def self.new_from_stream(stream)
        data = stream.read(512)
        fields = data.unpack( "A100" + # record name
                             "A8A8A8" +        # mode, uid, gid
                             "A12A12" +        # size, mtime
                             "A8A" +           # checksum, typeflag
                             "A100" +          # linkname
                             "A6A2" +          # magic, version
                             "A32" +           # uname
                             "A32" +           # gname
                             "A8A8" +          # devmajor, devminor
                             "A155"            # prefix
                            )
        name = fields.shift
        mode = fields.shift.oct
        uid = fields.shift.oct
        gid = fields.shift.oct
        size = fields.shift.oct
        mtime = fields.shift.oct
        checksum = fields.shift.oct
        typeflag = fields.shift
        linkname = fields.shift
        magic = fields.shift
        version = fields.shift.oct
        uname = fields.shift
        gname = fields.shift
        devmajor = fields.shift.oct
        devminor = fields.shift.oct
        prefix = fields.shift

        empty = (data == "\0" * 512)
        
        new(:name=>name, :mode=>mode, :uid=>uid, :gid=>gid, :size=>size, 
            :mtime=>mtime, :checksum=>checksum, :typeflag=>typeflag, :magic=>magic,
            :version=>version, :uname=>uname, :gname=>gname, :devmajor=>devmajor,
            :devminor=>devminor, :prefix=>prefix, :empty => empty )
    end
    
    def initialize(vals)
        unless vals[:name] && vals[:size] && vals[:prefix] && vals[:mode]
            raise Package::ArgumentError
        end
        vals[:uid] ||= 0
        vals[:gid] ||= 0
        vals[:mtime] ||= 0
        vals[:checksum] ||= ""
        vals[:typeflag] ||= "0"
        vals[:magic] ||= "ustar"
        vals[:version] ||= "00"
        vals[:uname] ||= "wheel"
        vals[:gname] ||= "wheel"
        vals[:devmajor] ||= 0
        vals[:devminor] ||= 0
        FIELDS.each {|x| instance_variable_set "@#{x.to_s}", vals[x]}
        @empty = vals[:empty]
    end

    def empty?
        @empty
    end

    def to_s
        update_checksum
        header(checksum)
    end

    def update_checksum
        h = header(" " * 8)
        @checksum = oct(calculate_checksum(h), 6)
    end

    private
    def oct(num, len)
        "%0#{len}o" % num
    end

    def calculate_checksum(hdr)
        hdr.unpack("C*").inject{|a,b| a+b}
    end

    def header(chksum)
#           struct tarfile_entry_posix {
#             char name[100];   # ASCII + (Z unless filled)
#             char mode[8];     # 0 padded, octal, null
#             char uid[8];      # ditto
#             char gid[8];      # ditto
#             char size[12];    # 0 padded, octal, null
#             char mtime[12];   # 0 padded, octal, null
#             char checksum[8]; # 0 padded, octal, null, space
#             char typeflag[1]; # file: "0"  dir: "5" 
#             char linkname[100]; # ASCII + (Z unless filled)
#             char magic[6];      # "ustar\0"
#             char version[2];    # "00"
#             char uname[32];     # ASCIIZ
#             char gname[32];     # ASCIIZ
#             char devmajor[8];   # 0 padded, octal, null
#             char devminor[8];   # o padded, octal, null
#             char prefix[155];   # ASCII + (Z unless filled)
#           };
        arr = [name, oct(mode, 7), oct(uid, 7), oct(gid, 7), oct(size, 11),
                oct(mtime, 11), chksum, " ", typeflag, linkname, magic, version,
                uname, gname, oct(devmajor, 7), oct(devminor, 7), prefix]
        str = arr.pack("a100a8a8a8a12a12" + # name, mode, uid, gid, size, mtime
                       "a7aaa100a6a2" + # chksum, typeflag, linkname, magic, version
                       "a32a32a8a8a155") # uname, gname, devmajor, devminor, prefix
        str + "\0" * ((512 - str.size) % 512)
    end
end

class TarWriter
    class FileOverflow < StandardError; end
    class BlockNeeded < StandardError; end

    class BoundedStream
        attr_reader :limit, :written
        def initialize(io, limit)
            @io = io
            @limit = limit
            @written = 0
        end

        def write(data)
            if data.size + @written > @limit
                raise FileOverflow, 
                    "You tried to feed more data than fits in the file." 
            end
            @io.write data
            @written += data.size
            data.size
        end
    end
    class RestrictedStream
        def initialize(anIO)
            @io = anIO
        end

        def write(data)
            @io.write data
        end
    end

    def self.new(anIO)
        writer = super(anIO)
        return writer unless block_given?
        begin
            yield writer
        ensure
            writer.close
        end
        nil
    end

    def initialize(anIO)
        @io = anIO
        @closed = false
    end

    def add_file_simple(name, mode, size)
        raise BlockNeeded unless block_given?
        raise ClosedIO if @closed
        name, prefix = split_name(name)
        header = TarHeader.new(:name => name, :mode => mode, 
                               :size => size, :prefix => prefix).to_s
        @io.write header
        os = BoundedStream.new(@io, size)
        yield os
        #FIXME: what if an exception is raised in the block?
        min_padding = size - os.written 
        @io.write("\0" * min_padding)
        remainder = (512 - (size % 512)) % 512
        @io.write("\0" * remainder)
    end

    def add_file(name, mode)
        raise BlockNeeded unless block_given?
        raise ClosedIO if @closed
        raise NonSeekableIO unless @io.respond_to? :pos=
        name, prefix = split_name(name)
        init_pos = @io.pos
        @io.write "\0" * 512 # placeholder for the header
        yield RestrictedStream.new(@io)
        #FIXME: what if an exception is raised in the block?
        #FIXME: what if an exception is raised in the block?
        size = @io.pos - init_pos - 512
        remainder = (512 - (size % 512)) % 512
        @io.write("\0" * remainder)
        final_pos = @io.pos
        @io.pos = init_pos
        header = TarHeader.new(:name => name, :mode => mode, 
                               :size => size, :prefix => prefix).to_s
        @io.write header
        @io.pos = final_pos
     end

    def mkdir(name, mode)
        raise ClosedIO if @closed
        name, prefix = split_name(name)
        header = TarHeader.new(:name => name, :mode => mode, :typeflag => "5",
                               :size => 0, :prefix => prefix).to_s
        @io.write header
        nil
    end

    def flush
        raise ClosedIO if @closed
        @io.flush if @io.respond_to? :flush
    end

    def close
        #raise ClosedIO if @closed
        return if @closed
        @io.write "\0" * 1024
        @closed = true
    end

    private
    def split_name name
        raise TooLongFileName if name.size > 256 
        if name.size <= 100
            prefix = ""
        else
            parts = name.split(/\//)
            newname = parts.pop
            nxt = ""
            loop do
                nxt = parts.pop
                break if newname.size + 1 + nxt.size > 100
                newname = nxt + "/" + newname
            end
            prefix = (parts + [nxt]).join "/"
            name = newname
            raise TooLongFileName if name.size > 100 || prefix.size > 155
        end
        return name, prefix
    end
end 

class TarReader
    include Gem::Package
    class UnexpectedEOF < StandardError; end
    module InvalidEntry
        def read(len=nil); raise ClosedIO; end
        def getc; raise ClosedIO;  end
        def rewind; raise ClosedIO;  end
    end
    class Entry
        TarHeader::FIELDS.each{|x| attr_reader x}

        def initialize(header, anIO)
            @io = anIO
            @name = header.name
            @mode = header.mode
            @uid = header.uid
            @gid = header.gid
            @size = header.size
            @mtime = header.mtime
            @checksum = header.checksum
            @typeflag = header.typeflag
            @linkname = header.linkname
            @magic = header.magic
            @version = header.version
            @uname = header.uname
            @gname = header.gname
            @devmajor = header.devmajor
            @devminor = header.devminor
            @prefix = header.prefix
            @read = 0
            @orig_pos = @io.pos
        end

        def read(len = nil)
            return nil if @read >= @size
            len ||= @size - @read
            max_read = [len, @size - @read].min
            ret = @io.read(max_read)
            @read += ret.size
            ret
        end

        def getc
            return nil if @read >= @size
            ret = @io.getc
            @read += 1 if ret
            ret
        end

        def is_directory?
            @typeflag == "5"
        end

        def is_file?
            @typeflag == "0"
        end

        def eof?
            @read >= @size
        end

        def pos
            @read
        end

        def rewind
            raise NonSeekableIO unless @io.respond_to? :pos=
            @io.pos = @orig_pos
            @read = 0
        end

        alias_method :is_directory, :is_directory?
        alias_method :is_file, :is_file

        def bytes_read
            @read
        end

        def full_name
            if @prefix != ""
                File.join(@prefix, @name)
            else
                @name
            end
        end

        def close
            invalidate
        end

        private
        def invalidate
            extend InvalidEntry
        end
    end

    def self.new(anIO)
        reader = super(anIO)
        return reader unless block_given?
        begin
            yield reader
        ensure
            reader.close
        end
        nil
    end

    def initialize(anIO)
        @io = anIO
        @init_pos = anIO.pos
    end

    def each(&block)
        each_entry(&block)
    end

    # do not call this during a #each or #each_entry iteration
    def rewind
        if @init_pos == 0
            raise NonSeekableIO unless @io.respond_to? :rewind
            @io.rewind
        else
            raise NonSeekableIO unless @io.respond_to? :pos=
            @io.pos = @init_pos
        end
    end

    def each_entry
        loop do
            return if @io.eof?
            header = TarHeader.new_from_stream(@io)
            return if header.empty?
            entry = Entry.new header, @io
            size = entry.size
            yield entry
            skip = (512 - (size % 512)) % 512
            if @io.respond_to? :seek
                # avoid reading...
                @io.seek(size - entry.bytes_read, IO::SEEK_CUR)
            else
                pending = size - entry.bytes_read
                while pending > 0
                    bread = @io.read([pending, 4096].min).size
                    raise UnexpectedEOF if @io.eof?
                    pending -= bread
                end
            end
            @io.read(skip) # discard trailing zeros
            # make sure nobody can use #read, #getc or #rewind anymore
            entry.close
        end
    end

    def close
    end
end

class TarInput
    include FSyncDir
    include Enumerable
    attr_reader :metadata
    class << self; private :new end

    def initialize(io, security_policy = nil)
        @io = io
        @tarreader = TarReader.new(@io)
        has_meta = false
        data_sig, meta_sig, data_dgst, meta_dgst = nil, nil, nil, nil
        dgst_algo = security_policy ? Gem::Security::OPT[:dgst_algo] : nil

        @tarreader.each do |entry|
            case entry.full_name 
            when "metadata"
                # (GS) Changed to line below: @metadata = YAML.load(entry.read) rescue nil
                @metadata = load_gemspec(entry.read)
                has_meta = true
                break
            when "metadata.gz"
                begin
                    # if we have a security_policy, then pre-read the
                    # metadata file and calculate it's digest
                    sio = nil
                    if security_policy 
                      Gem.ensure_ssl_available
                      sio = StringIO.new(entry.read)
                      meta_dgst = dgst_algo.digest(sio.string)
                      sio.rewind
                    end

                    gzis = Zlib::GzipReader.new(sio || entry)
                    # YAML wants an instance of IO 
                    # (GS) Changed to line below: @metadata = YAML.load(gzis) rescue nil
                    @metadata = load_gemspec(gzis)
                    has_meta = true
                ensure
                    gzis.close
                end
            when 'metadata.gz.sig'
              meta_sig = entry.read
            when 'data.tar.gz.sig'
              data_sig = entry.read
            when 'data.tar.gz'
              if security_policy
                Gem.ensure_ssl_available
                data_dgst = dgst_algo.digest(entry.read)
              end
            end
        end

        if security_policy 
          Gem.ensure_ssl_available
          # map trust policy from string to actual class (or a
          # serialized YAML file, if that exists)
          if (security_policy.is_a?(String))
            if Gem::Security.constants.index(security_policy)
              # load one of the pre-defined security policies
              security_policy = Gem::Security.const_get(security_policy)
            elsif File.exists?(security_policy)
              # FIXME: this doesn't work yet
              security_policy = YAML::load(File.read(security_policy))
            else
              raise Gem::Exception, "Unknown trust policy '#{security_policy}'"
            end
          end
          
          if data_sig && data_dgst && meta_sig && meta_dgst
            # the user has a trust policy, and we have a signed gem
            # file, so use the trust policy to verify the gem signature
            
            begin
              security_policy.verify_gem(data_sig, data_dgst, @metadata.cert_chain)
            rescue Exception => e
              raise "Couldn't verify data signature: #{e}"
            end

            begin
              security_policy.verify_gem(meta_sig, meta_dgst, @metadata.cert_chain)
            rescue Exception => e
              raise "Couldn't verify metadata signature: #{e}"
            end
          elsif security_policy.only_signed
            raise Gem::Exception, "Unsigned gem"
          else
            # FIXME: should display warning here (trust policy, but
            # either unsigned or badly signed gem file)
          end
        end

        @tarreader.rewind
        @fileops = FileOperations.new
        raise RuntimeError, "No metadata found!" unless has_meta
    end

    # Attempt to YAML-load a gemspec from the given _io_ parameter.  Return nil if it fails.
    def load_gemspec(io)
      Gem::Specification.from_yaml(io)
    rescue Gem::Exception
      nil
    end

    def self.open(filename, security_policy = nil, &block)
        open_from_io(File.open(filename, "rb"), security_policy, &block)
    end

    def self.open_from_io(io, security_policy = nil,  &block)
        raise "Want a block" unless block_given?
        begin
            is = new(io, security_policy)
            yield is
        ensure
            is.close if is
        end
    end

    def each(&block)
      @tarreader.each do |entry|
        next unless entry.full_name == "data.tar.gz"
        is = zipped_stream(entry)
        begin
          TarReader.new(is) do |inner|
            inner.each(&block)
          end
        ensure
          is.close if is
        end
      end
      @tarreader.rewind
    end

    # Return an IO stream for the zipped entry.
    #
    # If zlib is earlier than 1.2.1, then read the entry into memory
    # and create a string IO object from it.  This avoids a "buffer
    # error" problem on windows when using an earlier version of zlib.
    # This problem has not been observed in versions of zlib 1.2.1 or
    # later.
    def zipped_stream(entry)
      if Zlib::ZLIB_VERSION < '1.2.1'
        zis = Zlib::GzipReader.new entry
        dis = zis.read
        is = StringIO.new(dis)
      else
        is = Zlib::GzipReader.new entry
      end
    ensure
      zis.finish if zis
    end

    def extract_entry(destdir, entry, expected_md5sum = nil)
        if entry.is_directory?
            dest = File.join(destdir, entry.full_name)
            if file_class.dir? dest
                @fileops.chmod entry.mode, dest, :verbose=>false
            else
                @fileops.mkdir_p(dest, :mode => entry.mode, :verbose=>false)
            end
            fsync_dir dest 
            fsync_dir File.join(dest, "..")
            return
        end
        # it's a file
        md5 = Digest::MD5.new if expected_md5sum
        destdir = File.join(destdir, File.dirname(entry.full_name))
        @fileops.mkdir_p(destdir, :mode => 0755, :verbose=>false)
        destfile = File.join(destdir, File.basename(entry.full_name))
        @fileops.chmod(0600, destfile, :verbose=>false) rescue nil  # Errno::ENOENT
        file_class.open(destfile, "wb", entry.mode) do |os|
            loop do 
                data = entry.read(4096)
                break unless data
                md5 << data if expected_md5sum
                os.write(data)
            end
            os.fsync
        end
        @fileops.chmod(entry.mode, destfile, :verbose=>false)
        fsync_dir File.dirname(destfile)
        fsync_dir File.join(File.dirname(destfile), "..")
        if expected_md5sum && expected_md5sum != md5.hexdigest
            raise BadCheckSum
        end
    end

    def close
        @io.close
        @tarreader.close
    end
    
    private
    
    def file_class
        File
    end
end

class TarOutput
    
    class << self; private :new end

    class << self
        
    end
    
    def initialize(io)
        @io = io
        @external = TarWriter.new @io
    end

    def external_handle
        @external
    end

    def self.open(filename, signer = nil, &block)
        io = File.open(filename, "wb")
        open_from_io(io, signer, &block)
        nil
    end

    def self.open_from_io(io, signer = nil, &block)
        outputter = new(io)
        metadata = nil
        set_meta = lambda{|x| metadata = x}
        raise "Want a block" unless block_given?
        begin
            data_sig, meta_sig = nil, nil

            outputter.external_handle.add_file("data.tar.gz", 0644) do |inner|
                begin
                    sio = signer ? StringIO.new : nil
                    os = Zlib::GzipWriter.new(sio || inner)

                    TarWriter.new(os) do |inner_tar_stream| 
                        klass = class <<inner_tar_stream; self end
                        if RUBY_VERSION >= "1.9" then
                          klass.funcall(:define_method, :metadata=, &set_meta) 
                        else
                          klass.send(:define_method, :metadata=, &set_meta) 
                        end
                        block.call inner_tar_stream
                    end
                ensure
                    os.flush
                    os.finish
                    #os.close

                    # if we have a signing key, then sign the data
                    # digest and return the signature
                    data_sig = nil
                    if signer 
                      dgst_algo = Gem::Security::OPT[:dgst_algo]
                      dig = dgst_algo.digest(sio.string)
                      data_sig = signer.sign(dig)
                      inner.write(sio.string)
                    end
                end
            end
            
            # if we have a data signature, then write it to the gem too
            if data_sig
              sig_file = 'data.tar.gz.sig'
              outputter.external_handle.add_file(sig_file, 0644) do |os|
                os.write(data_sig)
              end
            end
            
            outputter.external_handle.add_file("metadata.gz", 0644) do |os|
                begin
                    sio = signer ? StringIO.new : nil
                    gzos = Zlib::GzipWriter.new(sio || os)
                    gzos.write metadata
                ensure
                    gzos.flush
                    gzos.finish

                    # if we have a signing key, then sign the metadata
                    # digest and return the signature
                    if signer
                      dgst_algo = Gem::Security::OPT[:dgst_algo]
                      dig = dgst_algo.digest(sio.string)
                      meta_sig = signer.sign(dig)
                      os.write(sio.string)
                    end
                end
            end

            # if we have a metadata signature, then write to the gem as
            # well
            if meta_sig
              sig_file = 'metadata.gz.sig'
              outputter.external_handle.add_file(sig_file, 0644) do |os|
                os.write(meta_sig)
              end
            end
            
        ensure
            outputter.close
        end
        nil
    end
    
    def close
        @external.close
        @io.close
    end
end
end # module Package


module Package    
    #FIXME: refactor the following 2 methods
    def self.open(dest, mode = "r", signer = nil, &block)
        raise "Block needed" unless block_given?

        case mode
        when "r"
            security_policy = signer
            TarInput.open(dest, security_policy, &block)
        when "w"
            TarOutput.open(dest, signer, &block)
        else
            raise "Unknown Package open mode"
        end
    end

    def self.open_from_io(io, mode = "r", signer = nil, &block)
        raise "Block needed" unless block_given?

        case mode
        when "r"
            security_policy = signer
            TarInput.open_from_io(io, security_policy, &block)
        when "w"
            TarOutput.open_from_io(io, signer, &block)
        else
            raise "Unknown Package open mode"
        end
    end

    def self.pack(src, destname, signer = nil)
        TarOutput.open(destname, signer) do |outp|
            dir_class.chdir(src) do 
                outp.metadata = (file_class.read("RPA/metadata") rescue nil)
                find_class.find('.') do |entry|
                    case 
                    when file_class.file?(entry)
                        entry.sub!(%r{\./}, "")
                        next if entry =~ /\ARPA\//
                        stat = File.stat(entry)
                        outp.add_file_simple(entry, stat.mode, stat.size) do |os|
                            file_class.open(entry, "rb") do |f| 
                                os.write(f.read(4096)) until f.eof? 
                            end
                        end
                    when file_class.dir?(entry)
                        entry.sub!(%r{\./}, "")
                        next if entry == "RPA"
                        outp.mkdir(entry, file_class.stat(entry).mode)
                    else
                        raise "Don't know how to pack this yet!"
                    end
                end
            end
        end
    end

    class << self
        def file_class
            File
        end

        def dir_class
            Dir
        end

        def find_class
            Find
        end
    end
end
    
end

