package tools

import (
	"fmt"
	"io"
	"net"
	"strconv"
	"time"

	"github.com/spf13/cobra"
)

func init() {
	Maincmd.AddCommand(netcmd)
	netcmd.AddCommand(netforwardcmd, natscmd, natccmd)
	netforwardcmd.Flags().BoolP("tcp", "t", false, "TCP转发")
	netforwardcmd.Flags().BoolP("udp", "u", false, "UDP转发")
	netforwardcmd.Flags().IntP("port", "p", 1127, "转发本地端口")
	netforwardcmd.Flags().StringP("address", "a", "0.0.0.0", "监听地址")

	natscmd.Flags().BoolP("tcp", "t", false, "TCP穿透")
	natscmd.Flags().BoolP("udp", "u", false, "UDP穿透")
	natscmd.Flags().IntP("port", "p", 1127, "监听本地端口开始点")
	natscmd.Flags().StringP("key", "k", "sakura", "认证密码")
	natscmd.Flags().StringP("address", "a", "0.0.0.0", "监听本地地址")
	natscmd.Flags().IntP("number", "n", 1, "监听端口数量")

	natccmd.Flags().BoolP("tcp", "t", false, "TCP穿透")
	natccmd.Flags().BoolP("udp", "u", false, "UDP穿透")
	natccmd.Flags().StringP("port", "p", "1127", "转发穿透端口")
	natccmd.Flags().StringP("key", "k", "1127", "认证密码")
	natccmd.Flags().StringP("address", "a", "127.0.0.1", "转发穿透地址")
}

var netcmd = &cobra.Command{
	Use:   "net",
	Short: "端口转发，穿透相关",
	Long:  "端口转发，穿透相关",
}

var netforwardcmd = &cobra.Command{
	Use:   "forward",
	Short: "TCP/UDP端口转发",
	Long:  "TCP/UDP端口转发",
	Run: func(cmd *cobra.Command, args []string) {
		t, _ := cmd.Flags().GetBool("tcp")
		u, _ := cmd.Flags().GetBool("udp")
		p, _ := cmd.Flags().GetInt("port")
		a, _ := cmd.Flags().GetString("address")
		if len(args) != 1 {
			cmd.Help()
			return
		}
		ra := args[0]
		if !t && !u {
			fmt.Println("请指定TCP或UDP协议")
			return
		}

		localtcp, err := net.ResolveTCPAddr("tcp", a+":"+strconv.Itoa(p))
		if err != nil {
			fmt.Println(err)
			return
		}
		remotetcp, err := net.ResolveTCPAddr("tcp", ra)
		if err != nil {
			fmt.Println(err)
			return
		}
		localudp, _ := net.ResolveUDPAddr("udp", a+":"+strconv.Itoa(p))
		remoteudp, _ := net.ResolveUDPAddr("udp", ra)
		if u {
			go UdpForward(localudp, remoteudp)
		}
		if t {
			go TcpForward(localtcp, remotetcp)
		}

		for {
			time.Sleep(time.Second * 100)
		}
	},
}

var natscmd = &cobra.Command{
	Use:   "nats",
	Short: "TCP/UDP内网穿透服务端",
	Long:  "TCP/UDP内网穿透服务端",
	Run: func(cmd *cobra.Command, args []string) {
		t, _ := cmd.Flags().GetBool("tcp")
		u, _ := cmd.Flags().GetBool("udp")
		p, _ := cmd.Flags().GetInt("port")
		k, _ := cmd.Flags().GetString("key")
		a, _ := cmd.Flags().GetString("address")
		n, _ := cmd.Flags().GetInt("number")
		/*udplocal, err := net.ResolveUDPAddr("udp", a+":"+p)
		if err != nil {
			fmt.Println(err)
			return
		}
		*/
		for i := 0; i < n; i++ {
			tcplocal, err := net.ResolveTCPAddr("tcp", a+":"+strconv.Itoa(p+i))
			if err != nil {
				fmt.Println(err)
				return
			}
			udplocal, _ := net.ResolveUDPAddr("udp", a+":"+strconv.Itoa(p+i))
			if !t && !u {
				fmt.Println("请至少指定一种网络协议！")
				return
			}
			if t {
				go TcpNatServer(tcplocal, k)
			}
			if u {
				go UdpNatServer(udplocal, k)
			}
		}
		for {
			time.Sleep(time.Second * 10)
		}
	},
}

