package credssp

import (
	"crypto/md5"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"io"
	"math/big"
	"strings"
)

// Options 控制 CredSSP 兼容策略，仅调试和特殊 Windows 版本需要覆盖 // Logan.Li
type Options struct {
	Version     int
	PubKeyStyle string
	ForceMIC    bool
	NoKeyExch   bool
	UseCBT      bool
	OmitVersion bool
}

// ChallengeInfo 暴露 NTLM challenge 诊断信息，不包含凭据 // Logan.Li
type ChallengeInfo struct {
	Version         int
	TargetName      string
	Flags           uint32
	TargetInfoBytes int
	RequiresMIC     bool
	DomainName      string
	ComputerName    string
}

// ReadChallengeInfo 只读取 CredSSP/NTLM challenge，不发送认证数据 // Logan.Li
func ReadChallengeInfo(conn *tls.Conn, version int) (ChallengeInfo, error) {
	if version == 0 {
		version = defaultCredSSPVersion
	}
	negotiate := buildNegotiate()
	req, err := encodeTSRequestVersion(version, [][]byte{negotiate}, nil, nil, nil)
	if err != nil {
		return ChallengeInfo{}, err
	}
	if _, err := conn.Write(req); err != nil {
		return ChallengeInfo{}, err
	}
	challengeReq, err := readTSRequest(conn)
	if err != nil {
		return ChallengeInfo{}, err
	}
	if len(challengeReq.NegoTokens) == 0 {
		return ChallengeInfo{}, fmt.Errorf("CredSSP challenge 缺少 nego token")
	}
	challenge, err := parseChallenge(challengeReq.NegoTokens[0].Data)
	if err != nil {
		return ChallengeInfo{}, err
	}
	return ChallengeInfo{
		Version:         challengeReq.Version,
		TargetName:      challenge.TargetName,
		Flags:           challenge.Flags,
		TargetInfoBytes: len(challenge.TargetInfo),
		RequiresMIC:     requiresMIC(challenge.TargetInfo),
		DomainName:      utf16leString(findAVPair(challenge.TargetInfo, msvAvNbDomainName)),
		ComputerName:    utf16leString(findAVPair(challenge.TargetInfo, msvAvNbComputerName)),
	}, nil
}

// Authenticate 在已升级 TLS 的 RDP 连接上执行最小 CredSSP/NTLMv2 认证 // Logan.Li
func Authenticate(conn *tls.Conn, domain, user, password string) error {
	return AuthenticateWithOptions(conn, domain, user, password, Options{Version: defaultCredSSPVersion, PubKeyStyle: "binding"})
}

