func TestAssociatePGPackets(t *testing.T) {
metric := &metrics.QueryMetric{
Query: "select * from table",
QueryNetUniqueIDs: []*metrics.QueryNetUniqueID{
&metrics.QueryNetUniqueID{
SrcIP: net.IPv4(111, 111, 111, 111),
Syn: uint32(43212),
},
},
}
combinedQueryMetrics := metrics.NewQueryMetrics()
combinedQueryMetrics.Add(metric)
responses := []*ResponsePacket{
&ResponsePacket{
DstIP: net.IPv4(111, 111, 111, 111),
Ack: uint32(43212),
Size: uint64(1000),
},
&ResponsePacket{
DstIP: net.IPv4(111, 111, 111, 111),
Ack: uint32(43212),
Size: uint64(1000),
},
}
AssociatePGPackets(combinedQueryMetrics, responses)
assert.Equal(t, uint64(2), combinedQueryMetrics.List[0].TotalResponsePackets)
assert.Equal(t, uint64(2000), combinedQueryMetrics.List[0].TotalNetworkLoad)
}
func ExtractPGPackets(handle *pcap.Handle) (*metrics.QueryMetrics, []*ResponsePacket) {
combinedQueryMetrics := metrics.NewQueryMetrics()
responses := []*ResponsePacket{}
// Sorts packets into queries or responses
packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
var raw string
for {
packet, err := packetSource.NextPacket()
if err == io.EOF {
break
} else if err != nil {
fmt.Println("Error: ", err)
continue
}
ipLayer := packet.Layer(layers.LayerTypeIPv4)
if ipLayer == nil {
continue
}
ip, _ := ipLayer.(*layers.IPv4)
tcpLayer := packet.Layer(layers.LayerTypeTCP)
if tcpLayer == nil {
continue
}
tcp, _ := tcpLayer.(*layers.TCP)
// If the destination port is 5432...
if tcp.DstPort == 5432 {
// And the packet payload starts with P...
raw = fmt.Sprintf("%s", tcp.Payload)
if strings.HasPrefix(raw, "P") {
// It is a (Parse) packet that contains a Query
combinedQueryMetrics.Add(
metrics.New(
NormalizeQuery(raw),
1,
ip.SrcIP,
tcp.Seq,
),
)
}
} else if tcp.SrcPort == 5432 && tcp.ACK {
responses = append(responses, &ResponsePacket{
DstIP: ip.DstIP,
Ack: tcp.Ack,
Size: uint64(len(tcp.Payload)),
})
}
}
return combinedQueryMetrics, responses
}
func main() {
app := cli.NewApp()
app.Name = "pgnetdetective"
app.Version = "0.1"
app.Usage = "Analyze Postgres Network Traffic Captures"
app.Flags = []cli.Flag{
cli.BoolFlag{
Name: "bytes",
Usage: "Display bytes instead of as Human-Readable",
},
cli.StringFlag{
Name: "output",
Value: "text",
Usage: "Specify output file format: [text|json|csv]",
},
cli.IntFlag{
Name: "limit",
Value: 0,
Usage: "Limit output based on NetworkLoad size in kilobytes",
},
}
app.Action = func(c *cli.Context) {
if len(c.Args()) != 1 {
cli.ShowAppHelp(c)
os.Exit(0)
}
path := c.Args()[0]
// Open the .cap file
handle, err := pcap.OpenOffline(path)
if err != nil {
panic(err)
}
combinedQueryMetrics, responses := processing.ExtractPGPackets(handle)
processing.AssociatePGPackets(combinedQueryMetrics, responses)
if c.Int("limit") > 0 {
limitedQueryMetrics := metrics.NewQueryMetrics()
limit := uint64(c.Int("limit"))
for _, m := range combinedQueryMetrics.List {
if m.TotalNetworkLoad/1000 >= limit {
limitedQueryMetrics.List = append(limitedQueryMetrics.List, m)
}
}
combinedQueryMetrics = limitedQueryMetrics
}
combinedQueryMetrics.DisplayBytes = c.Bool("bytes")
sort.Sort(combinedQueryMetrics)
if c.String("output") == "json" {
out, err := json.Marshal(combinedQueryMetrics)
if err != nil {
panic(err)
}
os.Stdout.Write(out)
} else if c.String("output") == "csv" {
w := csv.NewWriter(os.Stdout)
w.WriteAll(combinedQueryMetrics.CsvString())
if err := w.Error(); err != nil {
panic(err)
}
} else {
combinedQueryMetrics.PrintText()
}
}
app.Run(os.Args)
}