package main

import (
	"crypto/tls"
	"flag"
	"github.com/micro/go-config"
	"github.com/micro/go-config/source/env"
	"github.com/micro/go-config/source/file"
	mflag "github.com/micro/go-config/source/flag"
	log "github.com/sirupsen/logrus"
	"github.com/wercker/journalhook"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"net"
)

func main() {

	flag.String("ssh-privkey", "id_rsa", "ssh private key")
	flag.String("ssh-listen", "0.0.0.0:22", "ssh listen address")
	flag.String("tls-cert", "", "tls certificate")
	flag.String("tls-key", "", "tls certificate")
	flag.String("tls-listen", "", "tls listen address")
	flag.String("log-mode", "", "Logging mode (std, journal)")
	flag.Parse()

	config.Load(
		file.NewSource(
			file.WithPath("/etc/fwd.json"),
		),
		env.NewSource(),
		mflag.NewSource(),
	)

	if config.Get("log", "mode").String("std") == "journal" {
		journalhook.Enable()
	}

	sshConfig := &ssh.ServerConfig{
		NoClientAuth: !config.Get("ssh", "auth").Bool(false),
	}

	// You can generate a keypair with 'ssh-keygen -t rsa'
	key := config.Get("ssh", "privkey").String("id_rsa")

	privateBytes, err := ioutil.ReadFile(key)
	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")
	}

	sshConfig.AddHostKey(private)

	if config.Get("tls", "listen").String("") != "" {
		cer, err := tls.LoadX509KeyPair(
			config.Get("tls", "cert").String(""),
			config.Get("tls", "key").String(""),
		)
		if err != nil {
			log.Println(err)
			return
		}
		tlsConfig := &tls.Config{Certificates: []tls.Certificate{cer}}
		tlsListener, err := tls.Listen("tcp", config.Get("tls", "listen").String(""), tlsConfig)
		if err != nil {
			log.Fatalf("TLS listen failed: %s", err)
		}

		go handleListener(tlsListener, "tls", sshConfig)
	}

	listener, err := net.Listen("tcp", config.Get("ssh", "listen").String("0.0.0.0:22"))
	if err != nil {
		log.Fatalf("SSH listen failed: %s", err)
	}

	handleListener(listener, "ssh", sshConfig)
}

func handleListener(listener net.Listener, backend string, sshConfig *ssh.ServerConfig) {
	log.Printf("Listening on %s (%s)...", listener.Addr(), backend)
	for {
		tcpConn, err := listener.Accept()
		if err != nil {
			log.Printf("accept failed: %s", err)
			continue
		}

		go handleConnection(tcpConn, backend, sshConfig)
	}
}