//go:build !windows

package tcm

import (
	"b612.me/bcap"
	"b612.me/bcap/nfq"
	"b612.me/stario"
	"b612.me/starlog"
	"context"
	"encoding/hex"
	"fmt"
	"github.com/florianl/go-nfqueue/v2"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcapgo"
	"math/rand"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

type NfCap struct {
	count uint64
	// 按连接数最后一个历史报文
	cap *bcap.Packets
	// 监控目标ip地址列表
	target []string
	// 将军，下达命令吧！
	// 监控命令列表
	targetCmd   []string
	targetAsHex bool
	//保存为pcap格式的数据文件
	saveFile string
	// 写缓存
	packetCache chan gopacket.Packet
	//保存封禁的map
	blockMap sync.Map
	//交互模式
	interactive bool
	//展示所有报文信息，包括未在追踪列表的
	showAll bool
	//展示payload报文（utf-8直接输出）
	showPayload bool
	//payload展示为hex
	showAsHex bool
	//payload允许展示的最大字符
	maxShowPayloadSize int
	//不展示任何信息
	noShowMode bool
	// 全链路丢包率，百分比
	loss float64
	// 报文延迟抵达时间，毫秒
	delay int
	// 触发封禁词后，再过N个包再封禁
	packetDelay int
	// 触发封禁词后，使用RST重置链路
	useRST bool
	// RST模式，target=目标端单向RST，reverse=对端反向RST，both=双向RST
	rstMode                 string
	fastMode                bool //自探测
	printColor              []*starlog.Color
	logCache                chan loged
	monitorPort             int
	cache                   []string
	NFQNums                 int
	ctx                     context.Context
	fn                      context.CancelFunc
	requests                chan *handler
	allowRandomAck          bool
	onlyDropblackwordPacket bool
}

type loged struct {
	str        string
	stateLevel int
	logLevel   string
}

func (n *NfCap) inactve() {
	for {
		f := strings.Fields(stario.MessageBox("", "").MustString())
		if len(f) < 2 {
			continue
		}
		switch f[0] {
		case "allow":
			for _, v := range f[1:] {
				n.blockMap.Delete(v)
				starlog.Infof("允许%s报文\n", v)
			}
		case "drop":
			for _, v := range f[1:] {
				n.blockMap.Store(v, true)
				starlog.Infof("封禁%s报文\n", v)
			}
		case "delay":
			tmp, err := strconv.Atoi(f[1])
			if err != nil {
				starlog.Errorln("输入延迟无效：%v\n", err)
				continue
			}
			n.delay = tmp
			starlog.Infof("延迟生效：%v ms\n", n.delay)
		case "loss":
			tmp, err := strconv.ParseFloat(f[1], 64)
			if err != nil {
				starlog.Errorln("输入丢包百分比无效：%v\n", err)
				continue
			}
			n.loss = tmp
			starlog.Infof("丢包百分比生效：%v ms\n", n.loss)
		}
	}
}

func (t *NfCap) doIptables() error {
	if t.monitorPort != 0 {
		if _, err := exec.Command("iptables", "-t", "raw", "-A", "PREROUTING", "-p", "tcp", "--dport", strconv.Itoa(t.monitorPort), "-j", "NFQUEUE", "--queue-num", strconv.Itoa(t.NFQNums)).CombinedOutput(); err != nil {
			return err
		}
		t.cache = append(t.cache, "-t raw -D PREROUTING -p tcp --dport "+strconv.Itoa(t.monitorPort)+" -j NFQUEUE --queue-num "+strconv.Itoa(t.NFQNums))
		if _, err := exec.Command("iptables", "-t", "raw", "-A", "OUTPUT", "-p", "tcp", "--sport", strconv.Itoa(t.monitorPort), "-j", "NFQUEUE", "--queue-num", strconv.Itoa(t.NFQNums)).CombinedOutput(); err != nil {
			return err
		}
		t.cache = append(t.cache, "-t raw -D OUTPUT -p tcp --sport "+strconv.Itoa(t.monitorPort)+" -j NFQUEUE --queue-num "+strconv.Itoa(t.NFQNums))
		if _, err := exec.Command("iptables", "-t", "raw", "-A", "PREROUTING", "-p", "tcp", "--sport", strconv.Itoa(t.monitorPort), "-j", "NFQUEUE", "--queue-num", strconv.Itoa(t.NFQNums)).CombinedOutput(); err != nil {
			return err
		}
		t.cache = append(t.cache, "-t raw -D PREROUTING -p tcp --sport "+strconv.Itoa(t.monitorPort)+" -j NFQUEUE --queue-num "+strconv.Itoa(t.NFQNums))
		if _, err := exec.Command("iptables", "-t", "raw", "-A", "OUTPUT", "-p", "tcp", "--dport", strconv.Itoa(t.monitorPort), "-j", "NFQUEUE", "--queue-num", strconv.Itoa(t.NFQNums)).CombinedOutput(); err != nil {
			return err
		}
		t.cache = append(t.cache, "-t raw -D OUTPUT -p tcp --dport "+strconv.Itoa(t.monitorPort)+" -j NFQUEUE --queue-num "+strconv.Itoa(t.NFQNums))

	}
	return nil
}

func (t *NfCap) undoIptables() {
	for _, cmd := range t.cache {
		exec.Command("iptables", strings.Fields(cmd)...).CombinedOutput()
	}
}

func (t *NfCap) Run() error {
	stopSignal := make(chan os.Signal)
	starlog.Noticef("Starting nfqueue capture\n")
	nf := nfq.NewNfQueue(t.ctx, uint16(t.NFQNums), 65535)
	nf.SetRecall(t.handleRoute)
	if t.targetAsHex {
		for k, v := range t.targetCmd {
			tmp, _ := hex.DecodeString(v)
			t.targetCmd[k] = string(tmp)
		}
	}
	go func() {
		if err := nf.Run(); err != nil {
			starlog.Errorln(err)
			stopSignal <- syscall.SIGKILL
		}
	}()
	if t.saveFile != "" {
		f, err := os.Create(t.saveFile)
		if err != nil {
			starlog.Errorln("创建写入文件失败", err)
			os.Exit(4)
		}
		defer f.Close()
		go t.pcapWriter(t.ctx, f)
	}
	if t.monitorPort != 0 {
		if _, err := exec.Command("iptables", "-V").CombinedOutput(); err != nil {
			starlog.Warningln("iptables not found, cannot auto set iptables")
			return fmt.Errorf("iptables not found")
		}
		defer t.undoIptables()
		if err := t.doIptables(); err != nil {
			starlog.Errorf("Failed to set iptables:%v\n", err)
			return err
		}
		starlog.Infof("Set iptables success\n")
	}
	signal.Notify(stopSignal, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)
	go t.handleNfResult()
	go t.logPrint(t.ctx)
	select {
	//todo need nf.Stop()
	case <-stopSignal:
		starlog.Warningf("Received stop signal\n")

	case <-t.ctx.Done():
		starlog.Infoln("TCPMonitor Task Finished")

	}
	return nil

}

func NewNfCap() *NfCap {
	var nf = new(NfCap)
	nf.packetCache = make(chan gopacket.Packet, 2048)
	nf.logCache = make(chan loged, 2048)
	nf.printColor = []*starlog.Color{
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgHiCyan),                //1=tcp_connect_1,
		starlog.NewColor(starlog.FgHiCyan),                //2=tcp_connect_2,
		starlog.NewColor(starlog.FgHiCyan),                //3=tcp_connect_3,
		starlog.NewColor(starlog.FgCyan),                  //4=tcp_bye_bye
		starlog.NewColor(starlog.FgCyan),                  //5=tcp_bye_bye
		starlog.NewColor(starlog.FgCyan),                  //6=tcp_bye_bye
		starlog.NewColor(starlog.FgCyan),                  //7=tcp_bye_bye
		starlog.NewColor(starlog.FgCyan),                  //8=tcp_bye_bye
		starlog.NewColor(starlog.FgGreen),                 //9=tcp_ok
		starlog.NewColor(starlog.BgRed, starlog.FgYellow), //10=tcp_retrans
		starlog.NewColor(starlog.FgHiMagenta),             //ece
		starlog.NewColor(starlog.FgHiMagenta),             //cwr
		starlog.NewColor(starlog.FgRed),                   //rst
		starlog.NewColor(starlog.FgHiGreen),               //keepalive
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //20=udp
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
	}
	nf.cap = bcap.NewPackets()
	nf.ctx, nf.fn = context.WithCancel(context.Background())
	nf.requests = make(chan *handler, 2048)
	return nf
}
func (t *NfCap) handleNfResult() {
	for {
		select {
		case <-t.ctx.Done():
			return
		case info := <-t.requests:
			if info == nil {
				continue
			}
			if t.allowRandomAck {
				go info.p.SetVerdict(info.id, <-info.fin)
				continue
			}
			info.p.SetVerdict(info.id, <-info.fin)
		}
	}
}