var natccmd = &cobra.Command{
	Use:   "natc",
	Short: "TCP/UDP内网穿透客户端",
	Long:  "TCP/UDP内网穿透客户端",
	Run: func(cmd *cobra.Command, args []string) {
		t, _ := cmd.Flags().GetBool("tcp")
		u, _ := cmd.Flags().GetBool("udp")
		p, _ := cmd.Flags().GetString("port")
		k, _ := cmd.Flags().GetString("key")
		a, _ := cmd.Flags().GetString("address")
		/*udplocal, err := net.ResolveUDPAddr("udp", a+":"+p)
		if err != nil {
			fmt.Println(err)
			return
		}
		*/
		if len(args) != 1 {
			cmd.Help()
			return
		}
		if !t && !u {
			fmt.Println("请至少指定一种网络协议！")
			return
		}
		tcprmt, err := net.ResolveTCPAddr("tcp", args[0])
		if err != nil {
			fmt.Println(err)
			return
		}
		tcpforward, err := net.ResolveTCPAddr("tcp", a+":"+p)
		if err != nil {
			fmt.Println(err)
			return
		}
		udprmt, _ := net.ResolveUDPAddr("udp", args[0])
		udpforward, _ := net.ResolveUDPAddr("udp", a+":"+p)
		if t {
			go TcpNatClient(tcprmt, tcpforward, k)
		}
		if u {
			go UdpNatClient(udprmt, udpforward, k)
		}
		for {
			time.Sleep(time.Second * 10)
		}
	},
}

type tcpnat struct {
	Conn *net.TCPConn
	Msg  []byte
	Date int64
}

func UdpNatClient(udpremote, udpforward *net.UDPAddr, pwd string) {
	var trueconn *net.UDPConn
	var err error
	ismain := false
	dialudp := func(udpaddr *net.UDPAddr) (*net.UDPConn, error) {
		trueconn, err := net.DialUDP("udp", nil, udpaddr)
		if err != nil {
			return nil, err
		}
		return trueconn, nil
	}
	for {
		for trueconn, err = dialudp(udpremote); err != nil; trueconn, err = dialudp(udpremote) {
			fmt.Println(err)
			time.Sleep(5 * time.Second)
		}

		fmt.Println("远端服务器连接成功" + trueconn.RemoteAddr().String())
		trueconn.Write([]byte(pwd))
		go func() {
			for {
				time.Sleep(time.Second * 5)
				if !ismain {
					_, err := trueconn.Write([]byte(pwd))
					if err != nil {
						return
					}
				}
			}
		}()
		for {
			var conn1, conn2 *net.UDPConn
			buf := make([]byte, 7)
			n, err := trueconn.Read(buf)
			if n != 7 || err != nil {
				fmt.Println(err)
				trueconn.Close()
				ismain = false
				break
			}
			switch string(buf) {
			case "yesyesi":
				ismain = true
				fmt.Println("主节点认证成功")
				go func() {
					for {
						time.Sleep(time.Second * 2)
						trueconn.Write([]byte("vicregs"))
					}
				}()
			case "newconn":
				if !ismain {
					trueconn.Close()
					break
				}
				for conn1, err = dialudp(udpforward); err != nil; conn1, err = dialudp(udpforward) {
					time.Sleep(time.Second * 2)
				}
				fmt.Println("已连接转发远端：" + conn1.RemoteAddr().String())
				for conn2, err = dialudp(udpremote); err != nil; conn2, err = dialudp(udpremote) {
					time.Sleep(time.Second * 2)
				}
				conn2.Write([]byte("v%2^f&K"))
				fmt.Println("已连接主控远端：" + conn2.RemoteAddr().String())
				go UdpCopy(conn1, conn2)
				go UdpCopy(conn2, conn1)
			}
		}
	}
}

