package tcp

import (
	"bytes"
	"encoding/binary"

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

// RDP does not start TLS right away: the client first sends an X.224
// Connection Request and the server answers with a Connection Confirm, each in
// a TPKT packet, to agree on a security protocol ([MS-RDPBCGR] 1.3.1.1). Only
// then does the TLS handshake begin, on the same connection.
//
// The tls analyzer used to give up on the first byte of such a stream, which
// left RDP invisible to TLS rules, and with it anything that puts a fake RDP
// negotiation in front of its TLS. So the negotiation is parsed as an optional
// prelude of the hello, and published as rdp_req / rdp_resp.

const (
	tpktVersion    = 0x03
	tpktHeaderSize = 4 // version (1 byte) + reserved (1 byte) + length (2 bytes)
	// x224FixedSize is length indicator (1 byte) + code (1 byte) +
	// DST-REF (2 bytes) + SRC-REF (2 bytes) + class option (1 byte).
	x224FixedSize = 7

	x224CodeConnectionRequest = 0xe0
	x224CodeConnectionConfirm = 0xd0

	rdpNegReq     = 0x01
	rdpNegRsp     = 0x02
	rdpNegFailure = 0x03
	rdpNegSize    = 8
)

var rdpCookiePrefix = []byte("Cookie: ")

// readPrelude consumes the X.224 connection PDU that precedes the hello on an
// RDP connection, if there is one.
func (h *tlsHello) readPrelude() utils.LSMAction {
	b, ok := h.rec.GetByte(false)
	if !ok {
		return utils.LSMActionPause
	}
	if b == internal.RecordTypeHandshake {
		return utils.LSMActionNext
	}
	if b != tpktVersion {
		return utils.LSMActionCancel
	}
	header, ok := h.rec.Get(tpktHeaderSize, false)
	if !ok {
		return utils.LSMActionPause
	}
	n := int(binary.BigEndian.Uint16(header[2:4]))
	if header[1] != 0 || n < tpktHeaderSize+x224FixedSize {
		return utils.LSMActionCancel
	}
	pdu, ok := h.rec.Get(n, true)
	if !ok {
		return utils.LSMActionPause
	}
	m := parseX224Connection(pdu[tpktHeaderSize:], h.x224Code)
	if m == nil {
		return utils.LSMActionCancel
	}
	h.rdp = m
	h.rdpUpdated = true
	return utils.LSMActionNext
}

// parseX224Connection parses an X.224 Connection Request or Confirm TPDU
// carrying the RDP negotiation, or returns nil if tpdu is not one.
func parseX224Connection(tpdu []byte, code byte) analyzer.PropMap {
	// The length indicator does not count itself. The low nibble of the code
	// is the credit, and the high nibble of the class option the class, which
	// is always 0 for RDP.
	if int(tpdu[0]) != len(tpdu)-1 || tpdu[1]&0xf0 != code || tpdu[6]&0xf0 != 0 {
		return nil
	}
	rest := tpdu[x224FixedSize:]
	m := analyzer.PropMap{}
	if code == x224CodeConnectionRequest {
		// Both the cookie ("Cookie: mstshash=<user>") and the load balancer
		// routing token ("Cookie: msts=<token>") are terminated by CRLF.
		if bytes.HasPrefix(rest, rdpCookiePrefix) {
			i := bytes.Index(rest, []byte("\r\n"))
			if i < 0 {
				return nil
			}
			m["cookie"] = string(rest[len(rdpCookiePrefix):i])
			rest = rest[i+2:]
		}
		// A client that sends no RDP_NEG_REQ can only do standard RDP
		// security (PROTOCOL_RDP).
		m["protocols"] = 0
	} else {
		m["protocol"] = 0
	}
	if len(rest) == 0 {
		return m
	}
	// Anything after the negotiation structure, such as the correlation info
	// a client may append, is ignored.
	if len(rest) < rdpNegSize || binary.LittleEndian.Uint16(rest[2:4]) != rdpNegSize {
		return nil
	}
	v := int(binary.LittleEndian.Uint32(rest[4:8]))
	switch {
	case code == x224CodeConnectionRequest && rest[0] == rdpNegReq:
		m["protocols"] = v
	case code == x224CodeConnectionConfirm && rest[0] == rdpNegRsp:
		m["protocol"] = v
	case code == x224CodeConnectionConfirm && rest[0] == rdpNegFailure:
		delete(m, "protocol")
		m["failure"] = v
	default:
		return nil
	}
	return m
}