rate_limit.go

rate_limit.go

import (
        "bytes"
        "code.google.com/p/gopacket"
        "code.google.com/p/gopacket/examples/util"
        "code.google.com/p/gopacket/layers"
        "code.google.com/p/gopacket/pcap"
        "github.com/nsf/termbox-go"
        "log"
        "regexp"
)
func capturePackets(tcpData chan PacketStatsMap) {
        defer util.Run()()

        var handle *pcap.Handle
        var err error

        if handle, err = pcap.OpenLive(*networkInterface, int32(*snapLength), true, pcap.BlockForever); err != nil {
                        log.Fatal("error:", err)
         }

        }

        if err := handle.SetBPFFilter("tcp src port " + *BPFfilter); err != nil {
                panic(err)
        }

        packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
        packets := packetSource.Packets()

        mcs := *PacketStatsMap()

        for {
                select {
                case packet := <-packets:

                       if app := packet.ApplicationLayer(); app != nil && len(app.Payload()) > 6 {
                                if bytes.Equal(app.Payload()[0:6], []byte{86, 65, 76, 85, 69, 32}) {


                                        if match := mcValue.FindStringSubmatch(string(app.Payload())); match != nil {

                                                if _, ok := mcs.calls[match[1]]; ok {
                                                        mcs.calls[match[1]] += 1
                                                } else {
                                                        mcs.calls[match[1]] = 1
                                                }

                                                               tcpData <- mcs                       
                                        }
                                }
                        }
                }
        }
}
func main() {

        err := termbox.Init()
        if err != nil {
                panic(err)
        }
        defer termbox.Close()

        termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
        drawUI()
        termbox.Flush()


        tcpData := make(chan PacketStatsMap)
        go capturePackets(tcpData)

        tbEvent := make(chan termbox.Event)
        go func() {
                for {
                        tbEvent <- termbox.PollEvent()
                }
        }()
loop:
        for {
                select {
                case ev := <-tbEvent:
                        termbox.Sync()
                        if ev.Ch == 113 {
                                break loop
                        }

                case tcp := <-tcpData:
                     drawUpdates(tcp)
                }
        }

        close(tbEvent)
        close(tcpData)
        termbox.Close()
}