# frozen_string_literal: true

require "zip"

module Compression
  class SafeZipReader
    Error = Class.new(StandardError)
    EntryTooLargeError = Class.new(Error)
    MissingEntryError = Class.new(Error)
    TooManyEntriesError = Class.new(Error)
    SuspiciousEntryError = Class.new(Error)

    DEFAULT_READ_CHUNK_BYTES = 16 * 1024

    attr_reader :zip_file, :remaining_total_bytes

    def self.open(path, **kwargs)
      ::Zip::File.open(path) do |zip_file|
        reader = new(zip_file, **kwargs)
        reader.validate!

        yield(reader)
      end
    end

    def initialize(
      zip_file,
      max_entries: nil,
      max_total_bytes: nil,
      max_compression_ratio: nil,
      read_chunk_bytes: DEFAULT_READ_CHUNK_BYTES
    )
      @zip_file = zip_file
      @max_entries = max_entries
      @remaining_total_bytes = max_total_bytes
      @max_compression_ratio = max_compression_ratio
      @read_chunk_bytes = read_chunk_bytes
    end

    def validate!
      if @max_entries && entries.size > @max_entries
        raise TooManyEntriesError, "Zip archive has too many entries"
      end

      self
    end

    def entries
      zip_file.entries
    end

    def find_entry(name)
      zip_file.find_entry(name) || find_prefixed_entry(name)
    end

    def read_entry(entry_or_name, max_bytes:, required: false)
      entry = resolve_entry(entry_or_name)
      if required && (entry.nil? || entry.directory?)
        raise MissingEntryError, "Zip entry #{entry_name(entry_or_name)} is missing"
      end
      return if entry.nil? || entry.directory?

      data = +""
      stream_entry(entry, max_bytes: max_bytes) { |chunk| data << chunk }
      data
    end

    def stream_entry_to_file(entry_or_name, file, max_bytes:, required: false)
      entry = resolve_entry(entry_or_name)
      if required && (entry.nil? || entry.directory?)
        raise MissingEntryError, "Zip entry #{entry_name(entry_or_name)} is missing"
      end
      return if entry.nil? || entry.directory?

      stream_entry(entry, max_bytes: max_bytes) { |chunk| file.write(chunk) }
    end

    private

    # macOS Finder's "Compress" wraps the folder's contents in a root
    # directory and adds __MACOSX metadata entries, so lookups fall back to
    # resolving names through that wrapper directory.
    def find_prefixed_entry(name)
      root_prefix && zip_file.find_entry("#{root_prefix}#{name}")
    end

    def root_prefix
      return @root_prefix if defined?(@root_prefix)

      @root_prefix = compute_root_prefix
    end

    def compute_root_prefix
      names = entries.map(&:name).reject { |name| metadata_entry?(name) }
      first_segment, separator, _rest = names.first.to_s.partition("/")
      return if separator.empty?

      prefix = "#{first_segment}/"
      prefix if names.all? { |name| name.start_with?(prefix) }
    end

    def metadata_entry?(name)
      name.start_with?("__MACOSX/") || File.basename(name).start_with?(".")
    end

    def resolve_entry(entry_or_name)
      entry_or_name.respond_to?(:get_input_stream) ? entry_or_name : find_entry(entry_or_name)
    end

    def entry_name(entry_or_name)
      entry_or_name.respond_to?(:name) ? entry_or_name.name : entry_or_name
    end

    def stream_entry(entry, max_bytes:)
      limit = entry_limit(max_bytes)
      validate_entry_metadata!(entry, limit)

      bytes_read = 0
      entry.get_input_stream do |stream|
        while (chunk = stream.read(@read_chunk_bytes))
          bytes_read += chunk.bytesize
          raise EntryTooLargeError, "Zip entry #{entry.name} is too large" if bytes_read > limit

          yield chunk
        end
      end

      @remaining_total_bytes -= bytes_read if @remaining_total_bytes
      bytes_read
    end

    def entry_limit(max_bytes)
      limit = [max_bytes, @remaining_total_bytes].compact.min
      raise EntryTooLargeError, "Zip content is too large" if limit.nil? || limit <= 0

      limit
    end

    def validate_entry_metadata!(entry, limit)
      if entry.size && entry.size > limit
        raise EntryTooLargeError, "Zip entry #{entry.name} is too large"
      end

      return if @max_compression_ratio.nil?
      return if entry.compressed_size.nil? || entry.compressed_size <= 0
      return if entry.size.nil?

      if (entry.size.to_f / entry.compressed_size) > @max_compression_ratio
        raise SuspiciousEntryError, "Zip entry #{entry.name} has a suspicious compression ratio"
      end
    end
  end
end