func (t *NfCap) handleRoute(id uint32, q *nfqueue.Nfqueue, p nfq.Packet) {
	if t.requests == nil {
		q.SetVerdict(id, nfqueue.NfAccept)
		return
	}
	info, err := t.cap.ParsePacket(p.Packet)
	if err != nil {
		q.SetVerdict(id, nfqueue.NfAccept)
		return
	}
	fin := make(chan int)
	t.requests <- &handler{
		id:     id,
		p:      q,
		packet: p.Packet,
		attr:   p.Attr,
		fin:    fin,
	}
	go func() {
		fin <- t.handlePacket(info, p)
	}()
}

type handler struct {
	id     uint32
	p      *nfqueue.Nfqueue
	packet gopacket.Packet
	attr   nfqueue.Attribute
	fin    chan int
}

func (n *NfCap) logPrint(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case l := <-n.logCache:
			if n.noShowMode {
				fmt.Printf("已捕获报文数量:%v个\r", n.count)
				continue
			}
			switch l.logLevel {
			case "info":
				starlog.Info(l.str)
			case "notice":
				starlog.Notice(l.str)
			case "warning":
				starlog.Warning(l.str)
			case "error":
				starlog.Error(l.str)
			case "critical":
				starlog.Critical(l.str)
			case "debug":
				starlog.Debug(l.str)
			case "payload":
				fmt.Println(l.str)
			default:
				n.printColor[l.stateLevel].Print(l.str)
			}
		}
	}
}