func TcpNatClient(tcpremote, tcpforward *net.TCPAddr, pwd string) {
	var trueconn *net.TCPConn
	var err error
	ismain := false
	dialtcp := func(tcpaddr *net.TCPAddr) (*net.TCPConn, error) {
		trueconn, err := net.DialTCP("tcp", nil, tcpaddr)
		if err != nil {
			return nil, err
		}
		return trueconn, nil
	}
	for {
		for trueconn, err = dialtcp(tcpremote); err != nil; trueconn, err = dialtcp(tcpremote) {
			fmt.Println(err)
			time.Sleep(5 * time.Second)
		}

		fmt.Println("远端服务器连接成功" + trueconn.RemoteAddr().String())
		trueconn.Write([]byte(pwd))
		go func() {
			for {
				time.Sleep(time.Second * 5)
				if !ismain {
					_, err := trueconn.Write([]byte(pwd))
					if err != nil {
						return
					}
				}
			}
		}()
		for {
			var conn1, conn2 *net.TCPConn
			buf := make([]byte, 7)
			n, err := trueconn.Read(buf)
			if n != 7 || err != nil {
				fmt.Println(err)
				trueconn.Close()
				ismain = false
				break
			}
			switch string(buf) {
			case "yesyesi":
				ismain = true
				fmt.Println("主节点认证成功")
			case "newconn":
				if !ismain {
					trueconn.Close()
					break
				}
				for conn1, err = dialtcp(tcpforward); err != nil; conn1, err = dialtcp(tcpforward) {
					time.Sleep(time.Second * 2)
				}
				fmt.Println("已连接转发远端：" + conn1.RemoteAddr().String())
				for conn2, err = dialtcp(tcpremote); err != nil; conn2, err = dialtcp(tcpremote) {
					time.Sleep(time.Second * 2)
				}
				conn2.Write([]byte("v%2^f&K"))
				fmt.Println("已连接主控远端：" + conn2.RemoteAddr().String())
				go TcpCopy(conn1, conn2)
				go TcpCopy(conn2, conn1)
			}
		}
	}
}

func TcpCopy3(dst net.Conn, src *net.TCPConn) {
	defer dst.Close()
	defer src.Close()
	io.Copy(dst, src)
}

func TcpCopy2(dst *net.TCPConn, src net.Conn) {
	defer src.Close()
	defer dst.Close()
	io.Copy(dst, src)
}

func TcpCopy(dst, src *net.TCPConn) {
	defer dst.Close()
	defer src.Close()
	io.Copy(dst, src)
}

func UdpCopy(dst, src *net.UDPConn) {
	defer dst.Close()
	defer src.Close()
	for {
		buf := make([]byte, 131072)
		n, err := src.Read(buf)
		dst.Write(buf[0:n])
		if err != nil {
			fmt.Println(err)
			break
		}
	}
}

type udpnat struct {
	Conn      *net.UDPAddr
	Msg       []byte
	Married   bool
	LastTrans int64
}

