module Msf::Exploit::Git::SmartHttp
  class Response < Rex::Proto::Http::Response

    include Msf::Exploit::Git::PktLine

    attr_reader :code, :message, :uri, :type, :wants, :haves, :capabilities

    def initialize(opts = {})
      @body = opts[:body] || ''
      @code = opts[:code] || 200
      @message = opts[:message] || 'OK'
      @uri = opts[:uri] || '/'
      @type = opts[:type] || ''
      @wants = opts[:wants] || []
      @haves = opts[:haves] || []
      @capabilities = opts[:capabilities] || ''

      super(@code, @message)
      set_headers
    end

    def set_headers
      @headers['pragma'] = 'no-cache'
      @headers['Cache-Control'] = 'no-cache'

      case @type
      when 'ref-discovery'
        @headers['Content-Type'] = 'application/x-git-upload-pack-advertisement'
      when 'upload-pack'
        @headers['Content-Type'] = 'application/x-git-upload-pack-result'
      end
    end

    def refs
      refs = {}

      pkt_lines = Msf::Exploit::Git::PktLine.get_pkt_lines(@body)
      pkt_lines.each do |line|
        match = line.match(/(?<sha>[a-f\d]{44})\s(?<ref>refs\/heads\/\S+)/)
        next unless match && match['sha'] && match['ref']
        refs[match['ref']] = match['sha']
      end

      refs
    end

    def self.from_raw_response(response)
      return nil unless response

      opts = {}
      if response.body.include?('service=git-upload-pack')
        opts[:type] = 'ref-discovery'
      else
        opts[:type] = 'upload-pack'
      end

      opts[:message] = response.message
      opts[:code] = response.code
      opts[:body] = response.body

      pkt_lines = Msf::Exploit::Git::PktLine.get_pkt_lines(response.body)
      cap = pkt_lines.find { |line| line.include?('symref=HEAD') }
      opts[:capabilities] = cap || ''

      Response.new(opts)
    end
  end
end