func (n *NfCap) ipInRange(ip string) bool {
	if len(n.target) == 0 {
		return false
	}
	for _, v := range n.target {
		if v == ip {
			return true
		}
	}
	return false
}

func (n *NfCap) strInRange(str string) bool {
	if len(n.targetCmd) == 0 {
		return false
	}
	for _, v := range n.targetCmd {
		if v == "" {
			continue
		}
		if strings.Contains(str, v) {
			return true
		}
	}
	return false
}

func (n *NfCap) handlePacket(info bcap.PacketInfo, nfp nfq.Packet) int {
	p := nfp.Packet
	n.count++
	if n.saveFile != "" {
		n.packetCache <- p
	}
	var layer gopacket.Layer
	for _, layerType := range []gopacket.LayerType{
		layers.LayerTypeTCP, layers.LayerTypeUDP, layers.LayerTypeICMPv4, layers.LayerTypeICMPv6,
		layers.LayerTypeIPv4, layers.LayerTypeIPv6, layers.LayerTypeARP,
	} {
		layer = p.Layer(layerType)
		if layer == nil {
			continue
		}
		break
	}
	if layer == nil {
		n.logCache <- loged{
			str:        "无法定义layer类型\n",
			stateLevel: 0,
			logLevel:   "error",
		}
		return nfqueue.NfAccept
	}
	var shouldDisallow bool
	if n.fastMode {
		if n.delay > 0 || n.loss > 0 {
			if n.delay > 0 && (n.ipInRange(info.SrcIP) || n.ipInRange(info.DstIP)) {
				time.Sleep(time.Millisecond * time.Duration(n.delay))
			}
			if n.loss > 0 && (n.ipInRange(info.SrcIP) || n.ipInRange(info.DstIP)) {
				if rand.Intn(10000) < int(n.loss*100) {
					shouldDisallow = true
				}
			}
		}
		if shouldDisallow {
			return nfqueue.NfDrop
		}
		return nfqueue.NfAccept
	}
	if n.loss > 0 && (n.ipInRange(info.SrcIP) || n.ipInRange(info.DstIP)) {
		if rand.Intn(10000) < int(n.loss*100) {
			shouldDisallow = true
		}
	}
	if n.delay > 0 && (n.ipInRange(info.SrcIP) || n.ipInRange(info.DstIP)) {
		time.Sleep(time.Millisecond * time.Duration(n.delay))
	}
	var dec string
	switch info.StateDescript() {
	case 0:
		dec = "未知状态"
	case 1:
		dec = "SYN tcp建立一次握手"
	case 2:
		dec = "SYN,ACK tcp建立二次握手"
	case 3:
		dec = "ACK tcp建立三次握手"
	case 4:
		dec = "FIN tcp断开一次挥手"
	case 5:
		dec = "ACK tcp断开二次挥手"
	case 6:
		dec = "FIN,ACK tcp断开二次三次挥手"
	case 7:
		dec = "FIN tcp断开三次挥手"
	case 8:
		dec = "ACK tcp断开四次挥手"
	case 9:
		dec = "tcp报文"
	case 10:
		dec = "TCP重传"
	case 11:
		dec = "TCP ece"
	case 12:
		dec = "TCP cwr"
	case 13:
		dec = "TCP RST重置"
	case 14:
		dec = "TCP Keepalive"
	}
	if n.showAll || n.ipInRange(info.SrcIP) || n.ipInRange(info.DstIP) {
		n.logCache <- loged{
			str: fmt.Sprintf("%s %v:%v -> %v:%v %s seq=%v ack=%v win=%v len=%v\n", time.Now().Format("2006-01-02 15:04:05.000000"), info.SrcIP, info.SrcPort,
				info.DstIP, info.DstPort, dec, info.TcpSeq(), info.TcpAck(), info.TcpWindow(), info.TcpPayloads()),
			stateLevel: int(info.StateDescript()),
			logLevel:   "",
		}
		if n.showPayload {
			str := string(layer.LayerPayload())
			if n.maxShowPayloadSize > 0 {
				if len(str) > n.maxShowPayloadSize {
					str = str[:n.maxShowPayloadSize]
				}
			}
			if n.showAsHex {
				str = hex.EncodeToString([]byte(str))
			}
			n.logCache <- loged{
				str:        str,
				stateLevel: int(info.StateDescript()),
				logLevel:   "payload",
			}
		}
	}
	if info.Comment() != "" || n.cap.Key(info.ReverseKey).Comment() != "" {
		tmp := info.Comment()
		if tmp == "" {
			tmp = n.cap.Key(info.ReverseKey).Comment()
		}
		pkg, _ := strconv.Atoi(tmp)
		n.logCache <- loged{
			str:        fmt.Sprintf("current delay count:%v\n", pkg-1),
			stateLevel: 0,
			logLevel:   "warning",
		}
		if pkg-1 <= 0 {
			if n.useRST {
				RealSendRST(info, n.rstMode, 3)
			}
			if n.ipInRange(info.SrcIP) {
				n.blockMap.Store(info.SrcIP, true)
			} else {
				n.blockMap.Store(info.DstIP, true)
			}
			n.cap.SetComment(info.Key, "")
			n.cap.SetComment(info.ReverseKey, "")
		} else {
			n.cap.SetComment(info.Key, strconv.Itoa(pkg-1))
			n.cap.SetComment(info.ReverseKey, strconv.Itoa(pkg-1))
		}
	}
	if len(n.targetCmd) > 0 && (n.ipInRange(info.SrcIP) || n.ipInRange(info.DstIP)) && n.strInRange(string(layer.LayerPayload())) {
		n.logCache <- loged{
			str: fmt.Sprintf("%s:%s -> %s:%s ！！Match Keyword,will block！！\n", info.SrcIP, info.SrcPort,
				info.DstIP, info.DstPort),
			stateLevel: 0,
			logLevel:   "warning",
		}
		if n.onlyDropblackwordPacket {
			if n.useRST && info.StateDescript() == 13 {
				return nfqueue.NfAccept
			}
			n.logCache <- loged{
				str:        fmt.Sprintf("Block TCP %v:%v -> %v:%v,LEN=%d\n", info.SrcIP, info.SrcPort, info.DstIP, info.DstPort, info.TcpPayloads()),
				stateLevel: 0,
				logLevel:   "warning",
			}
			return nfqueue.NfDrop
		}
		if n.packetDelay > 0 && info.Comment() == "" {
			n.cap.SetComment(info.Key, strconv.Itoa(n.packetDelay))
			n.cap.SetComment(info.ReverseKey, strconv.Itoa(n.packetDelay))
		} else {
			if n.useRST {
				RealSendRST(info, n.rstMode, 3)
			}
			if n.ipInRange(info.SrcIP) {
				n.blockMap.Store(info.SrcIP, true)
			} else {
				n.blockMap.Store(info.DstIP, true)
			}
		}
	}

	_, ok1 := n.blockMap.Load(info.DstIP)
	_, ok2 := n.blockMap.Load(info.SrcIP)
	if ok1 || ok2 {
		if n.useRST && info.StateDescript() == 13 {
			return nfqueue.NfAccept
		}
		n.logCache <- loged{
			str:        fmt.Sprintf("Block TCP %v:%v -> %v:%v,LEN=%d\n", info.SrcIP, info.SrcPort, info.DstIP, info.DstPort, info.TcpPayloads()),
			stateLevel: 0,
			logLevel:   "warning",
		}
		return nfqueue.NfDrop
	}
	if shouldDisallow {
		n.logCache <- loged{
			str:        fmt.Sprintf("Block(loss) TCP %v:%v -> %v:%v,LEN=%d\n", info.SrcIP, info.SrcPort, info.DstIP, info.DstPort, info.TcpPayloads()),
			stateLevel: 0,
			logLevel:   "warning",
		}
		return nfqueue.NfDrop
	}
	return nfqueue.NfAccept
}