func UdpNatServer(udplocal *net.UDPAddr, pwd string) {
	var trueconn *net.UDPAddr
	isconn := false
	var lastcheck int64
	udplistener, err := net.ListenUDP("udp", udplocal)
	defer udplistener.Close()
	if err != nil {
		fmt.Println(err)
		return
	}
	udpool := make(map[string]*udpnat)
	udpdial := make(map[string]*net.UDPAddr)
	fmt.Println("UDP已建立监听：" + udplistener.LocalAddr().String())
	go func() {
		for {
			time.Sleep(time.Second * 300)
			pool := []string{}
			for k, v := range udpool {
				if time.Now().Unix()-v.LastTrans > 600 {
					pool = append(pool, k)
				}
			}
			for _, v := range pool {
				if udpool[v].Married {
					delete(udpdial, udpool[v].Conn.String())
				}
				delete(udpool, v)
			}
		}
	}()
	for {
		buf := make([]byte, 131072)
		n, conn, errp := udplistener.ReadFromUDP(buf)
		go func() {
			if !isconn {
				if errp != nil {
					return
				}
				if string(buf[0:len(pwd)]) == pwd {
					isconn = true
					trueconn = conn
					udplistener.WriteToUDP([]byte("yesyesi"), trueconn)
					fmt.Println("穿透客户端已注册")
					lastcheck = time.Now().Unix()
					go func() {
						for {
							time.Sleep(time.Second * 5)
							if time.Now().Unix()-lastcheck > 4 {
								fmt.Println("穿透客户端已下线")
								isconn = false
								break
							}
						}
					}()
				}
			} else {
				if string(buf[0:7]) == "vicregs" {
					lastcheck = time.Now().Unix()
					return
				}
				if string(buf[0:7]) == "v%2^f&K" {
					fmt.Println("穿透客户端已建立新连接")
					for k, v := range udpool {
						if !v.Married {
							v.Conn = conn
							udplistener.WriteToUDP(v.Msg, conn)
							v.Msg = []byte{}
							udpaddr, _ := net.ResolveUDPAddr("udp", k)
							udpdial[conn.String()] = udpaddr
							v.Married = true
							break
						}
					}
					return
				} else {
					if v, ok := udpool[conn.String()]; ok {
						if !v.Married {
							for i := 0; i < n; i++ {
								v.Msg = append(v.Msg, buf[i])
							}
							if len(v.Msg) > 10485760 {
								v.Msg = []byte{}
							}
						} else {
							_, err := udplistener.WriteToUDP(buf[0:n], v.Conn)
							v.LastTrans = time.Now().Unix()
							if errp != nil || err != nil {
								fmt.Println(errp, err)
								udplistener.WriteToUDP(nil, v.Conn)
								udplistener.WriteToUDP(nil, conn)
								delete(udpool, conn.String())
								delete(udpdial, v.Conn.String())
								return
							}
						}
						return
					}
					if v, ok := udpdial[conn.String()]; ok {
						_, err := udplistener.WriteToUDP(buf[0:n], v)
						udpool[v.String()].LastTrans = time.Now().Unix()
						if errp != nil || err != nil {
							fmt.Println(errp, err)
							udplistener.WriteToUDP(nil, v)
							udplistener.WriteToUDP(nil, conn)
							delete(udpool, v.String())
							delete(udpdial, conn.String())
							return
						}
						return
					}
					fmt.Println("链接已加入等待列表")
					udpool[conn.String()] = &udpnat{
						Msg:       buf[0:n],
						Married:   false,
						LastTrans: time.Now().Unix(),
					}
					_, err := udplistener.WriteToUDP([]byte("newconn"), trueconn)
					if err != nil {
						fmt.Println("穿透客户端已下线")
						isconn = false
					}
				}
			}
		}()
	}
	fmt.Println("执行完毕")
}

