module Terminus class Proxy class DriverBody ASYNC_BROWSERS = %w[Android] TEMPLATE = ERB.new(<<-HTML) <% if @async %> <% else %> <% end %> HTML def initialize(env, response) @env = env @response = response @body = response[2] @host = "http://#{@env['SERVER_NAME']}:#{Terminus.port}" @async = ASYNC_BROWSERS.include?(UserAgent.parse(env['HTTP_USER_AGENT']).browser) end def each(&block) script_injected = false @source = '' @body.each do |fragment| @source << fragment output = inject_script(fragment) script_injected ||= (output != fragment) block.call(output) end unless script_injected block.call(TEMPLATE.result(binding)) end end private def inject_script(fragment) fragment.gsub(/((?:^\s*)?<\/body>)/i) do TEMPLATE.result(binding) + ($1 || '') end end def page_source @source.inspect.gsub('', '') end end end end