package main

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

	flags "github.com/spf13/pflag"

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

var (
	address string
	isHelp  bool
	keyPath string
	isJson  bool
)

func init() {
	flags.StringVarP(&address, "address", "a", "0.0.0.0:22", "Address 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.BoolVar(&isJson, "json", false, "Log in JSON instead of plain text.")
	flags.Parse()
}

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

	privateKey, err := getKey(keyPath)
	if err != nil {
		log.Fatalf("Failed to get keys: %v\n", err)
	}
	hostKey, err := ssh.ParsePrivateKey(privateKey)
	if err != nil {
		log.Fatalf("Failed to parse private key: %v\n", err)
	}

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

	log.Printf("Starting ssh logger on %s...\n", address)
	listener, err := net.Listen("tcp", address)
	if err != nil {
		log.Fatalf("Failed to open listener on '%s': %v\n", address, err)
	}
	defer listener.Close()

	for {
		con, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept incoming connection: %v\n", err)
		}
		go func(con net.Conn, serverConfig ssh.ServerConfig) {
			err := connectionHandler(con, serverConfig)
			if err != nil {
				log.Printf("Failed to handle connection: %v\n", err)
			}
		}(con, serverConfig)
	}
}

func getKey(path string) ([]byte, error) {
	if path != "" {
		privateBytes, err := os.ReadFile(path)
		if err != nil {
			return nil, fmt.Errorf("load private key '%s': %w", path, err)
		}
		return privateBytes, nil
	}

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("generate RSA key: %w", err)
	}

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

func connectionHandler(con net.Conn, serverConfig ssh.ServerConfig) error {
	_, _, _, err := ssh.NewServerConn(con, &serverConfig)
	if err != nil {
		return fmt.Errorf("connection handler: %w", err)
	}
	return nil
}

func printConnectionData(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
	responseErr := fmt.Errorf("password rejected for %s", conn.User())
	address = conn.RemoteAddr().String()
	ip, _, err := net.SplitHostPort(address)
	if err != nil {
		log.Printf("Could not split address '%s' in ip and port: %v\n", address, 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()
}