func TcpNatServer(tcplocal *net.TCPAddr, pwd string) {
	var trueconn *net.TCPConn
	isconn := false
	tcplistener, err := net.ListenTCP("tcp", tcplocal)
	defer tcplistener.Close()
	var waitconn []tcpnat
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println("TCP已建立监听：" + tcplistener.Addr().String())
	for {
		conn, err := tcplistener.AcceptTCP()
		if err != nil {
			return
		}
		go func() {
			for {
				now := time.Now().Unix()
				if len(waitconn) != 0 {
					if now-waitconn[0].Date > 8 {
						waitconn[0].Conn.Close()
						if len(waitconn) > 1 {
							waitconn = waitconn[1:]
						} else {
							waitconn = []tcpnat{}
						}
					}
				}
				time.Sleep(time.Millisecond * 600)
			}
		}()
		go func() {
			if !isconn {
				buf := make([]byte, len(pwd))
				n, err := conn.Read(buf)
				if n != len(pwd) {
					conn.Close()
					return
				}
				if err != nil {
					conn.Close()
					return
				}
				if string(buf) == pwd {
					isconn = true
					trueconn = conn
					conn.Write([]byte("yesyesi"))
					fmt.Println("穿透客户端已注册")
					go func() {
						for {
							time.Sleep(time.Second * 2)
							_, err := trueconn.Write([]byte("vicb612"))
							if err != nil {
								trueconn.Close()
								isconn = false
								return
							}
						}
					}()
				}
				return
			}
		}()
		go func() {
			if isconn {
				buf := make([]byte, 7)
				n, err := conn.Read(buf)
				if n != 7 {
					conn.Close()
					return
				}
				if err != nil {
					conn.Close()
					return
				}
				if string(buf) == "v%2^f&K" {
					fmt.Println("穿透客户端已建立新连接")
					if len(waitconn) != 0 {
						conn.Write(waitconn[0].Msg)
						go TcpCopy(waitconn[0].Conn, conn)
						go TcpCopy(conn, waitconn[0].Conn)
						if len(waitconn) > 1 {
							waitconn = waitconn[1:]
						} else {
							waitconn = []tcpnat{}
						}
					} else {
						conn.Close()
					}
				} else {
					fmt.Println("链接已加入等待列表")
					tcpnats := tcpnat{Msg: buf, Conn: conn, Date: time.Now().Unix()}
					waitconn = append(waitconn, tcpnats)
					_, err := trueconn.Write([]byte("newconn"))
					if err != nil {
						fmt.Println("穿透客户端已下线")
						trueconn.Close()
						isconn = false
					}
				}
			}
		}()
	}

}

func TcpForward(tcplocal, tcpforward *net.TCPAddr) {
	tcplistener, err := net.ListenTCP("tcp", tcplocal)
	if err != nil {
		fmt.Println(err)
		return
	}
	dialtcp := func() (*net.Conn, error) {
		tcpclient, err := net.DialTimeout("tcp", tcpforward.String(), time.Second*10)
		if err != nil {
			return nil, err
		}
		return &tcpclient, nil
	}
	fmt.Println("TCP监听已建立：" + tcplocal.String())
	for {
		conn, err := tcplistener.AcceptTCP()
		if err != nil {
			continue
		}
		fmt.Println(conn.RemoteAddr().String() + " 已连接")
		tcpclient, err := dialtcp()
		if err != nil {
			fmt.Println(err)
			conn.Write([]byte("Dial To Remote Addr Error"))
			conn.Close()
			continue
		}
		fmt.Println("成功连接到转发端口，本地端口：" + (*tcpclient).LocalAddr().String())
		go TcpCopy2(conn, *tcpclient)
		go TcpCopy3(*tcpclient, conn)
	}
}

func UdpForward(udplocal, udpforward *net.UDPAddr) {
	udplistener, err := net.ListenUDP("udp", udplocal)
	if err != nil {
		fmt.Println(err)
		return
	}
	dialudp := func() (*net.UDPConn, error) {
		udpclient, err := net.DialUDP("udp", nil, udpforward)
		if err != nil {
			return nil, err
		}
		return udpclient, nil
	}
	fmt.Println("UDP监听已建立：" + udplocal.String())
	udppool := make(map[string]*net.UDPConn)
	for {
		buf := make([]byte, 131072)
		n, conn, err := udplistener.ReadFromUDP(buf)
		udpclient, ok := udppool[conn.String()]
		if !ok || udpclient == nil {
			fmt.Println(conn.String() + " 已连接")
			udpclient, err = dialudp()
			if err != nil {
				fmt.Println(err)
				udplistener.WriteToUDP([]byte("Dial To Remote Addr Error"), conn)
				continue
			}
			udppool[conn.String()] = udpclient
			fmt.Println("成功连接到转发端口，本地端口：" + udpclient.LocalAddr().String())
		}
		go udpclient.Write(buf[0:n])
		go func(conns *net.UDPAddr, udpclient *net.UDPConn) {
			bufs := make([]byte, 131072)
			n, err := udpclient.Read(bufs)
			udplistener.WriteToUDP(bufs[0:n], conns)
			if err != nil {
				udpclient.Close()
				udppool[conn.String()] = nil
			}
		}(conn, udpclient)
		if err != nil {
			udpclient.Close()
			udppool[conn.String()] = nil
		}
	}
}