package bcap

import (
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"net"
	"sync"
)

type Packets struct {
	sync.RWMutex
	cu map[string]PacketInfo
}

type PacketInfo struct {
	Key         string
	ReverseKey  string
	Type        string
	SrcMac      net.HardwareAddr
	DstMac      net.HardwareAddr
	SrcIP       string
	SrcPort     string
	DstIP       string
	DstPort     string
	comment     string
	packet      gopacket.Packet
	tcpSeq      uint32
	tcpAck      uint32
	tcpWindow   uint16
	tcpPayloads int
	finState    bool
	synState    bool
	isFirst     bool
	//stateDescript 0=unknown 1=tcp_connect_1 2=tcp_connect_2 3=tcp_connect_3
	// 4=tcp_disconnect_1 5=tcp_disconnect_2 6=tcp_disconnect_2+3 7=tcp_disconnect_3
	// 8=tcp_disconnect_4 9=tcp_ack_ok 10=tcp_retransmit 11=tcp_ece 12=tcp_cwr 13=tcp_rst
	// 14=tcp_keepalived 20=udp
	stateDescript uint8
}

func (p *Packets) Key(key string) PacketInfo {
	p.RLock()
	defer p.RUnlock()
	return p.cu[key]
}

func (p *Packets) SetComment(key, comment string) {
	p.Lock()
	c := p.cu[key]
	c.comment = comment
	p.cu[key] = c
	p.Unlock()
}

func (p *PacketInfo) SetComment(comment string) {
	p.comment = comment
}

func (p PacketInfo) StateDescript() uint8 {
	return p.stateDescript
}

func (p PacketInfo) TcpPayloads() int {
	return p.tcpPayloads
}

func (p PacketInfo) FinState() bool {
	return p.finState
}

func (p PacketInfo) SynState() bool {
	return p.synState
}

func (p PacketInfo) TcpWindow() uint16 {
	return p.tcpWindow
}

func (p PacketInfo) TcpAck() uint32 {
	return p.tcpAck
}

func (p PacketInfo) TcpSeq() uint32 {
	return p.tcpSeq
}

func (p PacketInfo) Packet() gopacket.Packet {
	return p.packet
}

func (p PacketInfo) Comment() string {
	return p.comment
}

func NewPackets() *Packets {
	return &Packets{
		cu: make(map[string]PacketInfo),
	}
}

func (p *Packets) ParsePacket(packet gopacket.Packet, opts ...interface{}) (PacketInfo, error) {
	var info PacketInfo
	//firstOpts is macaddr(nfqueue)
	if ew := packet.Layer(layers.LayerTypeEthernet); ew != nil {
		eth := ew.(*layers.Ethernet)
		info.SrcMac = eth.SrcMAC
		info.DstMac = eth.DstMAC
	}
	for k, v := range opts {
		switch tp := v.(type) {
		case *[]byte:
			if tp != nil && k == 0 {
				//nfqueue src mac addr
				info.SrcMac = net.HardwareAddr(*tp)
			}
		}
	}
	if nw := packet.NetworkLayer(); nw != nil {
		srcp, dstp := nw.NetworkFlow().Endpoints()
		info.SrcIP = srcp.String()
		info.DstIP = dstp.String()
	} else {
		return info, fmt.Errorf("Failed to parse packet,not a valid network info")
	}
	{
		//tcp valid
		layer := packet.Layer(layers.LayerTypeTCP)
		if layer != nil {
			if tcp, ok := layer.(*layers.TCP); ok {
				return p.parseTcp(info, packet, layer, tcp)
			}
		}
	}
	{
		layer := packet.Layer(layers.LayerTypeUDP)
		if layer != nil {
			if udp, ok := layer.(*layers.UDP); ok {
				return p.parseUdp(info, packet, layer, udp)
			}
		}
	}
	//icmp
	return info, fmt.Errorf("not support packet")
}

