package main

import (
	"fmt"
	log "github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"io"
	"net"
)

func handleConnection(tcpConn net.Conn, backend string, sshConfig *ssh.ServerConfig) {
	logger := log.WithFields(log.Fields{
		"client":  tcpConn.RemoteAddr(),
		"backend": backend,
	})
	logger.Debug("Incoming connection")

	// Before use, a handshake must be performed on the incoming net.Conn.
	sshConn, chans, reqs, err := ssh.NewServerConn(tcpConn, sshConfig)
	if err != nil {
		logger.WithFields(log.Fields{
			"err": err,
		}).Warning("SSH Handshake failed")
		return
	}

	logger.WithFields(log.Fields{
		"version": string(sshConn.ClientVersion()),
		"user":    sshConn.User(),
	}).Info("Incoming SSH connection")

	go handleRequests(reqs, sshConn)
	go handleChannels(chans, sshConn)
	go func() {
		err := sshConn.Wait()
		logger.WithFields(log.Fields{
			"err": err,
		}).Info("Connection closed")
	}()
}

type tcpIpForwardRequestPayload struct {
	Raddr string
	Rport uint32
}

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

func verifyForwardRequest(conn *ssh.ServerConn, payload tcpIpForwardRequestPayload) bool {
	return payload.Rport >= 1024
}

func handleRequests(reqs <-chan *ssh.Request, conn *ssh.ServerConn) {

	for req := range reqs {
		logger := log.WithFields(log.Fields{
			"client":  conn.RemoteAddr(),
			"request": req.Type,
		})

		switch req.Type {
		case "tcpip-forward":
			var payload tcpIpForwardRequestPayload
			err := ssh.Unmarshal(req.Payload, &payload)
			if err != nil {
				logger.Warn("Malormed forward request")
				continue
			}

			if !verifyForwardRequest(conn, payload) {
				req.Reply(false, nil)
				continue
			}

			listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", payload.Rport))
			if err != nil {
				logger.Warn("Failed to listen: %s", err)
				req.Reply(false, nil)
				continue
			}
			defer listener.Close()

			logger.WithFields(log.Fields{
				"forward": listener.Addr(),
			}).Info("Forwarding")

			go handleForward(listener, conn, payload)

			req.Reply(true, nil)
		case "keepalive@openssh.com":
			req.Reply(true, nil)
		default:
			logger.Debug("unhandled request")
			req.Reply(false, nil)
		}
	}
}

func handleForward(listener net.Listener, conn *ssh.ServerConn, payload tcpIpForwardRequestPayload) {
	logger := log.WithFields(log.Fields{
		"client":  conn.RemoteAddr(),
		"forward": listener.Addr(),
	})

	for {
		lconn, err := listener.Accept()
		if err != nil {
			logger.Warn("accept failed: %s", err)
			return
		}
		logger.WithFields(log.Fields{
			"remote": lconn.RemoteAddr(),
		}).Info("Forwarding connection")

		channelPayload := ssh.Marshal(&forwardedTcpIpRequestPayload{
			Raddr: payload.Raddr,
			Rport: payload.Rport,
			Laddr: "localhost",
			Lport: 1,
		})

		channel, reqs, err := conn.OpenChannel("forwarded-tcpip", channelPayload)
		if err != nil {
			logger.Warn("open channel failed: %s", err)
			lconn.Close()
			continue
		}

		go func() {
			defer lconn.Close()
			io.Copy(lconn, channel)
		}()

		go func() {
			defer channel.Close()
			io.Copy(channel, lconn)
		}()

		go ssh.DiscardRequests(reqs)
	}
}

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

func handleChannel(newChannel ssh.NewChannel, conn *ssh.ServerConn) {
	logger := log.WithFields(log.Fields{
		"client": conn.RemoteAddr(),
	})

	// 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 {
		logger.Error("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 {
			logger := log.WithFields(log.Fields{
				"client":  conn.RemoteAddr(),
				"request": 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)
				}
			case "exec":
				logger.Debug("exec: %v", req.Payload)
				var payload struct {
					Command string
				}

				err := ssh.Unmarshal(req.Payload, &payload)
				if err != nil {
					log.Error("Malormed exec request")
					continue
				}
				logger.Debug("exec: %v", payload)
			case "env":
				var payload struct {
					Name  string
					Value string
				}

				err := ssh.Unmarshal(req.Payload, &payload)
				if err != nil {
					log.Error("Malormed env request")
					continue
				}
				logger.Printf("env: %v", payload)
			default:
				logger.Debug("unhandled session request: %s", req.Type)
			}
		}
	}()
	connection.Write([]byte("hello!\r\n"))
	io.Copy(connection, connection)
	logger.Info("Session closed")
	return
}