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 }