package ssh

import (
	"golang.org/x/crypto/ssh"
	"io"
	"log"
	"net"
)

/*
	https://stackoverflow.com/questions/21417223/simple-ssh-port-forward-in-golang

	Forward from local port 9000 to remote 192.168.1.150 port 9999 by ssh-server e.g:

	localAddr   = "localhost:9000"
	remoteAddr  = "192.168.1.150:9999"
	proxyAddr   = "192.168.1.100:22"
	proxyUser   = "root"
	proxyPwd    = "password"
*/

// 创建新的 ssh 简单转发通道，本地随机端口；简单实现，不保证稳定；返回监听地址
func GBSshSetupForwardTunnelRandom(remoteAddr, proxyAddr, proxyUser, proxyPwd string) string {

	// Check input forward parameters
	if remoteAddr == "" || proxyAddr == "" || proxyUser == "" || proxyPwd == "" {
		log.Println("ssh.forward : invalid input parameters")
		return ""
	}
	// Resolve address
	addr, err := net.ResolveTCPAddr("tcp", "127.0.0.1:0")
	// 解析失败
	if err != nil {
		log.Println("ssh.forward : resolve tcp addr fail")
		return ""
	}
	// features for ListenTCP, if addr.port is 0, will select random port
	localListener, err := net.ListenTCP("tcp", addr)
	// setup local listener fail
	if err != nil {
		log.Println("ssh.forward : listen tcp fail on " + addr.String() + ", err = " + err.Error())
		return ""
	}
	// obtain local tcp addr
	localTcpAddr := localListener.Addr().(*net.TCPAddr)
	// logging
	log.Println("ssh.forward : setup local listener", localTcpAddr)
	// Serve local proxy server
	go gbSshServeAcceptOnLocal(localListener, remoteAddr, proxyAddr, proxyUser, proxyPwd)
	// Complete
	return localTcpAddr.String()
}

// 创建新的 ssh 简单转发通道，本机指定地址；简单实现，不保证稳定；返回监听地址
func GBSshSetupForwardTunnel(localAddr, remoteAddr, proxyAddr, proxyUser, proxyPwd string) string {

	// Check input forward parameters
	if localAddr == "" || remoteAddr == "" || proxyAddr == "" || proxyUser == "" || proxyPwd == "" {
		log.Println("ssh.forward : invalid input parameters")
		return ""
	}
	// Setup localListener (type net.Listener)
	localListener, err := net.Listen("tcp", localAddr)
	// Setup localListener fail
	if err != nil {
		log.Println("ssh.forward : listen tcp fail on " + localAddr + ", err = " + err.Error())
		return ""
	}
	// obtain local tcp addr
	localTcpAddr := localListener.Addr().(*net.TCPAddr)
	// logging
	log.Println("ssh.forward : setup local listener", localTcpAddr)
	// Serve local proxy server
	go gbSshServeAcceptOnLocal(localListener, remoteAddr, proxyAddr, proxyUser, proxyPwd)
	// Complete
	return localTcpAddr.String()
}

func gbSshServeAcceptOnLocal(localListener net.Listener, remoteAddr, proxyAddr, proxyUser, proxyPwd string) {

	// Setup SSH config (type *ssh.ClientConfig)
	config := &ssh.ClientConfig {
		Config:            ssh.Config{},
		User:              proxyUser,
		Auth: []ssh.AuthMethod {
			ssh.Password(proxyPwd),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		BannerCallback:    nil,
		ClientVersion:     "",
		HostKeyAlgorithms: nil,
		Timeout:           0,
	}

	for {
		// Setup localConn (type net.Conn)
		localConn, err := localListener.Accept()
		if err != nil {
			log.Printf("ssh.forward listen.Accept failed: %v", err)
		}
		// Go forward for this channel
		go gbSshDoForward(localConn, config, remoteAddr, proxyAddr)
	}
}

func gbSshDoForward(localConn net.Conn, config * ssh.ClientConfig, remoteAddr, proxyAddr string) {

	// Setup sshClientConn (type *ssh.ClientConn)
	sshClientConn, err := ssh.Dial("tcp", proxyAddr, config)
	if err != nil {
		log.Printf("ssh.forward ssh.Dial failed: %s", err)
		return
	}

	// Setup sshConn (type net.Conn)
	sshConn, err := sshClientConn.Dial("tcp", remoteAddr)
	// Setup fail
	if err != nil {
		log.Println("ssh.forward ssh.client.Dial failed", err)
		return
	}

	// Logging
	log.Printf("ss.forward chain : %v -> %v -> %v -> %v",
		localConn.RemoteAddr(), localConn.LocalAddr(), proxyAddr, remoteAddr)

	// Copy localConn.Reader to sshConn.Writer
	go func() {
		_, err = io.Copy(sshConn, localConn)
		if err != nil {
			log.Printf("ssh.forward io.Copy (%v -> %v) failed : %v",
				localConn.RemoteAddr(), sshConn.LocalAddr(), err)
		}
	}()

	// Copy sshConn.Reader to localConn.Writer
	go func() {
		_, err = io.Copy(localConn, sshConn)
		if err != nil {
			log.Printf("ssh.forward io.Copy (%v -> %v) failed : %v",
				sshConn.RemoteAddr(), localConn.LocalAddr(), err)
		}
	}()
}