#!/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'

include GTVS::Config

gtvsOpts = GTVS::Options.new

opts = OptionParser.new {|opts|
    opts.banner = "Usage: flowsdump [options] <trace name> <cmd> [<flow specs>...]"
    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

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

class HttpReqInfo
    attr_accessor :method, :url, :responseCode
    attr_accessor :contentLength, :contentType
    attr_accessor :userAgent, :host, :server
    attr_accessor :requestTs, :responseTs
end

class TcpTrace
    attr_reader :cmd, :requests

    def initialize(flowId, srcIp, srcPort, dstIp, dstPort, ipProto)
        @flowId = flowId
        @outputdir = "/tmp"
        @cmd = "#{TCPTRACE} -n -q -xHTTP[#{dstPort}] --output_dir=#{@outputdir} -hq stdin"
        @requests = []
    end

    def handle(pd, gtvs)
        while true
            l = pd.gets
            break if l == nil
            l = l.strip
            #puts l
            if l =~ /^(OPTIONS|GET|HEAD|POST|PUT|DELETE|TRACE) ([^\s]+) (.*)$/
                req = HttpReqInfo.new
                req.method = $1
                req.url = $2
                httpVer = $3

                l = pd.readline.strip
                l =~ /Response Code:\s+([0-9]+)/
                req.responseCode = $1

                l = pd.readline.strip # Request Length
                l = pd.readline.strip # Reply Length

                l = pd.readline.strip
                l =~ /Content Length:\s+(.+)$/
                req.contentLength = $1

                l = pd.readline.strip
                l =~ /Content Type:\s+(.+)$/
                req.contentType = $1

                l = pd.readline.strip
                l =~ /User Agent:\s+(.+)$/
                req.userAgent = $1

                l = pd.readline.strip
                l =~ /Host:\s+(.+)$/
                req.host = $1

                l = pd.readline.strip
                l =~ /Server:\s+(.+)$/
                req.server = $1

                l = pd.readline.strip
                l =~ /Time request sent:\s+.*\(([0-9\.]+)\)$/
                req.requestTs = $1

                l = pd.readline.strip
                l =~ /Time reply started:\s+.*\(([0-9\.]+)\)$/
                req.responseTs = $1

                l = pd.readline.strip # Time reply ACKed
                l = pd.readline.strip # Elapsed time
                l = pd.readline.strip # Elapsed time

                @requests << req
            end
        end

        #p @requests
        insert_requests(gtvs, @flowId, @requests)
    end

    def insert_requests(gtvs, flowId, requests)
        return if requests.size == 0
        sql = <<-SQL
    INSERT INTO #{gtvs.httpTransactions}
(FlowId, Method, Url, Host, UserAgent, ResponseCode, ContentType, ContentLength, Server, RequestTs, ResponseTs)
    VALUES
        SQL
        values = []
        requests.each {|req|
            values.push <<-SQL
(#{flowId}, '#{req.method}', '#{Mysql.quote(req.url)}', '#{req.host}', '#{Mysql.quote(req.userAgent)}', #{req.responseCode}, '#{req.contentType}', #{req.contentLength}, '#{req.server}', #{req.requestTs}, #{req.responseTs})
            SQL
        }
        sql += values.join(',') + ';'
        gtvs.db.query(sql) {|rs| }
    end

    def finish()
        system("rm #{@outputdir}/trace_* #{@outputdir}/*_contents.dat")
    end
    
end

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

sql = <<SQL
SELECT DataPath FROM Traces WHERE Name = '#{trace}'
SQL
datapath = nil
db.query(sql) {|rs|
    r = rs.fetch_row
    datapath = r[0]
}
abort("Undefined datapath.") unless datapath

sql = <<SQL
SELECT FlowId, INET_NTOA(SrcIp), SrcPort, INET_NTOA(DstIp), DstPort, IpProto, ROUND(FirstPktTs, 6), ROUND(LastPktTs, 6) FROM `#{gtvs.flows}`
SQL

if ARGV.length > 0
    sql += ' WHERE ' + ARGV.shift
end

db.query(sql) {|rs|
    rs.each {|r|
        flowId, srcIp, srcPort, dstIp, dstPort, ipProto, firstPktTs, lastPktTs = r
        handler = nil
        case program
        when "tcptrace"
            handler = TcpTrace.new(flowId, srcIp, srcPort, dstIp, dstPort, ipProto)
            cmd = handler.cmd
        else
            cmd = program
        end

        filter = <<-FILTER
"(src port #{srcPort} and dst port #{dstPort}) or
(src port #{dstPort} and dst port #{srcPort})"
        FILTER
        filter.tr!("\n", '')

        fullcmd = "#{PCAP_IDX} query #{datapath} #{srcIp} #{dstIp} #{ipProto} start #{firstPktTs} end #{lastPktTs} | " +
                  "#{TCPDUMP} -r - -s 0 -n -w - #{filter} | #{cmd}"
        #puts fullcmd
        puts "Processig flow #{flowId} #{srcIp}:#{srcPort} -> #{dstIp}:#{dstPort} (#{ipProto})"
        if handler != nil
            IO.popen(fullcmd) {|pd|
                begin
                    handler.handle(pd, gtvs)
                rescue => exc
                    $stderr.puts exc
                    $stderr.puts exc.backtrace.join("\n")
                end
            }
            handler.finish
        else
            system(fullcmd)
        end
    }
}

db.close
