Victorique/vtqe/tools/tcp.go

package tools

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"

	"b612.me/starainrt"

	"github.com/spf13/cobra"
)

var tcpcmd = &cobra.Command{
	Use:   "tcp",
	Short: "发送并监听tcp数据包",
	Long:  "发送并监听tcp数据包",
	Run: func(this *cobra.Command, args []string) {
		if len(args) != 1 {
			fmt.Println("请指定远程tcp地址")
			return
		}
		l, _ := this.Flags().GetString("lport")
		p, _ := this.Flags().GetString("rport")
		a, _ := this.Flags().GetString("laddr")
		s, _ := this.Flags().GetBool("local")
		b, _ := this.Flags().GetBool("byte")
		laddr, err := net.ResolveTCPAddr("tcp", a+":"+l)
		if err != nil {
			fmt.Println(err)
			return
		}
		if s {
			tcplisten, err := net.ListenTCP("tcp", laddr)
			if err != nil {
				fmt.Println(err)
				return
			}
			fmt.Println("监听已建立")
			defer tcplisten.Close()
			go func() {
				for {
					conn, err := tcplisten.AcceptTCP()
					if err != nil {
						fmt.Printf("Error Connect From %s : %s\n", conn.RemoteAddr(), err.Error())
						continue
					}
					fmt.Printf("Accept Connect From %s\n", conn.RemoteAddr())
					go func(conns *net.TCPConn) {
						for {
							buf := make([]byte, 204800)
							n, err := conns.Read(buf)
							if err != nil {
								fmt.Printf("Error from %s Where Message=%s\n", conns.RemoteAddr(), err.Error())
								conn.Close()
								return
							}
							fmt.Printf("Receive Msg From %s : %s\n", conns.RemoteAddr(), string(buf[0:n]))
							if b {
								fmt.Println(buf[0:n])
							}
						}

					}(conn)
				}
			}()
		}
		mytcp, err := net.DialTimeout("tcp", args[0]+":"+p, time.Second*15)
		if err != nil {
			fmt.Println(err)
			if s {
				for {
					time.Sleep(time.Second * 10)
				}
			}
			return
		}
		defer mytcp.Close()
		fmt.Println("TCP连接已建立")
		go func() {
			var err error
			for {
				txt := starainrt.MessageBox("", "")
				if txt == "" {
					continue
				}
				if !b {
					_, err = mytcp.Write([]byte(txt))
				} else {
					var sendbyte []byte
					bytes := strings.Split(txt, ",")
					for _, v := range bytes {
						ints, _ := strconv.Atoi(v)
						if ints < 0 || ints > 255 {
							continue
						}
						sendbyte = append(sendbyte, byte(ints))
					}
					_, err = mytcp.Write(sendbyte)
				}
				if err != nil {
					fmt.Printf("Error from %s Where Message=%s\n", mytcp.RemoteAddr().String(), err.Error())
					return
				}
			}
		}()
		for {
			buf := make([]byte, 204800)
			n, err := mytcp.Read(buf)
			if err != nil {
				fmt.Printf("Error from %s Where Message=%s\n", mytcp.RemoteAddr().String(), err.Error())
				return
			}

			fmt.Printf("Receive Msg From %s : %s\n", mytcp.RemoteAddr().String(), string(buf[0:n]))
			if b {
				fmt.Println(buf[0:n])
			}
		}
	},
}

/*

 */

func init() {
	Maincmd.AddCommand(tcpcmd)
	tcpcmd.Flags().BoolP("byte", "b", false, "发送二进制数据")
	tcpcmd.Flags().StringP("lport", "l", "1127", "本地监听端口")
	tcpcmd.Flags().StringP("rport", "p", "1127", "远程连接端口")
	tcpcmd.Flags().StringP("laddr", "a", "0.0.0.0", "本地监听ip")
	tcpcmd.Flags().BoolP("local", "s", false, "启动本地监听")
	tcpsendcmd.Flags().StringP("port", "p", "1127", "远程连接端口")
	tcpsendcmd.Flags().StringP("addr", "a", "0.0.0.0", "远程监听ip")
	tcpsendcmd.Flags().StringP("regexp", "r", "", "正则匹配字符串")
	tcprecvcmd.Flags().StringP("port", "p", "1127", "本地监听端口")
	tcprecvcmd.Flags().StringP("addr", "a", "0.0.0.0", "本地监听ip")
	tcpcmd.AddCommand(tcpsendcmd, tcprecvcmd)
}