func (p *Packets) parseTcp(info PacketInfo, packet gopacket.Packet, layer gopacket.Layer, tcp *layers.TCP) (PacketInfo, error) {
	info.Key = fmt.Sprintf("tcp://%s:%d=%s:%d", info.SrcIP, tcp.SrcPort, info.DstIP, tcp.DstPort)
	info.ReverseKey = fmt.Sprintf("tcp://%s:%d=%s:%d", info.DstIP, tcp.DstPort, info.SrcIP, tcp.SrcPort)
	info.Type = "tcp"
	info.SrcPort = fmt.Sprintf("%d", tcp.SrcPort)
	info.DstPort = fmt.Sprintf("%d", tcp.DstPort)
	info.packet = packet
	info.tcpSeq = tcp.Seq
	info.tcpAck = tcp.Ack
	info.tcpPayloads = len(layer.LayerPayload())
	info.finState = tcp.FIN
	info.synState = tcp.SYN
	info.tcpWindow = tcp.Window
	p.RLock()
	lastPacket := p.cu[info.Key]
	p.RUnlock()
	if lastPacket.Key != info.Key {
		lastPacket = PacketInfo{
			Key:           info.Key,
			ReverseKey:    info.ReverseKey,
			Type:          "tcp",
			SrcIP:         info.SrcIP,
			SrcPort:       info.SrcPort,
			DstIP:         info.DstIP,
			DstPort:       info.DstPort,
			comment:       "",
			packet:        packet,
			tcpSeq:        tcp.Seq,
			tcpAck:        tcp.Ack,
			tcpWindow:     tcp.Window,
			tcpPayloads:   len(layer.LayerPayload()),
			finState:      tcp.FIN,
			synState:      tcp.SYN,
			isFirst:       true,
			stateDescript: 0,
		}
		p.Lock()
		p.cu[info.Key] = lastPacket
		p.Unlock()
	}
	p.RLock()
	lastReverse := p.cu[info.ReverseKey]
	p.RUnlock()
	if !lastPacket.isFirst {
		info.comment = lastPacket.comment
		if lastPacket.SrcMac != nil && len(info.SrcMac) == 0 {
			info.SrcMac = lastPacket.SrcMac
		}
		if lastPacket.SrcMac != nil && len(info.SrcMac) == 0 {
			info.SrcMac = lastPacket.SrcMac
		}
	}
	if lastReverse.SrcMac != nil && len(info.DstMac) == 0 {
		info.DstMac = lastReverse.SrcMac
	}
	{
		//state judge
		if tcp.RST {
			info.stateDescript = 13
			p.Lock()
			delete(p.cu, info.Key)
			delete(p.cu, info.ReverseKey)
			p.Unlock()
			return info, nil
		}
		if tcp.SYN && !tcp.ACK {
			info.stateDescript = 1
		} else if tcp.SYN && tcp.ACK {
			info.stateDescript = 2
		} else if tcp.ACK {

			if !tcp.FIN {
				if lastReverse.tcpSeq+1 == tcp.Ack && lastReverse.stateDescript == 2 {
					info.stateDescript = 3
				} else if tcp.CWR {
					info.stateDescript = 12
				} else if tcp.ECE {
					info.stateDescript = 11
				}
				if info.stateDescript != 0 {
					goto savereturn
				}
				if info.tcpSeq == lastReverse.tcpAck-1 && info.tcpSeq == lastPacket.tcpSeq+uint32(lastPacket.tcpPayloads)-1 {
					//keepalive
					info.stateDescript = 14
					goto savereturn
				}
				if !lastPacket.isFirst {
					if info.tcpSeq < lastPacket.tcpSeq+uint32(lastPacket.tcpPayloads) {
						//retransmit
						info.stateDescript = 10
						goto savereturn
					}
				}
				if lastReverse.finState && lastPacket.finState {
					info.stateDescript = 8
					p.Lock()
					delete(p.cu, info.Key)
					delete(p.cu, info.ReverseKey)
					p.Unlock()
					return info, nil
				}
				if lastReverse.finState && lastReverse.tcpSeq+1 == info.tcpAck {
					info.stateDescript = 5
					goto savereturn
				}
				info.stateDescript = 9
			} else {
				if !lastReverse.finState {
					info.stateDescript = 4
				} else {
					if lastReverse.finState && lastReverse.tcpSeq+1 == info.tcpAck && lastPacket.tcpAck == info.tcpAck && lastPacket.tcpSeq == info.tcpSeq {
						info.stateDescript = 7
					} else {
						info.stateDescript = 6
					}
				}
			}
		}
	}
savereturn:
	p.Lock()
	p.cu[info.Key] = info
	p.Unlock()
	return info, nil
}

func (p *Packets) parseUdp(info PacketInfo, packet gopacket.Packet, layer gopacket.Layer, udp *layers.UDP) (PacketInfo, error) {
	info.Key = fmt.Sprintf("udp://%s:%d=%s:%d", info.SrcIP, udp.SrcPort, info.DstIP, udp.DstPort)
	info.ReverseKey = fmt.Sprintf("udp://%s:%d=%s:%d", info.DstIP, udp.DstPort, info.SrcIP, udp.SrcPort)
	info.Type = "udp"
	info.SrcPort = fmt.Sprintf("%d", udp.SrcPort)
	info.DstPort = fmt.Sprintf("%d", udp.DstPort)
	info.packet = packet
	info.tcpPayloads = len(layer.LayerPayload())
	p.Lock()
	p.cu[info.Key] = info
	p.Unlock()
	return info, nil
}