package net

import (
	"b612.me/starcrypto"
	"b612.me/starlog"
	"b612.me/starnet"
	"crypto/elliptic"
	"encoding/json"
	"fmt"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
	"net"
	"os"
	"os/signal"
	"strings"
)

var (
	listenAddr string
	keyFile    string
	KeyPasswd  string
	outpath    string
	curlUrl    string
	curlArg    []string
)

func init() {
	cmdSSHJar.Flags().StringVarP(&listenAddr, "listen", "l", "0.0.0.0:22", "监听地址")
	cmdSSHJar.Flags().StringVarP(&keyFile, "key", "k", "", "私钥文件")
	cmdSSHJar.Flags().StringVarP(&KeyPasswd, "passwd", "p", "", "私钥密码")
	cmdSSHJar.Flags().StringVarP(&outpath, "output", "o", "", "输出文件")
}

var cmdSSHJar = &cobra.Command{
	Use:   "sshjar",
	Short: "SSH蜜罐",
	Long:  "SSH蜜罐",
	Run: func(cmd *cobra.Command, args []string) {
		runSSHHoneyJar(listenAddr, keyFile, KeyPasswd, outpath)
	},
}

func runSSHHoneyJar(listenAddr, keyFile, KeyPasswd, outpath string) {
	var f *os.File
	var err error
	if outpath != "" {
		f, err = os.OpenFile(outpath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
		if err != nil {
			starlog.Errorf("Failed to open file %s (%s)", outpath, err)
			return
		}
	}
	defer f.Close()
	config := &ssh.ServerConfig{
		// 密码验证回调函数
		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
			starlog.Infof("Login attempt from %s with %s %s\n", c.RemoteAddr(), c.User(), string(pass))
			data := []string{c.RemoteAddr().String(), c.User(), string(pass)}
			bts, _ := json.Marshal(data)
			if f != nil {
				f.Write(bts)
				f.Write([]byte("\n"))
			}
			if curlUrl != "" {
				go func() {
					data := map[string]string{
						"ip":     c.RemoteAddr().String(),
						"user":   c.User(),
						"passwd": string(pass),
					}
					if curlArg != nil && len(curlArg) > 0 {
						for _, v := range curlArg {
							args := strings.SplitN(v, ":", 2)
							if len(args) == 2 {
								data[args[0]] = args[1]
							}
						}
						starnet.Curl(starnet.NewRequests(curlUrl, []byte(starnet.BuildQuery(data)), "POST"))
					}
				}()
			}
			return nil, fmt.Errorf("password rejected for %q", c.User())
		},
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			return nil, fmt.Errorf("public key rejected for %q", conn.User())
		},
	}
	if keyFile == "" {
		secKey, _, err := starcrypto.GenerateEcdsaKey(elliptic.P256())
		if err != nil {
			starlog.Errorf("Failed to generate ECDSA key (%s)", err)
			return
		}
		key, err := ssh.NewSignerFromKey(secKey)
		if err != nil {
			starlog.Errorf("Failed to generate signer from key (%s)", err)
			return
		}
		config.AddHostKey(key)
	} else {
		keyByte, err := os.ReadFile(keyFile)
		if err != nil {
			starlog.Errorf("Failed to read private key from %s (%s)", keyFile, err)
			return
		}
		var key ssh.Signer
		if KeyPasswd != "" {
			key, err = ssh.ParsePrivateKeyWithPassphrase(keyByte, []byte(KeyPasswd))
		} else {
			key, err = ssh.ParsePrivateKey(keyByte)
		}
		if err != nil {
			starlog.Errorf("Failed to load private key from %s (%s)", keyFile, err)
			return
		}
		config.AddHostKey(key)
	}
	listener, err := net.Listen("tcp", listenAddr)
	if err != nil {
		starlog.Errorf("Failed to listen on %s (%s)", listenAddr, err)
		return
	}
	starlog.Noticeln("SSH HoneyJar is listening on", listenAddr)
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)
	for {
		select {
		case <-sig:
			starlog.Noticef("SSH HoneyJar is shutting down")
			listener.Close()
			return
		default:
		}
		conn, err := listener.Accept()
		if err != nil {
			continue
		}
		starlog.Infof("New connection from %s\n", conn.RemoteAddr())
		go func(conn net.Conn) {
			ssh.NewServerConn(conn, config)
			conn.Close()
		}(conn)
	}

}