func (n *NfCap) pcapWriter(stopCtx context.Context, fp *os.File) error {
	w := pcapgo.NewWriter(fp)
	err := w.WriteFileHeader(65535, layers.LinkTypeRaw)
	if err != nil {
		return err
	}
	for {
		select {
		case <-stopCtx.Done():
			return nil
		case p := <-n.packetCache:
			w.WritePacket(gopacket.CaptureInfo{
				Timestamp:     time.Now(),
				CaptureLength: len(p.Data()),
				Length:        len(p.Data()),
			}, p.Data())
		}
	}
}

func RealSendRST(info bcap.PacketInfo, target string, number int) {
	for i := 0; i < number; i++ {
		if target == "both" || target == "target" {
			SendRST(info.SrcIP, info.SrcPort, info.DstIP, info.DstPort, info.TcpSeq()+(uint32(i)*uint32(info.TcpWindow())))
		}
		if target == "both" || target == "reverse" {
			SendRST(info.DstIP, info.DstPort, info.SrcIP, info.SrcPort, info.TcpAck()+(uint32(i)*uint32(info.TcpWindow())))
		}
	}
}

func SendRST(srcIP, srcPort, dstIP, dstPort string, seq uint32) error {
	if net.ParseIP(dstIP).To4() != nil {
		return sendIPv4(srcIP, srcPort, dstIP, dstPort, seq, false)
	}
	return sendIPv6(srcIP, srcPort, dstIP, dstPort, seq, false)
}