// AuthenticateWithOptions 使用指定 CredSSP 版本和 pubKeyAuth 算法执行认证 // Logan.Li
func AuthenticateWithOptions(conn *tls.Conn, domain, user, password string, opts Options) error {
	if opts.Version == 0 {
		opts.Version = defaultCredSSPVersion
	}
	if opts.PubKeyStyle == "" {
		opts.PubKeyStyle = "binding"
	}
	domain, user = splitDomainUser(domain, user)
	negotiate := buildNegotiate()
	req, err := encodeTSRequestVersion(opts.Version, [][]byte{negotiate}, nil, nil, nil)
	if err != nil {
		return err
	}
	if _, err := conn.Write(req); err != nil {
		return err
	}

	challengeReq, err := readTSRequest(conn)
	if err != nil {
		return fmt.Errorf("读取 NTLM challenge 失败: %w", err)
	}
	if len(challengeReq.NegoTokens) == 0 {
		return fmt.Errorf("CredSSP challenge 缺少 nego token")
	}
	challenge, err := parseChallenge(challengeReq.NegoTokens[0].Data)
	if err != nil {
		return err
	}
	ntlmOpts := ntlmOptions{
		ForceMIC:       opts.ForceMIC,
		DisableKeyExch: opts.NoKeyExch,
		OmitVersion:    opts.OmitVersion,
	}
	if opts.UseCBT {
		ntlmOpts.ChannelBindings = channelBindingsHash(conn)
	}
	auth, sec, err := buildAuthenticate(challenge, domain, user, password, "AIOPS", negotiate, ntlmOpts)
	if err != nil {
		return err
	}
	pubKey, err := peerPublicKey(conn, opts.PubKeyStyle)
	if err != nil {
		return err
	}
	clientNonce := randomBytes(32)
	pubKeyAuthPlain := buildPubKeyAuth(pubKey, clientNonce, opts.PubKeyStyle)
	pubKeyAuth := sec.encrypt(pubKeyAuthPlain)
	var nonce []byte
	if opts.PubKeyStyle == "binding" {
		nonce = clientNonce
	}
	req, err = encodeTSRequestVersion(opts.Version, [][]byte{auth}, nil, pubKeyAuth, nonce)
	if err != nil {
		return err
	}
	if _, err := conn.Write(req); err != nil {
		return err
	}

	serverReq, err := readTSRequest(conn)
	if err != nil {
		return fmt.Errorf("读取服务端公钥校验失败(flags=0x%08x mic=%v targetInfo=%d): %w", challenge.Flags, requiresMIC(challenge.TargetInfo), len(challenge.TargetInfo), err)
	}
	if serverReq.ErrorCode != 0 {
		return fmt.Errorf("CredSSP 返回错误码: %s", ntStatusMessage(serverReq.ErrorCode))
	}

	creds, err := encodePasswordCredentials(domain, user, password)
	if err != nil {
		return err
	}
	encryptedCreds := sec.encrypt(creds)
	req, err = encodeTSRequestVersion(opts.Version, nil, encryptedCreds, nil, nil)
	if err != nil {
		return err
	}
	_, err = conn.Write(req)
	return err
}

func buildPubKeyAuth(pubKey, clientNonce []byte, style string) []byte {
	switch style {
	case "raw-plus-one":
		out := append([]byte(nil), pubKey...)
		if len(out) > 0 {
			out[0]++
		}
		return out
	case "binding":
		return credSSPBindingHash(clientNonce, pubKey)
	default:
		return pubKey
	}
}

func ntStatusMessage(code int) string {
	u := uint32(code)
	switch u {
	case 0xC0000064:
		return "0xC0000064 用户不存在"
	case 0xC000006A:
		return "0xC000006A 密码错误"
	case 0xC000006D:
		return "0xC000006D 用户名或认证信息无效"
	case 0xC000006E:
		return "0xC000006E 账号登录限制"
	case 0xC0000070:
		return "0xC0000070 工作站限制"
	case 0xC0000071:
		return "0xC0000071 密码已过期"
	case 0xC0000072:
		return "0xC0000072 账号已禁用"
	case 0xC000015B:
		return "0xC000015B 当前用户未被授予远程登录权限"
	case 0xC0000193:
		return "0xC0000193 账号已过期"
	case 0xC0000234:
		return "0xC0000234 账号已锁定，请先在 Windows/域控侧解锁后再连接"
	default:
		return fmt.Sprintf("0x%08X", u)
	}
}

func credSSPBindingHash(clientNonce, pubKey []byte) []byte {
	h := sha256.New()
	h.Write([]byte("CredSSP Client-To-Server Binding Hash\x00"))
	h.Write(clientNonce)
	h.Write(pubKey)
	return h.Sum(nil)
}

