package main

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/go-admin-team/go-admin-core/sdk/api"
	"github.com/gorilla/websocket"
	"github.com/ziutek/telnet"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

type TgWs struct {
	api.Api
}

var upGrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024 * 1024 * 10,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type ShellInfoStruct struct {
	Proto  string `json:"proto"`
	IpAddr string `json:"ipaddr"`
	Port   string `json:"port"`
}

type wsWrapper struct {
	*websocket.Conn
}

func main() {
	telnetHandle(rw io.ReadWriter, ip, port string, errhandle func(string)) {

}

func init() {
	routerCheckRole = append(routerCheckRole, registerTgWsRouter)
}

// 需认证的路由代码
func registerTgWsRouter(v1 *gin.RouterGroup, authMiddleware *jwt.GinJWTMiddleware) {
	api := apis.TgWs{}
	r := v1.Group("")
	{
		// 协议、IP、端口
		r.GET("/tgws/:proto/:ipaddr/:port", api.TgWsWeb)
	}
}

func (e TgWs) TgWsWeb(c *gin.Context) {
	// 初始化返回信息
	err := e.MakeContext(c).Errors
	if err != nil {
		e.Logger.Error(err)
		e.Error(500, err, fmt.Sprintf(" %s ", err.Error()))
		return
	}

	// 升级为websocket
	wsConn, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		e.Logger.Error(err)
		e.Error(500, err, "websocket client connect error")
		return
	}
	defer wsConn.Close()

	proto := c.Param("proto")
	ipaddr := c.Param("ipaddr")
	port := c.Param("port")
	shellinfo := ShellInfoStruct{
		Proto:  proto,
		IpAddr: ipaddr,
		Port:   port,
	}
	quitChan := make(chan bool, 1)
	go websocketHandle(wsConn, shellinfo, quitChan)
	<-quitChan
}

func websocketHandle(con *websocket.Conn, shellinfo ShellInfoStruct, exitCh chan bool) {
	defer setQuit(exitCh)
	rw := io.ReadWriter(&wsWrapper{con})
	webprintln := func(data string) {
		rw.Write([]byte(data + "\r\n"))
	}
	con.SetCloseHandler(func(code int, text string) error {
		con.Close()
		return nil
	})
	switch shellinfo.Proto {
	case "ssh":
		sshHandle(rw, shellinfo.IpAddr, shellinfo.Port, "XXX", "XXX", webprintln)
	case "telnet":
		telnetHandle(rw, shellinfo.IpAddr, shellinfo.Port, webprintln)
	case "bind_shell":
		bindShellHandler(rw, shellinfo.IpAddr, shellinfo.Port, webprintln)
	default:
		webprintln("Not Support Protocol '" + shellinfo.Proto + "'")
	}
	return
}

func (wsw *wsWrapper) Write(p []byte) (n int, err error) {
	writer, err := wsw.Conn.NextWriter(websocket.TextMessage)
	if err != nil {
		return 0, err
	}
	defer writer.Close()
	return writer.Write(p)
}

func (wsw *wsWrapper) Read(p []byte) (n int, err error) {
	for {
		msgType, reader, err := wsw.Conn.NextReader()
		if err != nil {
			return 0, err
		}
		if msgType != websocket.TextMessage {
			continue
		}
		return reader.Read(p)
	}
}

// SSH连接
func sshHandle(rw io.ReadWriter, ip, port, user, passwd string, errhandle func(string)) {
	sshConfig := &ssh.ClientConfig{
		User:    user,
		Auth:    []ssh.AuthMethod{ssh.Password(passwd)},
		Timeout: 6 * time.Second,
	}
	sshConfig.HostKeyCallback = ssh.InsecureIgnoreHostKey()
	if port == "" {
		ip = ip + ":22"
	} else {
		ip = ip + ":" + port
	}
	client, err := ssh.Dial("tcp", ip, sshConfig)
	if err != nil {
		errhandle(err.Error())
		return
	}
	defer client.Close()
	session, err := client.NewSession()
	if err != nil {
		errhandle(err.Error())
		return
	}
	defer session.Close()
	fd := int(os.Stdin.Fd())
	session.Stdout = rw
	session.Stderr = rw
	session.Stdin = rw
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}
	termWidth, termHeight, err := terminal.GetSize(fd)
	err = session.RequestPty("xterm", termHeight, termWidth, modes)
	if err != nil {
		errhandle(err.Error())
	}
	err = session.Shell()
	if err != nil {
		errhandle(err.Error())
	}
	err = session.Wait()
	if err != nil {
		errhandle(err.Error())
	}
	return
}

// telnet连接
func telnetHandle(rw io.ReadWriter, ip, port string, errhandle func(string)) {
	if port == "" {
		ip = ip + ":23"
	} else {
		ip = ip + ":" + port
	}
	con, err := telnet.Dial("tcp", ip)
	if err != nil {
		errhandle(err.Error())
		return
	}
	defer con.Close()
	buf := make([]byte, 16*1024)

	// 从远端读取返回结果并回显页面
	go func() {
		for {
			n, err := con.Read(buf)
			if err != nil {
				errhandle(err.Error())
				break
			}
			_, err = rw.Write(buf[:n])
			if err != nil {
				errhandle(err.Error())
				break
			}
		}
	}()

	for {
		// 从页面读取命令
		n, err := rw.Read(buf)
		if err != nil {
			errhandle(err.Error())
			break
		}
		if buf[0] == 13 { // 处理换行
			data := []byte{telnet.CR, telnet.LF}
			_, err = con.Write(data)
		} else {
			_, err = con.Write(buf[:n])
		}
		if err != nil {
			errhandle(err.Error())
			break
		}
	}
	return
}

// 正向shell
func bindShellHandler(rw io.ReadWriter, ip, port string, errhandle func(string)) {
	server := ip + ":" + port
	//获取命令行参数 socket地址
	addr, err := net.ResolveTCPAddr("tcp4", server)
	if err != nil {
		errhandle(err.Error())
		return
	}
	//建立tcp连接
	con, err := net.DialTCP("tcp4", nil, addr)
	if err != nil {
		errhandle(err.Error())
		return
	}
	rw.Write([]byte("reverse shell connected " + "\r\n"))
	defer con.Close()
	buf := make([]byte, 16*1024)

	go func() {
		for {
			n, err := con.Read(buf)
			if err != nil {
				errhandle(err.Error())
				break
			}
			_, err = rw.Write(buf[:n])
			if err != nil {
				errhandle(err.Error())
				break
			}
		}
	}()

	for {
		n, err := rw.Read(buf)
		if err != nil {
			errhandle(err.Error())
			break
		}
		_, err = rw.Write(buf[:n])
		if err != nil {
			errhandle(err.Error())
			break
		}

		if buf[0] == 13 {
			data := []byte{telnet.LF}
			_, err = con.Write(data)
			rw.Write([]byte("\r\n"))
		} else {
			_, err = con.Write(buf[:n])
		}
		if err != nil {
			errhandle(err.Error())
			break
		}
	}
	return
}

func setQuit(ch chan bool) {
	ch <- true
}