func SendSYN(srcIP, srcPort, dstIP, dstPort string, seq uint32) error {
	if net.ParseIP(dstIP).To4() != nil {
		return sendIPv4(srcIP, srcPort, dstIP, dstPort, seq, true)
	}
	return sendIPv6(srcIP, srcPort, dstIP, dstPort, seq, true)
}

func sendIPv4(srcIP, srcPort, dstIP, dstPort string, seq uint32, isSyn bool) error {
	fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_RAW, syscall.IPPROTO_RAW)
	if err != nil {
		return err
	}
	defer syscall.Close(fd)

	dstNetIP := net.ParseIP(dstIP)
	iPv4 := layers.IPv4{
		SrcIP:    net.ParseIP(srcIP),
		DstIP:    dstNetIP,
		Version:  4,
		TTL:      64,
		Protocol: layers.IPProtocolTCP,
	}
	sPort, err := strconv.Atoi(srcPort)
	if err != nil {
		return err
	}
	dPort, err := strconv.Atoi(dstPort)
	if err != nil {
		return err
	}
	tcp := layers.TCP{
		SrcPort: layers.TCPPort(sPort),
		DstPort: layers.TCPPort(dPort),
		Seq:     seq,
		RST:     !isSyn,
		SYN:     isSyn,
	}

	if err = tcp.SetNetworkLayerForChecksum(&iPv4); err != nil {
		return err
	}
	buffer := gopacket.NewSerializeBuffer()
	options := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	if err = gopacket.SerializeLayers(buffer, options, &iPv4, &tcp); err != nil {
		return err
	}
	addr := syscall.SockaddrInet4{
		Port: dPort,
		Addr: [4]byte{dstNetIP.To4()[0], dstNetIP.To4()[1], dstNetIP.To4()[2], dstNetIP.To4()[3]},
	}
	return syscall.Sendto(fd, buffer.Bytes(), 0, &addr)
}