func channelBindingsHash(conn *tls.Conn) []byte {
	state := conn.ConnectionState()
	if len(state.PeerCertificates) == 0 {
		return nil
	}
	cert := state.PeerCertificates[0]
	alg := cert.SignatureAlgorithm.String()
	var digest []byte
	switch {
	case strings.Contains(alg, "MD5"):
		sum := md5.Sum(cert.Raw)
		digest = sum[:]
	case strings.Contains(alg, "SHA1"):
		sum := sha1.Sum(cert.Raw)
		digest = sum[:]
	default:
		sum := sha256.Sum256(cert.Raw)
		digest = sum[:]
	}
	appData := append([]byte("tls-server-end-point:"), digest...)
	gss := make([]byte, 16+len(appData))
	copy(gss[16:], appData)
	gss[12] = byte(len(appData))
	gss[13] = byte(len(appData) >> 8)
	gss[14] = byte(len(appData) >> 16)
	gss[15] = byte(len(appData) >> 24)
	sum := md5.Sum(gss)
	return sum[:]
}

func readTSRequest(r io.Reader) (*tsRequest, error) {
	der, err := readDER(r)
	if err != nil {
		return nil, err
	}
	return decodeTSRequest(der)
}

func readDER(r io.Reader) ([]byte, error) {
	header := make([]byte, 2)
	if _, err := io.ReadFull(r, header); err != nil {
		return nil, err
	}
	if header[0] != 0x30 {
		return nil, fmt.Errorf("unexpected DER sequence tag: 0x%x", header[0])
	}
	length := int(header[1])
	lenBytes := 0
	if header[1]&0x80 != 0 {
		lenBytes = int(header[1] & 0x7f)
		if lenBytes == 0 || lenBytes > 4 {
			return nil, fmt.Errorf("invalid DER length bytes: %d", lenBytes)
		}
		raw := make([]byte, lenBytes)
		if _, err := io.ReadFull(r, raw); err != nil {
			return nil, err
		}
		length = 0
		for _, b := range raw {
			length = length<<8 | int(b)
		}
		header = append(header, raw...)
	}
	body := make([]byte, length)
	if _, err := io.ReadFull(r, body); err != nil {
		return nil, err
	}
	return append(header, body...), nil
}

func peerPublicKey(conn *tls.Conn, style string) ([]byte, error) {
	state := conn.ConnectionState()
	if len(state.PeerCertificates) == 0 {
		return nil, fmt.Errorf("RDP TLS 证书为空")
	}
	cert := state.PeerCertificates[0]
	if style == "rsa-compat" {
		if pub, ok := cert.PublicKey.(*rsa.PublicKey); ok {
			return asn1.Marshal(rsaPublicKeyCompat{N: pub.N, E: pub.E})
		}
	}
	// MS-CSSP v5/v6 绑定哈希必须使用 X.509 SubjectPublicKey，而非完整 SPKI // Logan.Li
	return subjectPublicKey(cert)
}

func subjectPublicKey(cert *x509.Certificate) ([]byte, error) {
	if pub, ok := cert.PublicKey.(*rsa.PublicKey); ok {
		return asn1.Marshal(*pub)
	}
	var spki struct {
		Algorithm pkix.AlgorithmIdentifier
		PublicKey asn1.BitString
	}
	if _, err := asn1.Unmarshal(cert.RawSubjectPublicKeyInfo, &spki); err != nil {
		return nil, fmt.Errorf("解析 TLS SubjectPublicKey 失败: %w", err)
	}
	if len(spki.PublicKey.Bytes) == 0 {
		return nil, fmt.Errorf("TLS SubjectPublicKey 为空")
	}
	return spki.PublicKey.Bytes, nil
}

type rsaPublicKeyCompat struct {
	N *big.Int `asn1:"explicit,tag:0"`
	E int      `asn1:"explicit,tag:1"`
}

func splitDomainUser(domain, user string) (string, string) {
	if strings.Contains(user, "\\") {
		parts := strings.Split(user, "\\")
		if len(parts) >= 2 {
			return parts[0], parts[len(parts)-1]
		}
	}
	if strings.Contains(user, "/") {
		parts := strings.Split(user, "/")
		if len(parts) >= 2 {
			return parts[0], parts[len(parts)-1]
		}
	}
	return domain, user
}