package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"

	"golang.org/x/crypto/ssh"
)

func main() {

	// In the latest version of crypto/ssh (after Go 1.3), the SSH server type has been removed
	// in favour of an SSH connection type. A ssh.ServerConn is created by passing an existing
	// net.Conn and a ssh.ServerConfig to ssh.NewServerConn, in effect, upgrading the net.Conn
	// into an ssh.ServerConn

	config := &ssh.ServerConfig{
		NoClientAuth: true,
	}

	// You can generate a keypair with 'ssh-keygen -t rsa'
	privateBytes, err := ioutil.ReadFile("id_rsa")
	if err != nil {
		log.Fatal("Failed to load private key (./id_rsa)")
	}

	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		log.Fatal("Failed to parse private key")
	}

	config.AddHostKey(private)

	// Once a ServerConfig has been configured, connections can be accepted.
	listener, err := net.Listen("tcp", "0.0.0.0:2200")
	if err != nil {
		log.Fatalf("Failed to listen on 2200 (%s)", err)
	}

	// Accept all connections
	log.Print("Listening on 2200...")
	for {
		tcpConn, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept incoming connection (%s)", err)
			continue
		}
		// Before use, a handshake must be performed on the incoming net.Conn.
		sshConn, chans, reqs, err := ssh.NewServerConn(tcpConn, config)
		if err != nil {
			log.Printf("Failed to handshake (%s)", err)
			continue
		}

		log.Printf("New SSH connection from %s (%s)", sshConn.RemoteAddr(), sshConn.ClientVersion())
		go handleRequests(reqs, sshConn)
		go handleChannels(chans)
	}
}

type tcpIpForwardRequestPayload struct {
	Raddr string
	Rport uint32
}

type forwardedTcpIpRequestPayload struct {
	Raddr string
	Rport uint32
	Laddr string
	Lport uint32
}

func handleRequests(reqs <-chan *ssh.Request, conn *ssh.ServerConn) {
	for req := range reqs {
		log.Printf("req: %s", req.Type)
		switch req.Type {
		case "tcpip-forward":
			var payload tcpIpForwardRequestPayload

			err := ssh.Unmarshal(req.Payload, &payload)
			if err != nil {
				log.Printf("Malormed forward request")
				continue
			}

			log.Printf("%v", payload)

			listener, err := net.Listen("tcp", "0.0.0.0:0")
			if err != nil {
				log.Printf("Failed to listen")
				req.Reply(false, nil)
				continue
			}
			defer listener.Close()

			log.Printf("Listening on %v", listener.Addr())

			go func() {
				for {
					lconn, err := listener.Accept()
					log.Printf("%v %v", lconn, err)

					channelPayload := ssh.Marshal(&forwardedTcpIpRequestPayload{
						Raddr: payload.Raddr,
						Rport: payload.Rport,
						Laddr: "localhost",
						Lport: 8000,
					})
					channel, reqs, err := conn.OpenChannel("forwarded-tcpip", channelPayload)
					if err != nil {
						log.Printf("%s: open channel failed", err)
						return
					}
					go func() {
						defer lconn.Close()
						io.Copy(lconn, channel)
					}()
					go func() {
						defer channel.Close()
						io.Copy(channel, lconn)
					}()
					go ssh.DiscardRequests(reqs)
				}
			}()

			req.Reply(true, nil)
		case "keepalive@openssh.com":
			req.Reply(true, nil)
		default:
			log.Printf("Ignoring request")
			req.Reply(false, nil)
		}
	}
}

func handleChannels(chans <-chan ssh.NewChannel) {
	// Service the incoming Channel channel in go routine
	for newChannel := range chans {
		go handleChannel(newChannel)
	}
}

func handleChannel(newChannel ssh.NewChannel) {
	// Since we're handling a shell, we expect a
	// channel type of "session". The also describes
	// "x11", "direct-tcpip" and "forwarded-tcpip"
	// channel types.
	if t := newChannel.ChannelType(); t != "session" {
		newChannel.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
		return
	}

	// At this point, we have the opportunity to reject the client's
	// request for another logical connection
	connection, requests, err := newChannel.Accept()
	if err != nil {
		log.Printf("Could not accept channel (%s)", err)
		return
	}

	// Sessions have out-of-band requests such as "shell", "pty-req" and "env"
	go func() {
		for req := range requests {
			log.Printf(req.Type)
			switch req.Type {
			case "shell":
				// We only accept the default shell
				// (i.e. no command in the Payload)
				if len(req.Payload) == 0 {
					req.Reply(true, nil)
				}
			}
		}
	}()
	connection.Write([]byte("hello!\r\n"))
	io.Copy(connection, connection)
	log.Printf("Session closed")
	return
}