probe/pcap/pcap_scan.go

package pcap

import (
	"fmt"
	"sort"
	"sync"

	oms "git.loafle.net/overflow/model/scan"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

func newPCapScanner(zone *oms.Zone) PCapScanner {
	ps := &pCapScan{
		zone: zone,
	}

	return ps
}

type PCapScanner interface {
	OpenARP() chan *layers.ARP
	CloseARP(ch chan *layers.ARP)

	OpenTCP(ip string) chan *layers.TCP
	CloseTCP(ip string, ch chan *layers.TCP)

	OpenUDP(ip string) chan gopacket.Packet
	CloseUDP(ip string, ch chan gopacket.Packet)

	WritePacketData(data []byte) (err error)

	start() error
	stop()

	retain()
	release() bool
}

type pCapScan struct {
	pCapHandle *pcap.Handle
	zone       *oms.Zone

	arpListenerChanMtx sync.RWMutex
	arpListenerChans   []chan *layers.ARP

	tcpListenerChanMtx sync.RWMutex
	tcpListenerChans   map[string][]chan *layers.TCP

	udpListenerChanMtx sync.RWMutex
	udpListenerChans   map[string][]chan gopacket.Packet

	refCount int
	stopChan chan struct{}
	stopWg   sync.WaitGroup
}

func (ps *pCapScan) start() error {
	if ps.stopChan != nil {
		return fmt.Errorf("PCapScanner: already running. Stop it before starting it again")
	}

	// new pcap handle
	var h *pcap.Handle
	var err error
	if h, err = pcap.OpenLive(ps.zone.Iface, 65536, true, pcap.BlockForever); nil != err {
		return err
	}
	// set filter
	// todo add tcp, udp filter
	if err = h.SetBPFFilter("arp and src net " + ps.zone.Network + " or (((tcp[tcpflags] & (tcp-syn|tcp-ack) != 0) or (tcp[tcpflags] & (tcp-rst) != 0)) and port 60000) or udp "); nil != err {
		h.Close()
		return err
	}
	ps.pCapHandle = h

	ps.arpListenerChans = make([]chan *layers.ARP, 0)
	ps.tcpListenerChans = make(map[string][]chan *layers.TCP, 0)
	ps.udpListenerChans = make(map[string][]chan gopacket.Packet, 0)

	ps.stopChan = make(chan struct{})

	ps.stopWg.Add(1)
	go handleReceive(ps)

	return nil
}

func (ps *pCapScan) stop() {
	if ps.stopChan == nil {
		panic("PCapScanner: must be started before stopping it")
	}
	close(ps.stopChan)
	ps.stopWg.Wait()
	ps.stopChan = nil
}

func (ps *pCapScan) OpenARP() chan *layers.ARP {
	ps.arpListenerChanMtx.Lock()
	defer ps.arpListenerChanMtx.Unlock()

	c := make(chan *layers.ARP, 0)
	ps.arpListenerChans = append(ps.arpListenerChans, c)
	return c
}
func (ps *pCapScan) CloseARP(ch chan *layers.ARP) {
	ps.arpListenerChanMtx.Lock()
	defer ps.arpListenerChanMtx.Unlock()

	i := sort.Search(len(ps.arpListenerChans), func(i int) bool {
		return ch == ps.arpListenerChans[i]
	})

	if -1 != i {
		close(ch)
		ps.arpListenerChans = append(ps.arpListenerChans[:i], ps.arpListenerChans[i+1:]...)
	}
}

func (ps *pCapScan) OpenTCP(ip string) chan *layers.TCP {
	ps.tcpListenerChanMtx.Lock()
	defer ps.tcpListenerChanMtx.Unlock()

	if _, ok := ps.tcpListenerChans[ip]; !ok {
		ps.tcpListenerChans[ip] = make([]chan *layers.TCP, 0)
	}

	ch := make(chan *layers.TCP, 0)
	ps.tcpListenerChans[ip] = append(ps.tcpListenerChans[ip], ch)

	return ch
}
func (ps *pCapScan) CloseTCP(ip string, ch chan *layers.TCP) {
	ps.tcpListenerChanMtx.Lock()
	defer ps.tcpListenerChanMtx.Unlock()

	if _, ok := ps.tcpListenerChans[ip]; !ok {
		return
	}

	chs := ps.tcpListenerChans[ip]
	i := sort.Search(len(chs), func(i int) bool {
		return ch == chs[i]
	})

	if -1 != i {
		close(ch)
		ps.tcpListenerChans[ip] = append(ps.tcpListenerChans[ip][:i], ps.tcpListenerChans[ip][i+1:]...)
	}

}

func (ps *pCapScan) OpenUDP(ip string) chan gopacket.Packet {
	ps.udpListenerChanMtx.Lock()
	defer ps.udpListenerChanMtx.Unlock()

	if _, ok := ps.udpListenerChans[ip]; !ok {
		ps.udpListenerChans[ip] = make([]chan gopacket.Packet, 0)
	}

	ch := make(chan gopacket.Packet, 0)
	ps.udpListenerChans[ip] = append(ps.udpListenerChans[ip], ch)

	return ch

}
func (ps *pCapScan) CloseUDP(ip string, ch chan gopacket.Packet) {
	ps.udpListenerChanMtx.Lock()
	defer ps.udpListenerChanMtx.Unlock()

	if _, ok := ps.udpListenerChans[ip]; !ok {
		return
	}

	chs := ps.udpListenerChans[ip]
	i := sort.Search(len(chs), func(i int) bool {
		return ch == chs[i]
	})

	if -1 != i {
		close(ch)
		ps.udpListenerChans[ip] = append(ps.udpListenerChans[ip][:i], ps.udpListenerChans[ip][i+1:]...)
	}
}

func (ps *pCapScan) WritePacketData(data []byte) (err error) {
	return ps.pCapHandle.WritePacketData(data)
}

func (ps *pCapScan) retain() {
	ps.refCount++
}
func (ps *pCapScan) release() bool {
	ps.refCount--
	if 0 > ps.refCount {
		ps.refCount = 0
	}

	return 0 == ps.refCount
}

func (ps *pCapScan) destroy() {
	ps.tcpListenerChanMtx.Lock()
	for k, v := range ps.tcpListenerChans {
		for _, ch := range v {
			close(ch)
		}
		v = v[:0]
		delete(ps.tcpListenerChans, k)
	}
	ps.tcpListenerChanMtx.Unlock()

	ps.udpListenerChanMtx.Lock()
	for k, v := range ps.udpListenerChans {
		for _, ch := range v {
			close(ch)
		}
		v = v[:0]
		delete(ps.udpListenerChans, k)
	}
	ps.udpListenerChanMtx.Unlock()

	ps.arpListenerChanMtx.Lock()
	for _, v := range ps.arpListenerChans {
		close(v)
	}
	ps.arpListenerChans = ps.arpListenerChans[:0]
	ps.arpListenerChanMtx.Unlock()

	ps.pCapHandle.Close()
}

func handleReceive(ps *pCapScan) {
	defer ps.stopWg.Done()

	pSrc := gopacket.NewPacketSource(ps.pCapHandle, layers.LayerTypeEthernet)
	inPacket := pSrc.Packets()

	for {
		select {
		case packet := <-inPacket:
			handlePacket(ps, packet)
		case <-ps.stopChan:
			ps.destroy()
			return
		}
	}
}