package ussher

import (
	"fmt"
	"log"
	"net"
	"strings"
	"time"

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

func RunSSHServer(config *Config) error {
	sshConfig := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			return config.Keystore.CheckPublicKey(conn.User(), key)
		},
	}

	private, err := config.GetPrivateKey()
	if err != nil {
		return err
	}
	sshConfig.AddHostKey(private)

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

	// Accept all connections
	log.Printf("Listening on %s...", listenAddr)
	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, sshConfig)
		if err != nil {
			log.Printf("Failed to handshake (%s)", err)
			continue
		}

		log.Printf("New SSH connection from %s (%s)", sshConn.RemoteAddr(), sshConn.ClientVersion())
		// Discard all global out-of-band Requests
		go ssh.DiscardRequests(reqs)
		// Accept all channels
		go handleChannels(chans, sshConn, config)
	}
}

func handleChannels(chans <-chan ssh.NewChannel, conn *ssh.ServerConn, config *Config) {
	for newChannel := range chans {
		go handleChannel(newChannel, conn, config)
	}
}

func handleChannel(newChannel ssh.NewChannel, conn *ssh.ServerConn, config *Config) {
	// We're emulating a shell command; so this is a "session"
	if t := newChannel.ChannelType(); t != "session" {
		newChannel.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
		return
	}

	connection, requests, err := newChannel.Accept()
	if err != nil {
		log.Printf("Could not accept channel (%s)", err)
		return
	}
	go func() {
		var exit bool
		for {
			if exit {
				break
			}
			select {
			case req := <-requests:
				switch req.Type {
				case "shell":
					log.Printf("Requested shell; none can be provided")
					connection.Close()
					exit = true
				case "exec":
					appstring := string(req.Payload[4:]) // skip the first four bytes, which are length of string
					appspaced := strings.SplitN(appstring, " ", 2)
					if len(appspaced) == 0 {
						exit = true
						break
					}
					v, ok := config.SSHApps[appspaced[0]]
					if ok {
						err := v.HandleExec(appstring, conn, connection)
						if err != nil {
							fmt.Fprintln(connection.Stderr(), err)
							log.Printf("Got error while handling ssh session: %s", err)
						}
					} else {
						log.Printf("Can't find app %s", appspaced[0])
					}
					exit = true
				default:
					// Ignore things like 'pty' and 'env' variables.
					continue
				}
			case <-time.After(5 * time.Second):
				fmt.Println("connection timeout")
				exit = true
			}
		}
		connection.Close()
	}()
}