package net

import (
	"b612.me/starcrypto"
	"b612.me/starlog"
	"b612.me/starnet"
	"crypto/elliptic"
	"encoding/csv"
	"fmt"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
	"math/rand"
	"net"
	"os"
	"strings"
	"time"
)

var (
	listenAddr    string
	keyFile       string
	KeyPasswd     string
	outpath       string
	curlUrl       string
	serverVersion string
	curlArg       []string
	passwds       []string
	allowAny      bool
	logPath       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", "", "输出文件")
	cmdSSHJar.Flags().StringVarP(&serverVersion, "version", "v", "SSH-2.0-OpenSSH_8.0", "SSH版本")
	cmdSSHJar.Flags().StringVarP(&curlUrl, "curl", "c", "", "Curl URL")
	cmdSSHJar.Flags().StringSliceVarP(&passwds, "allow-passwds", "P", nil, "密码列表，格式：[用户名]:[密码]")
	cmdSSHJar.Flags().BoolVarP(&allowAny, "allow-any", "A", false, "允许任意密码登录")
	cmdSSHJar.Flags().StringVarP(&logPath, "log", "L", "", "日志文件")
}

var cmdSSHJar = &cobra.Command{
	Use:   "sshjar",
	Short: "SSH蜜罐",
	Long:  "SSH蜜罐",
	Run: func(cmd *cobra.Command, args []string) {
		var mypwds [][]string
		for _, v := range passwds {
			args := strings.SplitN(v, ":", 2)
			if len(args) == 2 {
				mypwds = append(mypwds, args)
			}
		}
		runSSHHoneyJar(SSHJar{
			listenAddr: listenAddr,
			keyFile:    keyFile,
			keyPasswd:  KeyPasswd,
			outpath:    outpath,
			logpath:    logPath,
			version:    serverVersion,
			passwds:    mypwds,
			allowAny:   allowAny,
		})
	},
}

type SSHJar struct {
	listenAddr string
	keyFile    string
	keyPasswd  string
	outpath    string
	logpath    string
	version    string
	passwds    [][]string
	allowAny   bool
}

func runSSHHoneyJar(jar SSHJar) {
	if jar.logpath != "" {
		starlog.SetLogFile(jar.logpath, starlog.Std, true)
	}
	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
		}
	}
	conn := csv.NewWriter(f)
	defer f.Close()
	defer conn.Flush()
	config := &ssh.ServerConfig{
		ServerVersion: jar.version,
		// 密码验证回调函数
		PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
			starlog.Infof("Login attempt from %s with %s by %s\n", c.RemoteAddr(), c.User(), string(pass))
			data := []string{time.Now().Format("2006-01-02 15:04:05"), c.RemoteAddr().String(), c.User(), string(pass)}
			if f != nil {
				conn.Write(data)
				conn.Flush()
			}
			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.NewSimpleRequest(curlUrl, "POST").SetBodyDataBytes([]byte(starnet.BuildQuery(data))).Do()
					}
				}()
			}
			perm := &ssh.Permissions{
				Extensions: map[string]string{
					"user":   c.User(),
					"passwd": string(pass),
				},
			}
			if jar.allowAny {
				return perm, nil
			}
			for _, v := range jar.passwds {
				if c.User() == v[0] && string(pass) == v[1] {
					return perm, nil
				}
			}
			return nil, fmt.Errorf("password rejected for %q", c.User())
		},
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			if jar.allowAny {
				return &ssh.Permissions{
					Extensions: map[string]string{
						"user": conn.User(),
					},
				}, nil
			}
			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)
	for {
		conn, err := listener.Accept()
		if err != nil {
			continue
		}
		starlog.Infof("New connection from %s\n", conn.RemoteAddr())
		go func(conn net.Conn) {
			sConn, chans, reqs, err := ssh.NewServerConn(conn, config)
			if err != nil {
				starlog.Errorf("SSH handshake failed: %v\n", err)
				return
			}
			defer sConn.Close()
			defer starlog.Noticef("Connection from %s closed\n", sConn.RemoteAddr())
			go ssh.DiscardRequests(reqs)
			for newChannel := range chans {
				if newChannel.ChannelType() != "session" {
					newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
					continue
				}
				channel, requests, err := newChannel.Accept()
				if err != nil {
					starlog.Errorf("Failed to accept channel: %v", err)
					continue
				}
				go handleSession(channel, requests, sConn)
			}
		}(conn)
	}

}

func handleSession(channel ssh.Channel, requests <-chan *ssh.Request, conn *ssh.ServerConn) {
	defer channel.Close()
	term := terminal.NewTerminal(channel, "$ ") // 设置 shell 提示符
	term.AutoCompleteCallback = nil             // 禁用自动补全
	for req := range requests {
		switch req.Type {
		case "pty-req":
			// 接受伪终端请求（攻击者希望获得交互式体验）
			req.Reply(true, nil)
			term.SetSize( // 简单设置终端尺寸
				24, // 行
				80, // 列
			)
		case "shell":
			req.Reply(true, nil)
			go func() {
				for {
					line, err := term.ReadLine()
					if err != nil {
						break
					}
					starlog.Infof("[%s %s] Command: %s\n", conn.RemoteAddr(), conn.Permissions.Extensions["user"], line)
					time.Sleep(time.Millisecond * 200)
					term.Write([]byte(FakeCommand(line)))
				}
			}()
		case "exec":
			// 处理非交互式命令（如 ssh user@host 'ls -l'）
			var payload struct{ Command string }
			ssh.Unmarshal(req.Payload, &payload)
			req.Reply(true, nil)

			// 记录并返回假输出
			starlog.Infof("[%s %s] Exec: %s\n", conn.RemoteAddr(), conn.Permissions.Extensions["user"], payload.Command)
			term.Write([]byte(FakeCommand(payload.Command)))
			channel.Close()
		default:
			req.Reply(false, nil)
		}
	}
}

