discovery/core/pcapwrapper/pcapwrapper.go

package pcap

import (
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"git.loafle.net/overflow/discovery/discovery/types"
	"sync"
)

type PcapWrapper struct {
	m        *sync.RWMutex
	HND      *pcap.Handle
	refCount int
	stop     chan bool
	arps     []chan *layers.ARP
	tcps     map[string][]chan *layers.TCP
	udps     map[string][]chan gopacket.Packet
}

func newPcapWrapper(zone *types.DiscoveryZone) (*PcapWrapper, error) {

	// new pcap handle
	handle, err := pcap.OpenLive(zone.Iface, 65536, true, pcap.BlockForever)
	if err != nil {
		return nil, err
	}
	// set filter
	// todo add tcp, udp filter
	err = handle.SetBPFFilter("arp and src net " + zone.CidrStr() + " or (((tcp[tcpflags] & (tcp-syn|tcp-ack) != 0) or (tcp[tcpflags] & (tcp-rst) != 0)) and port 60000) or udp ")
	if err != nil {
		return nil, err
	}
	// make pcap wrapper
	w := &PcapWrapper{
		HND:      handle,
		arps:     make([]chan *layers.ARP, 0),
		tcps:     make(map[string][]chan *layers.TCP, 0),
		udps:     make(map[string][]chan gopacket.Packet, 0),
		stop:     make(chan bool),
		refCount: 1,
		m:        new(sync.RWMutex),
	}

	// start recv goroutine
	go w.recvPacket()
	return w, err
}

// arp channel funcs
func (p *PcapWrapper) OpenARPChannel() chan *layers.ARP {
	c := make(chan *layers.ARP, 0)
	p.arps = append(p.arps, c)
	return c
}

func (p *PcapWrapper) CloseARPChannel(ch chan *layers.ARP) {

	var n int = -1
	for index, value := range p.arps {
		if ch == value {
			close(ch)
			n = index
			break
		}
	}
	if n != -1 {
		p.arps = append(p.arps[:n], p.arps[n+1:]...)
	}
}

// tcp channel funcs
func (p *PcapWrapper) OpenTCPChannel(host *types.DiscoveryHost) chan *layers.TCP {
	p.m.Lock()
	defer p.m.Unlock()
	_, ok := p.tcps[host.Ip_]
	if !ok {
		p.tcps[host.Ip_] = make([]chan *layers.TCP, 0)
	}

	c := make(chan *layers.TCP, 0)
	p.tcps[host.Ip_] = append(p.tcps[host.Ip_], c)

	return c
}
func (p *PcapWrapper) CloseTCPChannel(host *types.DiscoveryHost, ch chan *layers.TCP) {
	p.m.Lock()
	defer p.m.Unlock()

	_, ok := p.tcps[host.Ip_]
	if ok {
		var n int = -1
		for index, value := range p.tcps[host.Ip_] {
			if ch == value {
				close(ch)
				n = index
				break
			}
		}
		if n != -1 {
			p.tcps[host.Ip_] = append(p.tcps[host.Ip_][:n], p.tcps[host.Ip_][n+1:]...)
		}
	}
}

// udp channel funcs
func (p *PcapWrapper) OpenUDPChannel(host *types.DiscoveryHost) chan gopacket.Packet {
	p.m.Lock()
	defer p.m.Unlock()

	_, ok := p.udps[host.Ip_]
	if !ok {
		p.udps[host.Ip_] = make([]chan gopacket.Packet, 0)
	}

	c := make(chan gopacket.Packet, 0)
	p.udps[host.Ip_] = append(p.udps[host.Ip_], c)

	return c
}

func (p *PcapWrapper) CloseUDPChannel(host *types.DiscoveryHost, ch chan gopacket.Packet) {
	p.m.Lock()
	defer p.m.Unlock()

	_, ok := p.udps[host.Ip_]
	if ok {
		var n int = -1
		for index, value := range p.udps[host.Ip_] {
			if ch == value {
				close(ch)
				n = index
				break
			}
		}
		if n != -1 {
			p.udps[host.Ip_] = append(p.udps[host.Ip_][:n], p.udps[host.Ip_][n+1:]...)
		}
	}
}

func (p *PcapWrapper) incRefCount() {
	p.refCount++
}

func (p *PcapWrapper) release() bool {
	p.refCount--
	if p.refCount == 0 {
		p.destroy()
		return true
	}
	return false
}

func (p *PcapWrapper) closeChannels() {

	// close stop
	if p.stop != nil {
		close(p.stop)
		p.stop = nil
	}
	// close tcp channels
	for k, v := range p.tcps {
		for _, ch := range v {
			close(ch)
		}
		v = v[:0]
		delete(p.tcps, k)
	}
	// close udp channels
	for k, v := range p.udps {
		for _, ch := range v {
			close(ch)
		}
		v = v[:0]
		delete(p.udps, k)
	}
	// close arp channels
	for _, v := range p.arps {
		close(v)
	}
	p.arps = p.arps[:0]
}

func (p *PcapWrapper) destroy() {
	p.m.Lock()
	defer p.m.Unlock()
	p.closeChannels()
	go p.HND.Close()
}

const arptype = 0
const tcptype = 1
const udptype = 2

func (p *PcapWrapper) recvPacket() {

	src := gopacket.NewPacketSource(p.HND, layers.LayerTypeEthernet)
	in := src.Packets()

	for {
		var packet gopacket.Packet
		select {
		case <-p.stop:
			return
		case packet = <-in:
			ptype := CheckProtocol(packet)

			if ptype == arptype {
				if p.arps != nil {
					arpLayer := packet.Layer(layers.LayerTypeARP)
					arp := arpLayer.(*layers.ARP)
					for _, c := range p.arps {
						c <- arp
					}
				}
			} else if ptype == tcptype {
				ipLayer := packet.Layer(layers.LayerTypeIPv4)
				ip := ipLayer.(*layers.IPv4).SrcIP.String()
				p.m.RLock()
				chans, ok := p.tcps[ip]
				if ok {
					layer := packet.Layer(layers.LayerTypeTCP)
					tcp, _ := layer.(*layers.TCP)
					for _, c := range chans {
						c <- tcp
					}
				}
				p.m.RUnlock()
			} else if ptype == udptype {
				ipLayer := packet.Layer(layers.LayerTypeIPv4)
				if ipLayer != nil {
					ip := ipLayer.(*layers.IPv4).SrcIP.String()
					p.m.RLock()
					chans, ok := p.udps[ip]

					if ok {
						for _, c := range chans {
							c <- packet
						}
					}

					p.m.RUnlock()
				}
			}
		}
	}
}
func CheckProtocol(packet gopacket.Packet) int {
	if packet == nil {
		return -1
	}
	layer := packet.Layer(layers.LayerTypeARP)
	if layer != nil {
		return arptype
	}

	layer = packet.Layer(layers.LayerTypeTCP)
	if layer != nil {
		if _, ok := layer.(*layers.TCP); ok {
			return tcptype
		}
	}

	layer = packet.Layer(layers.LayerTypeUDP)
	if layer != nil {
		if _, ok := layer.(*layers.UDP); ok {
			return udptype
		}
	}
	return -1
}