package fishfinding
import (
	"net"
	"os"
	"time"
	"github.com/kataras/golog"
	"golang.org/x/net/bpf"
	"golang.org/x/net/ipv4"
)
type TCPScanner struct {
	src     net.IP
	srcPort int
	dstPort int
	input   chan string
	output  chan string
}
func NewTCPScanner(srcPort, dstPort int, input chan string, output chan string) *TCPScanner {
	localIP := GetLocalIP()
	s := &TCPScanner{
		input:   input,
		output:  output,
		src:     net.ParseIP(localIP).To4(),
		srcPort: srcPort,
		dstPort: dstPort,
	}
	return s
}
func (s *TCPScanner) Scan() {
	go s.recv()
	go s.send(s.input)
}

func (s *TCPScanner) send(input chan string) error {
	defer func() {
		time.Sleep(5 * time.Second)
		close(s.output)
		golog.Infof("send goroutine exit")
	}()
	// 创建原始套接字
	conn, err := net.ListenPacket("ip4:tcp", s.src.To4().String())
	if err != nil {
		golog.Fatal(err)
	}
	defer conn.Close()
	pconn := ipv4.NewPacketConn(conn)
	// 不接收数据
	filter := createEmptyFilter()
	if assembled, err := bpf.Assemble(filter); err == nil {
		pconn.SetBPF(assembled)
	}
	seq := uint32(os.Getpid())
	for ip := range input {
		dstIP := net.ParseIP(ip)
		if dstIP == nil {
			golog.Errorf("failed to resolve IP address %s", ip)
			continue
		}
		// 构造 TCP SYN 包
		tcpHeader := &TCPHeader{
			Source:      uint16(s.srcPort), // 源端口
			Destination: uint16(s.dstPort), // 目标端口(这里探测80端口)
			SeqNum:      seq,
			AckNum:      0,
			Flags:       0x002, // SYN
			Window:      65535,
			Checksum:    0,
			Urgent:      0,
		}
		// 计算校验和
		tcpHeader.Checksum = tcpChecksum(tcpHeader, s.src, dstIP)
		// 序列化 TCP 头
		packet := tcpHeader.Marshal()
		// 发送 TCP SYN 包
		_, err = conn.WriteTo(packet, &net.IPAddr{IP: dstIP})
		if err != nil {
			golog.Errorf("failed to send TCP packet: %v", err)
		}
	}
	return nil
}

func (s *TCPScanner) recv() error {
	defer recover()
	// 创建原始套接字
	conn, err := net.ListenIP("ip4:tcp", &net.IPAddr{IP: s.src})
	if err != nil {
		golog.Fatal(err)
	}
	defer conn.Close()
	pconn := ipv4.NewPacketConn(conn)
	if err := pconn.SetControlMessage(ipv4.FlagSrc|ipv4.FlagDst|ipv4.FlagInterface, true); err != nil {
		golog.Fatalf("set control message error: %v\n", err)
	}
	seq := uint32(os.Getpid()) + 1
	buf := make([]byte, 1024)
	for {
		n, peer, err := conn.ReadFrom(buf)
		if err != nil {
			golog.Errorf("failed to read: %v", err)
			continue
		}
		if n < tcpHeaderLength {
			continue
		}
		// 解析 TCP 头
		tcpHeader := ParseTCPHeader(buf[:n])
		if tcpHeader.Destination != uint16(s.srcPort) || tcpHeader.Source != uint16(s.dstPort) {
			continue
		}
		// golog.Info("peer: %s, flags: %d", peer.String(), tcpHeader.Flags)
		// 检查是否是 SYN+ACK, 同时检查ACK是否和发送的seq对应
		if tcpHeader.Flags == 0x012 && tcpHeader.AckNum == seq { // SYN + ACK
			s.output <- peer.String()
		}
	}
}

package fishfinding
import (
	"context"
	"fmt"
	"runtime"
	"sync"
	"time"
	"github.com/ClickHouse/clickhouse-go/v2"
	_ "github.com/ClickHouse/clickhouse-go/v2"
	"github.com/kataras/golog"
)
type ClickHouseChecker struct {
	wg   *sync.WaitGroup
	port int
	input  chan string
	output chan string
}
func NewClickHouseChecker(port int, input chan string, output chan string, wg *sync.WaitGroup) *ClickHouseChecker {
	s := &ClickHouseChecker{
		port:   port,
		input:  input,
		output: output,
		wg:     wg,
	}
	return s
}
func (s *ClickHouseChecker) Check() {
	parallel := runtime.NumCPU()
	for i := 0; i < parallel; i++ {
		s.wg.Add(1)
		go s.check()
	}
}
func (s *ClickHouseChecker) check() {
	defer s.wg.Done()
	for ip := range s.input {
		if ip == "splitting" || ip == "failed" {
			continue
		}
		if isClickHouse(ip, s.port) {
			s.output <- ip
		}
	}
}
func isClickHouse(ip string, port int) bool {
	conn, err := clickhouse.Open(&clickhouse.Options{
		Addr: []string{fmt.Sprintf("%s:%d", ip, port)},
		// Auth: clickhouse.Auth{
		// 	Database: "default",
		// 	Username: "default",
		// 	Password: "",
		// },
		Settings: clickhouse.Settings{
			"max_execution_time": 1,
		},
		DialTimeout:          time.Second,
		MaxOpenConns:         1,
		MaxIdleConns:         1,
		ConnMaxLifetime:      time.Duration(1) * time.Minute,
		ConnOpenStrategy:     clickhouse.ConnOpenInOrder,
		BlockBufferSize:      10,
		MaxCompressionBuffer: 1024,
	})
	if err != nil {
		golog.Errorf("open %s:%d failed: %v", ip, port, err)
		return false
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	err = conn.Ping(ctx)
	if err != nil {
		golog.Warnf("ping %s:%d failed: %v", ip, port, err)
		return false
	}
	return true
}