func FakeCommand(cmd string) string {
	// 按命令类型分级模拟
	switch {
	//---------------- 系统信息探测类 ----------------
	case strings.Contains(cmd, "uname -a"):
		return "Linux core 6.1.0-21-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.90-1 (2024-05-03) x86_64 GNU/Linux\n\n"

	case strings.Contains(cmd, "cat /etc/os-release"):
		return `PRETTY_NAME="Debian GNU/Linux 11 (bullseye)"
NAME="Debian GNU/Linux"
VERSION_ID="11"
VERSION="11 (bullseye)"
VERSION_CODENAME=bullseye
ID=debian
HOME_URL="https://www.debian.org/"
SUPPORT_URL="https://www.debian.org/support"
BUG_REPORT_URL="https://bugs.debian.org/"
`

	case strings.Contains(cmd, "free -h"):
		return `              total        used        free      shared  buff/cache   available
Mem:           1022Gi        12Gi       3.0Gi       1.0Gi        47Gi        480Gi
Swap:            0B          0B          0B
`

	//---------------- 敏感文件诱导类 ----------------
	case strings.Contains(cmd, "ls /home"):
		return "admin  backup  devops  secret\n"

	case strings.Contains(cmd, "ls /var/log"):
		return `auth.log  apache2  payment_system.log  database_backup.log
`
	case strings.Contains(cmd, "ls"):
		return "password.txt\n"

	case strings.Contains(cmd, "cat /etc/passwd"):
		return `root:x:0:0:root:/root:/bin/bash
admin:x:1000:1000:,,,:/home/admin:/bin/bash
mysql:x:106:113:MySQL Server,,,:/nonexistent:/bin/false
core:x:0:0:,,,:/root:/bin/bash
`

	//---------------- 网络配置诱导类 ----------------
	case strings.Contains(cmd, "ifconfig") || strings.Contains(cmd, "ip addr"):
		return `eth0: flags=4163<UP,BROADCAST,RUNNING,MULTICAST>  mtu 1500
        inet 192.168.1.105  netmask 255.255.255.0  broadcast 192.168.1.255
        inet6 fe80::250:56ff:fec0:8888  prefixlen 64  scopeid 0x20<link>
        ether 00:50:56:c0:88:88  txqueuelen 1000  (Ethernet)
        RX packets 123456  bytes 123456789 (117.7 MiB)
        RX errors 0  dropped 0  overruns 0  frame 0
        TX packets 98765  bytes 9876543 (9.4 MiB)
        TX errors 0  dropped 0 overruns 0  carrier 0  collisions 0

lo: flags=73<UP,LOOPBACK,RUNNING>  mtu 65536
        inet 127.0.0.1  netmask 255.0.0.0
        loop  txqueuelen 1000  (Local Loopback)
`

	case strings.Contains(cmd, "netstat -antp"):
		return `Active Internet connections (servers and established)
Proto Recv-Q Send-Q Local Address           Foreign Address         State       PID/Program name
tcp        0      0 0.0.0.0:22              0.0.0.0:*               LISTEN      1234/sshd
tcp        0      0 192.168.1.105:5432      203.0.113.5:43892       ESTABLISHED 5678/postgres
tcp6       0      0 :::8080                 :::*                    LISTEN      91011/java
`

	//---------------- 凭证钓鱼类 ----------------
	case strings.Contains(cmd, "mysql -u root -p"):
		return "ERROR 1045 (28000): Access denied\n"

	case strings.Contains(cmd, "sudo -l"):
		return `Matching Defaults entries for admin on this host:
    env_reset, mail_badpass, secure_path=/usr/local/sbin\:/usr/local/bin\:/usr/sbin\:/usr/bin\:/sbin\:/bin

User admin may run the following commands on honeypot:
    (ALL) NOPASSWD: /usr/bin/vim /etc/shadow
`

	//---------------- 进程服务类 ----------------
	case strings.Contains(cmd, "ps aux"):
		return `USER       PID %CPU %MEM    VSZ   RSS TTY      STAT START   TIME COMMAND
root         1  0.0  0.1 169020 13184 ?        Ss   May01   0:12 /sbin/init
admin     1234  0.3  2.1 1123456 178912 ?      Sl   May01  12:34 /opt/payment_system/payment_processor --debug
`

	//---------------- 定制化陷阱 ----------------
	case strings.Contains(cmd, "find / -name *.db"):
		return `/var/lib/mysql/transactions.db
/home/backup/internal_users.db
`

	case strings.Contains(cmd, "curl"):
		return `<!DOCTYPE HTML PUBLIC "-//IETF//DTD HTML 2.0//EN">
<html><head>
<title>404 Not Found</title>
</head><body>
<h1>Not Found</h1>
<p>The requested URL was not found on this server.</p>
</body></html>
`

	default:
		// 模糊响应增加真实感
		if rand.Intn(100) > 70 { // 30%概率返回"command not found"
			return "sh: command not found: " + strings.Fields(cmd)[0] + "\n"
		}
		return "\n"
	}
}