package tcp

import (
	"reflect"
	"testing"

	"github.com/apernet/OpenGFW/analyzer"
)

// rdpConnectionRequest is what mstsc sends first: a cookie with the user name
// and an RDP_NEG_REQ asking for PROTOCOL_SSL | PROTOCOL_HYBRID |
// PROTOCOL_HYBRID_EX.
var rdpConnectionRequest = []byte{
	0x03, 0x00, 0x00, 0x2c, // TPKT, length 44
	0x27, 0xe0, 0x00, 0x00, 0x00, 0x00, 0x00, // X.224 CR
	'C', 'o', 'o', 'k', 'i', 'e', ':', ' ', 'm', 's', 't', 's', 'h', 'a', 's', 'h', '=', 'e', 'l', 't', 'o', 'n', 's', '\r', '\n',
	0x01, 0x00, 0x08, 0x00, 0x0b, 0x00, 0x00, 0x00, // RDP_NEG_REQ
}

func rdpConnectionConfirm(negType byte, v byte) []byte {
	return []byte{
		0x03, 0x00, 0x00, 0x13, // TPKT, length 19
		0x0e, 0xd0, 0x00, 0x00, 0x12, 0x34, 0x00, // X.224 CC
		negType, 0x00, 0x08, 0x00, v, 0x00, 0x00, 0x00,
	}
}

var rdpTLSClientHelloProps = analyzer.PropMap{"cookie": "mstshash=eltons", "protocols": 11}

func TestTlsRDPClientHello(t *testing.T) {
	data := append(append([]byte{}, rdpConnectionRequest...), tlsExampleClientHello...)
	for name, segs := range map[string][][]byte{
		"one segment":  {data},
		"byte by byte": splitEvery(data, 1),
		"pdu, hello":   {rdpConnectionRequest, tlsExampleClientHello},
	} {
		props, _ := feedTLS(newTLSStream(nil), false, segs...)
		if got := props.Get("rdp_req"); !reflect.DeepEqual(got, rdpTLSClientHelloProps) {
			t.Errorf("%s: rdp_req = %v, want %v", name, got, rdpTLSClientHelloProps)
		}
		if got := props.Get("req"); !reflect.DeepEqual(got, tlsExampleClientHelloProps) {
			t.Errorf("%s: req = %v, want %v", name, got, tlsExampleClientHelloProps)
		}
	}
}

// The negotiation is published before any TLS shows up, so that rules can
// already see the stream is RDP.
func TestTlsRDPNegotiationReportedEarly(t *testing.T) {
	s := newTLSStream(nil)
	props, done := feedTLS(s, false, rdpConnectionRequest)
	if done {
		t.Fatal("stream done after connection request")
	}
	if got := props.Get("rdp_req"); !reflect.DeepEqual(got, rdpTLSClientHelloProps) {
		t.Errorf("rdp_req = %v, want %v", got, rdpTLSClientHelloProps)
	}
	props, done = feedTLS(s, true, rdpConnectionConfirm(0x02, 0x02))
	if done {
		t.Fatal("stream done after connection confirm")
	}
	if got, want := props.Get("rdp_resp"), (analyzer.PropMap{"protocol": 2}); !reflect.DeepEqual(got, want) {
		t.Errorf("rdp_resp = %v, want %v", got, want)
	}
}

func TestTlsRDPConnectionConfirm(t *testing.T) {
	for name, tc := range map[string]struct {
		pdu  []byte
		want analyzer.PropMap
	}{
		"ssl":        {rdpConnectionConfirm(0x02, 0x01), analyzer.PropMap{"protocol": 1}},
		"hybrid_ex":  {rdpConnectionConfirm(0x02, 0x08), analyzer.PropMap{"protocol": 8}},
		"failure":    {rdpConnectionConfirm(0x03, 0x05), analyzer.PropMap{"failure": 5}},
		"no neg rsp": {[]byte{0x03, 0x00, 0x00, 0x0b, 0x06, 0xd0, 0x00, 0x00, 0x12, 0x34, 0x00}, analyzer.PropMap{"protocol": 0}},
	} {
		props, _ := feedTLS(newTLSStream(nil), true, tc.pdu)
		if got := props.Get("rdp_resp"); !reflect.DeepEqual(got, tc.want) {
			t.Errorf("%s: rdp_resp = %v, want %v", name, got, tc.want)
		}
	}
}

// Standard RDP security has no TLS: the client goes on with MCS, which is not
// a TLS record, so the analyzer stops with only the negotiation published.
func TestTlsRDPWithoutTLS(t *testing.T) {
	mcs := []byte{0x03, 0x00, 0x01, 0xa0, 0x02, 0xf0, 0x80, 0x7f, 0x65}
	props, done := feedTLS(newTLSStream(nil), false, rdpConnectionRequest, mcs)
	if !done {
		t.Error("stream not done")
	}
	if props.Get("rdp_req") == nil || props.Get("req") != nil {
		t.Errorf("props = %v, want rdp_req only", props)
	}
}

// truncatedNegReq is rdpConnectionRequest with the last two bytes of the
// RDP_NEG_REQ cut off and the lengths adjusted to match.
func truncatedNegReq() []byte {
	pdu := append([]byte{}, rdpConnectionRequest[:len(rdpConnectionRequest)-2]...)
	pdu[3] -= 2
	pdu[4] -= 2
	return pdu
}

func TestTlsRDPRejectsMalformedPrelude(t *testing.T) {
	mutate := func(i int, b byte) []byte {
		pdu := append([]byte{}, rdpConnectionRequest...)
		pdu[i] = b
		return pdu
	}
	for name, pdu := range map[string][]byte{
		"reserved set":      mutate(1, 0x01),
		"wrong code":        mutate(5, 0xd0),
		"length mismatch":   mutate(4, 0x26),
		"class 2":           mutate(10, 0x20),
		"cookie w/o crlf":   mutate(len(rdpConnectionRequest)-10, 'x'),
		"bad neg type":      mutate(len(rdpConnectionRequest)-8, 0x02),
		"tpkt too short":    {0x03, 0x00, 0x00, 0x05, 0x00},
		"not tpkt nor tls":  {0x05, 0x01, 0x00},
		"truncated neg req": truncatedNegReq(),
	} {
		props, done := feedTLS(newTLSStream(nil), false, pdu)
		if !done {
			t.Errorf("%s: stream not done", name)
		}
		if len(props) != 0 {
			t.Errorf("%s: props = %v, want none", name, props)
		}
	}
}