func sendIPv6(srcIP, srcPort, dstIP, dstPort string, seq uint32, isSyn bool) error {
	fd, err := syscall.Socket(syscall.AF_INET6, syscall.SOCK_RAW, syscall.IPPROTO_RAW)
	if err != nil {
		return err
	}
	defer syscall.Close(fd)

	dstNetIP := net.ParseIP(dstIP)
	iPv6 := layers.IPv6{
		SrcIP:      net.ParseIP(srcIP),
		DstIP:      dstNetIP,
		Version:    6,
		NextHeader: layers.IPProtocolTCP,
		HopLimit:   64,
	}
	sPort, err := strconv.Atoi(srcPort)
	if err != nil {
		return err
	}
	dPort, err := strconv.Atoi(dstPort)
	if err != nil {
		return err
	}
	tcp := layers.TCP{
		SrcPort: layers.TCPPort(sPort),
		DstPort: layers.TCPPort(dPort),
		Seq:     seq,
		RST:     !isSyn,
		SYN:     isSyn,
	}

	if err := tcp.SetNetworkLayerForChecksum(&iPv6); err != nil {
		return err
	}
	buffer := gopacket.NewSerializeBuffer()
	options := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	if err = gopacket.SerializeLayers(buffer, options, &iPv6, &tcp); err != nil {
		return err
	}
	addr := syscall.SockaddrInet6{
		Port: dPort,
		Addr: [16]byte(dstNetIP.To16()),
	}
	return syscall.Sendto(fd, buffer.Bytes(), 0, &addr)
}