package main

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"net"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/ccding/go-stun/stun"
)

var natTable = make(map[string]int, 4)

var (
	NATType  string
	Nat3Stun bool

	servers []string
	ports   []string
)

type NATInfo struct {
	NAT                 int    `json:"nat"`
	NATType             string `json:"nat_type"`
	UPNPUsable          string `json:"upnp_usable"`
	UPNPPublicIP        string `json:"upnp_public_ip"`
	NATStunUsable       string `json:"nat_stun_usable"`
	UPNP2NAT1StunUsable string `json:"upnp2nat1_stun_usable"`
}

func init() {
	initMapServer()
	initSlicePorts()

	natTable[nat0] = 0
	natTable[nat1] = 1
	natTable[nat2] = 2
	natTable[nat3] = 3
	natTable[nat4] = 4
}

func main() {

	c := stun.NewClient()
	c.SetServerAddr("newstun.bkdomain.cn:3478")
	nat, _, err := c.Discover()
	if err != nil {
		fmt.Println("check nat type failed:", err)
		return
	}

	natType := nat.String()
	natType = strings.ToLower(natType)
	if strings.Contains(natType, "full cone") {
		NATType = nat1
	} else if strings.Contains(natType, "symmetric") {
		NATType = nat4
	} else if strings.Contains(natType, "restricted") {
		if strings.Contains(natType, "port") {
			NATType = nat3
			//if Nat3StunCheck(NATType) {
			//	Nat3Stun = true
			//}
		} else {
			NATType = nat2
		}
	} else if strings.Contains(natType, "public") {
		NATType = nat0
	} else {
		NATType = "unknown"
	}
	info := &NATInfo{
		NATType:             NATType,
		NAT:                 natTable[NATType],
		UPNPUsable:          "false",
		UPNPPublicIP:        "false",
		NATStunUsable:       "false",
		UPNP2NAT1StunUsable: "false",
	}
	data, err := json.Marshal(info)
	if err != nil {
		return
	}

	fmt.Println(string(data))
}

func initMapServer() {
	servers = make([]string, 0)
	servers = []string{"47.95.114.47", "101.201.103.209", "47.95.204.13", "39.105.60.154", "47.94.17.198", "59.110.157.142"}
}

func initSlicePorts() {
	ports = make([]string, 0)
	ports = []string{"35101", "35102", "35103", "5001", "5002", "5003"}
}

func Nat3StunCheck(natType string) bool {
	if natType != "port restricted" {
		fmt.Println("the device nat type is not 'Port Restricted'")
		return false
	}
	fmt.Println("nat3 device check stun enable.")
	localIP := getLocalIP()
	rand.Seed(time.Now().UnixMilli())
	lport := 9600 + rand.Intn(200)
	usedAddr := make([]string, 0)
	lastAddr, flag := "", 0
	count := 0

	for i := 0; i < 10; i++ {
		saltIP := rand.Intn(len(servers))
		saltPORT := rand.Intn(len(ports))
		ip, port := servers[saltIP], ports[saltPORT]
		used := false
		serverAddr := ip + ":" + port
		for _, addr := range usedAddr {
			if addr == serverAddr {
				used = true
			}
		}
		if used {
			continue
		}

		usedAddr = append(usedAddr, serverAddr)

		externADDR, err := udpClient(lport, localIP, ip, port)
		if err != nil {
			fmt.Println("get nat3 stun check external add failed:", err)
			time.Sleep(1 * time.Second)
			continue
		}
		fmt.Println("interior_port:", port, "externalAddr:", externADDR)
		if flag == 0 {
			lastAddr = externADDR
			count++
			flag = 1
		} else {
			if lastAddr != externADDR {
				fmt.Println("the device is nat3 which can't push tunnel used udp proto")
				return false
			}
			count++
			if count >= 3 {
				fmt.Println("the device is nat3 which can push tunnel used udp proto")
				return true
			}
		}
	}

	return false
}

func getLocalIP() string {
	output, err := exec.Command("ip", "route", "get", "223.5.5.5").Output()
	if err != nil {

	}
	rs := strings.Fields(string(output))
	for i, r := range rs {
		if r == "src" && len(rs) >= i+1 {
			return rs[i+1]
		}
	}
	return ""
}

func udpClient(lport int, lip, ip, port string) (string, error) {
	fmt.Println("checking nat3 stun client used udp...")
	defer func() {
		if p := recover(); p != nil {
			fmt.Println("udpClient panic:", p)
		}
		fmt.Println("exit checking nat3 stun client")
	}()
	server := ip + ":" + port
	raddr, err := net.ResolveUDPAddr("udp4", server)
	if err != nil {
		fmt.Println(server, "UDP parsing of local address failed：", err)
		return "", err
	}

	client := lip + ":" + strconv.Itoa(lport)
	laddr, err := net.ResolveUDPAddr("udp4", client)
	if err != nil {
		fmt.Println(laddr, "UDP parsing of local address failed：", err)
		return "", err
	}
	conn, err := net.DialUDP("udp4", laddr, raddr)
	if err != nil {
		fmt.Println("failed to connect UDP:", err)
		return "", err
	}
	sendBuffer := make(map[string]int)
	sendBuffer["source_port"] = lport
	buff, err := json.Marshal(sendBuffer)
	if err != nil {
		fmt.Println("marshal failed:special:", err)
		b := fmt.Sprintf("{\"source_port\": %d}", lport)
		buff = []byte(b)
	}

	err = conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
	if err != nil {
		fmt.Println("Error setting write deadline:", err)
		return "", err
	}

	if _, err = conn.Write(buff); err != nil {
		fmt.Println("nat3 stun check write failed:", err)
		return "", err
	}
	defer conn.Close()
	err = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
	if err != nil {
		fmt.Println("Error setting read deadline:", err)
		return "", err
	}
	b := make([]byte, 1024)
	n, err := conn.Read(b)
	if err != nil {
		fmt.Println("nat3 stun check read failed:", err)
		return "", err
	}
	fmt.Println("nat3Stun outputs:", string(b))
	recvBuffer := make(map[string]interface{})
	if err = json.Unmarshal(b[:n], &recvBuffer); err != nil {
		fmt.Println("nat stun check unmarshal read content failed:", err)
		return "", err
	}
	externIP, ok := recvBuffer["Addr"].(string)
	if !ok {
		fmt.Println("nat3 stun check unmarshal read content not include extern ip")
		fmt.Println("nat3 stun check read content:", recvBuffer)
		return "", err
	}
	externPORT, ok := recvBuffer["Port"].(float64)
	if !ok {
		fmt.Println("nat stun check unmarshal read content not include extern port")
		fmt.Println("nat3 stun check read content:", recvBuffer)
		return "", err
	}
	return fmt.Sprintf("%s:%d", externIP, int(externPORT)), nil

}