#!/usr/bin/env ruby

$:.unshift File.dirname(__FILE__)
$:.unshift File.join(File.dirname(__FILE__), 'lib')
$:.unshift File.join(File.dirname(__FILE__), 'third_parties')

require 'optparse'
require 'gtvs'


gtvsOpts = GTVS::Options.new

opts = OptionParser.new {|opts|
    opts.banner = "Usage: resolve-hn-http [options] <trace name>"
    opts.separator ""
    opts.separator "Specific options:"

    gtvsOpts.set_options(opts)

    opts.separator ""
    opts.separator "Common options:"
    opts.on_tail("-?", "--help", "Show this message") {
        puts opts
        exit
    }
}
opts.parse!(ARGV)

if ARGV.length < 1
  $stderr.puts 'Missing trace name (try --help)'
  exit 1
end
trace = ARGV.shift

db = GTVS::Connection.get(gtvsOpts)
gtvs = GTVS::GTVS.new(db, trace, gtvsOpts)

module HTTP
  HostName = begin
    sub_domain = '[^\\x00-\\x20\\x22\\x28\\x29\\x2c\\x2e\\x3a-' +
      '\\x3c\\x3e\\x40\\x5b-\\x5d\\x7f-\\xff]+'
    domain = "#{sub_domain}(?:\\x2e#{sub_domain})*"
    pattern = "(#{domain})(?::([0-9]+))?"
  end
end

re = Regexp.new('Host: ' + HTTP::HostName + '\\x0d\\x0a', Regexp::MULTILINE)

usql = <<-SQL
UPDATE #{gtvs.hostNames} SET HostName = ?
WHERE Ip = INET_ATON(?)
SQL

ust = db.prepare(usql) || abort("Prepared statement?")

# NOTE: here you can use GtProto instead of L7Mark
sql = <<-SQL
SELECT f.FlowId, INET_NTOA(f.DstIp), hn.HostName
FROM #{gtvs.flows} AS f, #{gtvs.hostNames} AS hn
WHERE f.DstIp = hn.Ip
AND L7Mark = 'http'
SQL

db.query(sql) {|rs|
    rs.each {|r|
        httpHost = nil
        flowId, dstIp, hostName = r

        gtvs.read_flow(flowId) {|pkt|
            next unless pkt.tcp?
            next if httpHost != nil
            data = pkt.tcp_data
#p data
            md = re.match(data)
            next if md == nil
#p md
            httpHost = md[1]
        }

        next if httpHost == nil
#puts 'in db: ' + hostName
        if hostName == nil || hostName == ""
            got = true
            hostName = httpHost
        else
            hostNames = hostName.split(' ')
            if hostNames.index(httpHost) == nil
                hostNames << httpHost
            end
            newHostName = hostNames.sort.join(' ')
            if newHostName != hostName
                got = true
                hostName = newHostName
            end
        end
        if got
#puts 'modified:' + hostName
            puts "Found #{hostName} in flow #{flowId} to #{dstIp}"
            ust.execute(hostName, dstIp)
            #puts "Database updated, #{ust.affected_rows} affected rows."
        end
    }
}

db.close

