package rdp

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"time"

	"ezops/pkg/rdp/credssp"
	"ezops/pkg/rdp/transport"
)

// ProbeHandshake 执行 TCP + TPKT/X.224 + RDP 安全协议协商 // Logan.Li
func ProbeHandshake(ctx context.Context, cfg Config) (HandshakeInfo, error) {
	port := cfg.Port
	if port == 0 {
		port = 3389
	}
	address := fmt.Sprintf("%s:%d", cfg.Host, port)
	start := time.Now()

	var d net.Dialer
	conn, err := d.DialContext(ctx, "tcp", address)
	if err != nil {
		return HandshakeInfo{Address: address, Duration: time.Since(start)}, err
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(15 * time.Second))

	if err := transport.WriteTPKT(conn, transport.BuildConnectionRequest()); err != nil {
		return HandshakeInfo{Address: address, Duration: time.Since(start)}, err
	}
	payload, err := transport.ReadTPKT(conn)
	if err != nil {
		return HandshakeInfo{Address: address, Duration: time.Since(start)}, err
	}
	result, err := transport.ParseConnectionConfirm(payload)
	if err != nil {
		return HandshakeInfo{Address: address, Duration: time.Since(start)}, err
	}
	return HandshakeInfo{
		Address:          address,
		SelectedProtocol: result.ProtocolName(),
		Duration:         time.Since(start),
	}, nil
}

// ProbeTLSHandshake 执行 X.224 协商后升级 TLS，用于验证 Windows 10/Server NLA 前置链路 // Logan.Li
func ProbeTLSHandshake(ctx context.Context, cfg Config) (HandshakeInfo, error) {
	port := cfg.Port
	if port == 0 {
		port = 3389
	}
	address := fmt.Sprintf("%s:%d", cfg.Host, port)
	start := time.Now()

	var d net.Dialer
	conn, err := d.DialContext(ctx, "tcp", address)
	if err != nil {
		return HandshakeInfo{Address: address, Duration: time.Since(start)}, err
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(15 * time.Second))

	if err := transport.WriteTPKT(conn, transport.BuildConnectionRequestWithCookie(cfg.Username)); err != nil {
		return HandshakeInfo{Address: address, Duration: time.Since(start)}, err
	}
	payload, err := transport.ReadTPKT(conn)
	if err != nil {
		return HandshakeInfo{Address: address, Duration: time.Since(start)}, err
	}
	result, err := transport.ParseConnectionConfirm(payload)
	if err != nil {
		return HandshakeInfo{Address: address, Duration: time.Since(start)}, err
	}
	if result.SelectedProtocol != transport.ProtocolTLS && result.SelectedProtocol != transport.ProtocolCredSSP {
		return HandshakeInfo{Address: address, SelectedProtocol: result.ProtocolName(), Duration: time.Since(start)}, nil
	}
	tlsConn, err := transport.UpgradeTLS(conn, cfg.Host, 15*time.Second)
	if err != nil {
		return HandshakeInfo{Address: address, SelectedProtocol: result.ProtocolName(), Duration: time.Since(start)}, err
	}
	state := tlsConn.ConnectionState()
	certSubject := ""
	if len(state.PeerCertificates) > 0 {
		certSubject = state.PeerCertificates[0].Subject.String()
	}
	_ = tlsConn.SetReadDeadline(time.Now().Add(300 * time.Millisecond))
	_, _ = io.ReadAll(tlsConn)
	return HandshakeInfo{
		Address:          address,
		SelectedProtocol: result.ProtocolName(),
		Duration:         time.Since(start),
		TLSCipher:        transport.TLSCipherSuiteName(state.CipherSuite),
		CertSubject:      certSubject,
	}, nil
}

// ProbeCredSSP 执行到 CredSSP/NTLMv2 登录完成，用于验证 Windows 10/Server NLA 凭据 // Logan.Li
func ProbeCredSSP(ctx context.Context, cfg Config) (HandshakeInfo, error) {
	port := cfg.Port
	if port == 0 {
		port = 3389
	}
	address := fmt.Sprintf("%s:%d", cfg.Host, port)
	start := time.Now()

	strategies := []credssp.Options{
		{Version: 6, PubKeyStyle: "binding"},
		{Version: 6, PubKeyStyle: "binding", UseCBT: true},
		{Version: 6, PubKeyStyle: "binding", ForceMIC: true},
		{Version: 2, PubKeyStyle: "raw"},
	}
	var lastErr error
	for _, opts := range strategies {
		info, err := probeCredSSPOnce(ctx, cfg, address, opts)
		if err == nil {
			info.Duration = time.Since(start)
			return info, nil
		}
		lastErr = err
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("CredSSP 认证失败")
	}
	return HandshakeInfo{Address: address, Duration: time.Since(start)}, lastErr
}

func probeCredSSPOnce(ctx context.Context, cfg Config, address string, opts credssp.Options) (HandshakeInfo, error) {
	tlsConn, result, err := dialCredSSPOnce(ctx, cfg, address, opts)
	if err != nil {
		return HandshakeInfo{Address: address}, err
	}
	defer tlsConn.Close()
	state := tlsConn.ConnectionState()
	certSubject := ""
	if len(state.PeerCertificates) > 0 {
		certSubject = state.PeerCertificates[0].Subject.String()
	}
	return HandshakeInfo{
		Address:          address,
		SelectedProtocol: result.ProtocolName(),
		TLSCipher:        transport.TLSCipherSuiteName(state.CipherSuite),
		CertSubject:      certSubject,
	}, nil
}

func dialCredSSPOnce(ctx context.Context, cfg Config, address string, opts credssp.Options) (*tls.Conn, transport.NegotiationResult, error) {
	var d net.Dialer
	conn, err := d.DialContext(ctx, "tcp", address)
	if err != nil {
		return nil, transport.NegotiationResult{}, err
	}
	_ = conn.SetDeadline(time.Now().Add(30 * time.Second))

	if err := transport.WriteTPKT(conn, transport.BuildConnectionRequestWithCookie(cfg.Username)); err != nil {
		_ = conn.Close()
		return nil, transport.NegotiationResult{}, err
	}
	payload, err := transport.ReadTPKT(conn)
	if err != nil {
		_ = conn.Close()
		return nil, transport.NegotiationResult{}, err
	}
	result, err := transport.ParseConnectionConfirm(payload)
	if err != nil {
		_ = conn.Close()
		return nil, transport.NegotiationResult{}, err
	}
	if result.SelectedProtocol != transport.ProtocolCredSSP {
		_ = conn.Close()
		return nil, result, fmt.Errorf("目标未选择 CredSSP/NLA: %s", result.ProtocolName())
	}
	tlsConn, err := transport.UpgradeTLS(conn, cfg.Host, 30*time.Second)
	if err != nil {
		_ = conn.Close()
		return nil, result, err
	}
	if err := credssp.AuthenticateWithOptions(tlsConn, cfg.Domain, cfg.Username, cfg.Password, opts); err != nil {
		_ = tlsConn.Close()
		return nil, result, err
	}
	_ = tlsConn.SetDeadline(time.Time{})
	return tlsConn, result, nil
}