var cansend, backsck bool = false, false
var tcpsendcmd = &cobra.Command{
	Use:   "send",
	Short: "通过tcp发送文件",
	Long:  "通过tcp接受文件",
	Run: func(this *cobra.Command, args []string) {
		stop := false
		p, _ := this.Flags().GetString("port")
		a, _ := this.Flags().GetString("addr")
		r, _ := this.Flags().GetString("regexp")
		if len(args) != 1 {
			this.Help()
			return
		}
		mytcp, err := net.DialTimeout("tcp", a+":"+p, time.Second*15)
		if err != nil {
			fmt.Println(err)
			return
		}
		fmt.Println("TCP连接已建立")
		defer mytcp.Close()
		go func() {
			for {
				buf := make([]byte, 1048576)
				n, err := mytcp.Read(buf)
				if err != nil {
					if !stop {
						fmt.Println(err)
						os.Exit(1)
					}
					continue
				}
				data := string(buf[0:n])
				if data == "oksend!" {
					cansend = true
				} else if data == "backsck" {
					backsck = true
				} else {
					fmt.Println(data)
					os.Exit(1)
				}
			}
		}()
		file := args[0]
		if file[len(file)-1:] == "\\" || file[len(file)-1:] == "/" {
			file = file[0 : len(file)-1]
		}
		file, _ = filepath.Abs(file)
		Scan(file, "", mytcp, r)
		stop = true
		mytcp.Write([]byte("jane&0&0"))
		return
	},
}

func tcpupload(file, prefix string, mytcp net.Conn) error {
	if runtime.GOOS == "windows" {
		file = strings.Replace(file, "/", "\\", -1)
		prefix = strings.Replace(prefix, "/", "\\", -1)
	}
	fmt.Println(file)
	fpsrc, err := os.Open(file)
	if err != nil {
		fmt.Println(err)
		return err
	}
	stat, _ := os.Stat(file)
	filebig := float64(stat.Size())
	_, err = mytcp.Write([]byte("victorique&" + stat.Name() + "&" + prefix))
	if err != nil {
		fmt.Println(err)
		return err
	}
	var sakura int = 0
	for !cansend {
		time.Sleep(time.Millisecond * 100)
		sakura++
		if sakura > 50 {
			mytcp.Write([]byte{1, 9, 9, 6, 1, 1, 2, 7})
			time.Sleep(time.Millisecond * 100)
			mytcp.Write([]byte("victorique&" + stat.Name() + "&" + prefix))
		}
	}
	cansend = false
	sum := 0
	for {
		buf := make([]byte, 524288)
		n, err := fpsrc.Read(buf)
		if err != nil {
			if err == io.EOF {
				err = nil
				mytcp.Write([]byte{1, 9, 9, 6, 1, 1, 2, 7})
				for !backsck {
					time.Sleep(time.Millisecond * 150)
				}
				backsck = false
				break
			}
			fmt.Println(err)
			return err
		}
		_, err = mytcp.Write(buf[0:n])
		if err != nil {
			fmt.Println(err)
			return err
		}
		sum += n
		fmt.Printf("当前已传输：%f\r", float64(sum)/filebig*100)
	}
	fmt.Printf("当前已传输：%f\n\n", 100.000000)
	fpsrc.Close()
	return nil
}

