rastating/wordpress-exploit-framework

# frozen_string_literal: true

require 'fileutils'

module Wpxf
  module Utility
    # A super strong body builder capable of building formatted form bodies for
    # use with the {Wpxf::Net::HttpClient} mixin.
    class BodyBuilder
      def initialize
        @fields = {}
        @temp_dir = File.join(Dir.tmpdir, "wpxf_#{object_id}")
      end

      # Add a key-value pair to the field list.
      # @param name the name of the form item.
      # @param value the value of the form item.
      # @return [Hash] the newly added form item.
      def add_field(name, value)
        @fields[name] = { type: :normal, value: value }
      end

      # Add a file to the field list.
      # @param name the name of the form item.
      # @param path the local path of the file to be uploaded.
      # @param [optional] remote_name the file name to transmit the file as.
      # @return [Hash] the newly added form item.
      def add_file(name, path, remote_name = nil)
        @fields[name] = { type: :file, path: path, remote_name: remote_name }
      end

      # Add a file to the field list that will upload a specific string as its
      # file contents, rather than reading from disk.
      # @param name the name of the form item.
      # @param value the contents of the file.
      # @param remote_name the file name to transmit the file as.
      # @return [Hash] the newly added form item.
      def add_file_from_string(name, value, remote_name)
        @fields[name] = {
          type: :mem_file,
          value: value,
          remote_name: remote_name
        }
      end

      # Generate a ZIP file and add it to the field list.
      # @param name the name of the form item.
      # @param files [Hash] a hash of file names as keys and their
      #   content as values.
      # @param remote_name [String] the file name to transmit the file as.
      def add_zip_file(name, files, remote_name)
        @fields[name] = {
          type: :mem_zip,
          value: files,
          remote_name: remote_name
        }
      end

      # Create the body string and pass it to the specified block.
      # @yieldparam body the {Wpxf::Net::HttpClient} compatible body string.
      # @return [Nil] nothing, the body must be accessed by using
      #    a block when calling the method.
      def create
        body = _prepare_fields
        yield(body)
      ensure
        _cleanup_temp_files(body)
        nil
      end

      private

      def _cleanup_temp_files(fields)
        return if fields.nil?
        fields.each { |_k, v| v.close if v.is_a?(File) }
        FileUtils.rm_rf @temp_dir.to_s
      end

      def _create_tmp_directory(name)
        directory = File.join(@temp_dir, name)
        FileUtils.mkdir_p(directory)
      end

      def _copy_file_to_temp_path(src, dest, parent_name)
        directory = _create_tmp_directory(parent_name)
        temp_path = File.join(directory, dest)
        FileUtils.cp(src, temp_path)
        temp_path
      end

      def _create_tmp_file_from_string(content, dest, parent_name)
        path = File.join(_create_tmp_directory(parent_name), dest)
        File.open(path, 'w') { |f| f.write(content) }
        path
      end

      def _generate_zip_file(zip_name, field_name, fields)
        path = File.join(_create_tmp_directory(field_name), zip_name)
        Zip::File.open(path, Zip::File::CREATE) do |zipfile|
          fields.each do |field, content|
            zipfile.get_output_stream(field) do |os|
              os.write content
            end
          end
        end
        path
      end

      def _prepare_zip_file_field(name, field)
        zip = _generate_zip_file(field[:remote_name], name, field[:value])
        File.open(zip, 'r')
      end

      def _prepare_mem_file_field(name, field)
        path = _create_tmp_file_from_string(
          field[:value],
          field[:remote_name],
          name
        )
        File.open(path, 'r')
      end

      def _prepare_file_field(name, field)
        path = field[:path]
        unless field[:remote_name].nil?
          path = _copy_file_to_temp_path(field[:path], field[:remote_name], name)
        end
        File.open(path, 'r')
      end

      def _field_prep_map
        {
          file: :_prepare_file_field,
          mem_file: :_prepare_mem_file_field,
          mem_zip: :_prepare_zip_file_field
        }
      end

      def _prepare_fields
        fields = {}
        @fields.each do |name, field|
          if field[:type] == :normal
            fields[name] = field[:value]
          else
            fields[name] = send(_field_prep_map[field[:type]], name, field)
          end
        end
        fields
      end
    end
  end
end