package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"log"
	"net"
	"os"
	"sync"
	"time"

	flags "github.com/spf13/pflag"

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

var (
	port     string
	isHelp   bool
	keyPath  string
	onlyIPv4 bool
	isJson   bool
)

func init() {
	flags.StringVarP(&port, "port", "p", "22", "Port to listen for incoming connections.")
	flags.StringVarP(&keyPath, "key", "k", "", "Path to the host key for the ssh server.\nIf absent it will automatically generate a new one for each run.")
	flags.BoolVarP(&isHelp, "help", "h", false, "Prints this help message and exits.")
	flags.BoolVarP(&onlyIPv4, "onlyIPv4", "4", false, "Only listens on IPv4.")
	flags.BoolVar(&isJson, "json", false, "Log in JSON instead of plain text.")
	flags.Parse()
}

func main() {

	if isHelp {
		printHelp()
		os.Exit(0)
	}

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

	serverConfig := ssh.ServerConfig{PasswordCallback: printConnectionData}
	serverConfig.AddHostKey(hostKey)

	log.Printf("Starting ssh logger on port %s...\n", port)
	ipv4Listener, err := net.Listen("tcp", "0.0.0.0:"+port)
	if err != nil {
		panic(err)
	}
	defer ipv4Listener.Close()

	wg := new(sync.WaitGroup)
	wg.Add(1)
	go startAccepting(ipv4Listener, serverConfig)

	if !onlyIPv4 {
		ipv6Listener, err := net.Listen("tcp", "[::1]:"+port)
		if err != nil {
			panic(err)
		}
		defer ipv6Listener.Close()

		wg.Add(1)
		go startAccepting(ipv6Listener, serverConfig)
	}

	wg.Wait() // Waits until it gets terminated
}

func startAccepting(listener net.Listener, serverConfig ssh.ServerConfig) {
	for {
		con, err := listener.Accept()
		if err != nil {
			log.Println(err)
		}
		go connectionHandler(con, serverConfig)
	}
}

func getKey(path string) []byte {
	if path != "" {
		privateBytes, err := os.ReadFile(path)
		if err != nil {
			log.Fatalf("Failed to load private key: %s", path)
		}
		return privateBytes
	}

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		log.Fatalf("Cannot generate RSA key\n")
	}

	privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privateKeyPem := pem.EncodeToMemory(
		&pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: privateKeyBytes,
		},
	)
	return privateKeyPem
}

func connectionHandler(con net.Conn, serverConfig ssh.ServerConfig) {
	_, _, _, err := ssh.NewServerConn(con, &serverConfig)
	if err != nil {
		return
	}
}

func printConnectionData(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
	responseErr := fmt.Errorf("password rejected for %s", conn.User())
	ip, _, err := net.SplitHostPort(conn.RemoteAddr().String())
	if err != nil {
		fmt.Println(err)
		return nil, responseErr
	}

	if isJson {
		fmt.Printf(
			"{\"date\": \"%s\", \"src\": \"%s\", \"username\": \"%s\", \"password\": \"%s\"}\n",
			time.Now().Format(time.RFC3339),
			ip,
			conn.User(),
			string(password),
		)
	} else {
		log.Printf("SRC=%s USERNAME=%s PASSWORD=%s\n", ip, conn.User(), string(password))
	}

	return nil, responseErr
}

func printHelp() {
	fmt.Println(`A small tool to log IPs, usernames and passwords from incoming ssh-auth requests.

USAGE:
	sshlog [FLAGS]
	
FLAGS:`)
	flags.PrintDefaults()
}