func Scan(path, prefix string, mytcp net.Conn, reg string) {
	var err error
	var regp *regexp.Regexp
	if prefix != "" {
		prefix += "/"
	}
	if reg != "" {
		regp, err = regexp.Compile(reg)
		if err != nil {
			fmt.Println(err)
			return
		}
	}
	if runtime.GOOS == "windows" {
		path = strings.Replace(path, "/", "\\", -1)
		prefix = strings.Replace(prefix, "/", "\\", -1)
	}
	if starainrt.IsFile(path) {
		if reg != "" {
			if regp.MatchString(path) {
				tcpupload(path, prefix, mytcp)
			}
		} else {
			tcpupload(path, prefix, mytcp)
		}
	} else if starainrt.IsFolder(path) {
		dir, err := ioutil.ReadDir(path)
		if err != nil {
			fmt.Println(err)
			return
		}
		for _, v := range dir {
			if v.IsDir() && (v.Name() != "." || v.Name() != "..") {
				mytcp.Write([]byte("b612&" + prefix + v.Name() + "&0"))
				time.Sleep(time.Millisecond * 10)
				Scan(path+"/"+v.Name(), prefix+v.Name(), mytcp, reg)
			} else {
				if reg != "" {
					if regp.MatchString(path + "/" + v.Name()) {
						err = tcpupload(path+"/"+v.Name(), prefix, mytcp)
					}
				} else {
					err = tcpupload(path+"/"+v.Name(), prefix, mytcp)
				}
				if err != nil {
					fmt.Println(err)
					return
				}
			}
		}
	}

}

var tcprecvcmd = &cobra.Command{
	Use:   "recv",
	Short: "通过tcp接收文件",
	Long:  "通过tcp接收文件",
	Run: func(this *cobra.Command, args []string) {
		p, _ := this.Flags().GetString("port")
		a, _ := this.Flags().GetString("addr")
		laddr, err := net.ResolveTCPAddr("tcp", a+":"+p)
		if err != nil {
			fmt.Println(err)
			return
		}
		tcplisten, err := net.ListenTCP("tcp", laddr)
		if err != nil {
			fmt.Println(err)
			return
		}
		fmt.Println("TCP监听已建立")
		defer tcplisten.Close()
		for {
			conn, err := tcplisten.AcceptTCP()
			if err != nil {
				fmt.Printf("Error Connect From %s : %s\n", conn.RemoteAddr(), err.Error())
				continue
			}
			fmt.Printf("Accept Connect From %s\n", conn.RemoteAddr())
			go func(conns *net.TCPConn) {
				canrecv := false
				name := ""
				var fpdst *os.File
				defer conns.Close()
				var sum int
				for {
					buf := make([]byte, 393216)
					n, err := conns.Read(buf)

					if !canrecv {
						if err != nil {
							fmt.Println(err)
							return
						}
						data := string(buf[0:n])
						str := strings.Split(data, "&")
						if len(str) != 3 {
							continue
						}
						if str[0] == "victorique" {
							canrecv = true
							name = str[1]
							prefix := str[2]
							if runtime.GOOS != "windows" {
								prefix = strings.Replace(prefix, "\\", "/", -1)
							}
							if prefix == "" {
								fpdst, err = os.Create("./" + name)
							} else {
								fpdst, err = os.Create(prefix + name)
							}
							if err != nil {
								fmt.Println(err)
								conns.Write([]byte(err.Error()))
								return
							}
							conns.Write([]byte("oksend!"))
						} else if str[0] == "b612" {
							os.MkdirAll(str[1], 0644)
						} else if str[0] == "jane" {
							return
						}
					} else {
						if buf[n-1] == byte(7) && buf[n-2] == byte(2) && buf[n-3] == byte(1) {
							ok := func() bool {
								for k, v := range []byte{1, 9, 9, 6, 1, 1, 2, 7} {
									if buf[n-8+k] != v {
										return false
									}
								}
								return true
							}()
							if ok {
								if n > 8 {
									fpdst.Write(buf[0 : n-8])
									sum += n - 8
								}
								fmt.Printf("已写入:%d 字节\n", sum)
								fmt.Println(name + "接收成功\n")
								fpdst.Close()
								canrecv = false
								sum = 0
								conns.Write([]byte("backsck"))
								continue
							}
						}
						fpdst.Write(buf[0:n])
						sum += n
						fmt.Printf("已写入:%d 字节\r", sum)
					}
				}
			}(conn)
		}

	},
}