package protocol

import (
	"encoding/hex"
	"errors"
	"strings"
	"testing"
)

func TestDetectTLS(t *testing.T) {
	atomicFormats.Store([]format{})
	registerTLS()

	// A SSLv3 Client Hello
	sslV3ClientHello := testMustDecodeHexString(`
		                  16 03 00 00 51 01 00 00 4d 03
		00 50 42 b2 29 1f cf 52 a0 94 87 05 e7 0b 63 08
		12 a2 6c 59 f7 f5 72 2b 57 14 a7 07 95 cb ce e5
		e4 00 00 26 00 04 00 05 00 2f 00 33 00 32 00 0a
		fe ff 00 16 00 13 00 66 00 09 fe fe 00 15 00 12
		00 03 00 08 00 06 00 14 00 11 01 00`)

	// A synthesized TLS 1.1 ClientHello
	tls11ClientHello := []byte{
		// --- Record Layer (5 bytes) ---
		0x16,       // Content Type: Handshake (22)
		0x03, 0x02, // Version: TLS 1.1 (Major=3, Minor=2)
		0x00, 0x45, // Length of the handshake message below (69 bytes)

		// --- Handshake Protocol: ClientHello (69 bytes) ---
		0x01,             // Handshake Type: ClientHello (1)
		0x00, 0x00, 0x41, // Length of the rest of the message (65 bytes)
		0x03, 0x02, // Client Version: TLS 1.1 (Major=3, Minor=2)

		// Random (32 bytes):
		// In TLS 1.1+, the entire 32 bytes are fully random.
		// The timestamp structure from SSLv3 is removed.
		0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88,
		0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00,
		0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88,
		0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00,

		// Session ID
		0x00, // Session ID Length: 0 (new session)

		// Cipher Suites
		0x00, 0x04, // Cipher Suites Length: 4 bytes (2 suites)
		0x00, 0x2F, // Cipher Suite: TLS_RSA_WITH_AES_128_CBC_SHA
		0x00, 0x35, // Cipher Suite: TLS_RSA_WITH_AES_256_CBC_SHA

		// Compression Methods
		0x01, // Compression Methods Length: 1 byte
		0x00, // Compression Method: NULL (0)

		// --- Extensions (20 bytes) ---
		0x00, 0x14, // Extensions Length: 20 bytes
		// Extension: Server Name Indication (SNI)
		0x00, 0x00, // Extension Type: server_name (0)
		0x00, 0x10, // Extension Length: 16 bytes
		0x00, 0x0E, // Server Name List Length: 14 bytes
		0x00,       // Server Name Type: host_name (0)
		0x00, 0x0B, // Server Name Length: 11 bytes
		// Server Name: "example.com"
		'e', 'x', 'a', 'm', 'p', 'l', 'e', '.', 'c', 'o', 'm',
	}

	// A synthesized partial TLS 1.1 ClientHello
	tls11ClientHelloPartial := []byte{
		// --- Record Layer (5 bytes) ---
		0x16,       // Content Type: Handshake (22)
		0x03, 0x02, // Version: TLS 1.1 (Major=3, Minor=2)
		0x00, 0x45, // Length of the handshake message below (69 bytes)

		// --- Handshake Protocol: ClientHello (69 bytes) ---
		0x01,             // Handshake Type: ClientHello (1)
		0x00, 0x00, 0x41, // Length of the rest of the message (65 bytes)
		0x03, 0x02, // Client Version: TLS 1.1 (Major=3, Minor=2)
	}

	// A synthesized TLSv1.2 ClientHello
	tls12ClientHello := []byte{
		// --- Record Layer (5 bytes) ---
		0x16,       // Content Type: Handshake (22)
		0x03, 0x03, // Version: TLS 1.2 (Major=3, Minor=3)
		0x00, 0x61, // Length of handshake message below (97 bytes)

		// --- Handshake Protocol: ClientHello (97 bytes) ---
		0x01,             // Handshake Type: ClientHello (1)
		0x00, 0x00, 0x5D, // Length of the rest of the message (93 bytes)
		0x03, 0x03, // Client Version: TLS 1.2 (Major=3, Minor=3)

		// Random (32 bytes)
		0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00, 0x11,
		0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99,
		0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00, 0x11,
		0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99,

		// Session ID
		0x00, // Session ID Length: 0 (new session)

		// Cipher Suites (using modern GCM suites)
		0x00, 0x04, // Cipher Suites Length: 4 bytes (2 suites)
		0xC0, 0x2F, // Cipher Suite: TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256
		0xC0, 0x30, // Cipher Suite: TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384

		// Compression Methods
		0x01, // Compression Methods Length: 1 byte
		0x00, // Compression Method: NULL (0)

		// --- Extensions (48 bytes) ---
		0x00, 0x30, // Extensions Length: 48 bytes
		// Extension: Server Name Indication (SNI) (20 bytes)
		0x00, 0x00, 0x00, 0x10, 0x00, 0x0E, 0x00, 0x00, 0x0B,
		'e', 'x', 'a', 'm', 'p', 'l', 'e', '.', 'c', 'o', 'm',
		// Extension: Signature Algorithms (10 bytes) - CRITICAL for TLS 1.2
		0x00, 0x0D, // Extension Type: signature_algorithms (13)
		0x00, 0x06, // Extension Length: 6 bytes
		0x00, 0x04, // Hash/Signature Algorithm List Length: 4 bytes
		0x04, 0x01, // Algorithm: rsa_pkcs1_sha256
		0x05, 0x01, // Algorithm: rsa_pkcs1_sha384
		// Extension: Application-Layer Protocol Negotiation (ALPN) (18 bytes)
		0x00, 0x10, // Extension Type: application_layer_protocol_negotiation (16)
		0x00, 0x0E, // Extension Length: 14 bytes
		0x00, 0x0C, // ALPN Extension Length: 12 bytes
		// ALPN Protocol: "h2" (HTTP/2)
		0x02, 'h', '2',
		// ALPN Protocol: "http/1.1"
		0x08, 'h', 't', 't', 'p', '/', '1', '.', '1',
	}

	// A valid TLS 1.3 Client Hello (captured from a real connection)
	tls13ClientHello := []byte{
		0x16,       // Content Type: Handshake (22)
		0x03, 0x01, // Version: TLS 1.0 (for compatibility in a 1.3 hello)
		0x01, 0x3a, // Length
		0x01, 0x00, 0x01, 0x36, 0x03, 0x03, 0xb1,
		0x40, 0xd3, 0xf1, 0x7d, 0xa3, 0xb8, 0x33, 0xac, 0xad, 0x21, 0x79, 0x9c,
		0xbe, 0x39, 0x96, 0x08, 0x49, 0x3b, 0x53, 0x75, 0xa0, 0x1b, 0xee, 0x6e,
		0x6a, 0xbe, 0x6c, 0x41, 0xdf, 0x6c, 0xf4, 0x20, 0xa4, 0xaa, 0x0c, 0xca,
		0xd4, 0x37, 0x76, 0x5f, 0x49, 0xc6, 0x06, 0x9b, 0xac, 0x90, 0x89, 0x76,
		0x1c, 0xc7, 0xc4, 0x12, 0xb4, 0x4a, 0xe0, 0x27, 0x72, 0x89, 0x97, 0x85,
		0x76, 0xf8, 0xc8, 0x83, 0x00, 0x62, 0x13, 0x03, 0x13, 0x02, 0x13, 0x01,
		0xcc, 0xa9, 0xcc, 0xa8, 0xcc, 0xaa, 0xc0, 0x30, 0xc0, 0x2c, 0xc0, 0x28,
		0xc0, 0x24, 0xc0, 0x14, 0xc0, 0x0a, 0x00, 0x9f, 0x00, 0x6b, 0x00, 0x39,
		0xff, 0x85, 0x00, 0xc4, 0x00, 0x88, 0x00, 0x81, 0x00, 0x9d, 0x00, 0x3d,
		0x00, 0x35, 0x00, 0xc0, 0x00, 0x84, 0xc0, 0x2f, 0xc0, 0x2b, 0xc0, 0x27,
		0xc0, 0x23, 0xc0, 0x13, 0xc0, 0x09, 0x00, 0x9e, 0x00, 0x67, 0x00, 0x33,
		0x00, 0xbe, 0x00, 0x45, 0x00, 0x9c, 0x00, 0x3c, 0x00, 0x2f, 0x00, 0xba,
		0x00, 0x41, 0xc0, 0x11, 0xc0, 0x07, 0x00, 0x05, 0x00, 0x04, 0xc0, 0x12,
		0xc0, 0x08, 0x00, 0x16, 0x00, 0x0a, 0x00, 0xff, 0x01, 0x00, 0x00, 0x8b,
		0x00, 0x2b, 0x00, 0x09, 0x08, 0x03, 0x04, 0x03, 0x03, 0x03, 0x02, 0x03,
		0x01, 0x00, 0x33, 0x00, 0x26, 0x00, 0x24, 0x00, 0x1d, 0x00, 0x20, 0x2c,
		0x4b, 0xaa, 0xb4, 0xb3, 0xc8, 0x93, 0xcd, 0x5c, 0x24, 0xb9, 0x9b, 0xd4,
		0x59, 0x04, 0xfe, 0x69, 0xaf, 0x68, 0xb9, 0xa6, 0x36, 0xbb, 0xab, 0x87,
		0xfa, 0x15, 0x59, 0xea, 0xdd, 0x38, 0x68, 0x00, 0x00, 0x00, 0x0e, 0x00,
		0x0c, 0x00, 0x00, 0x09, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x68, 0x6f, 0x73,
		0x74, 0x00, 0x0b, 0x00, 0x02, 0x01, 0x00, 0x00, 0x0a, 0x00, 0x0a, 0x00,
		0x08, 0x00, 0x1d, 0x00, 0x17, 0x00, 0x18, 0x00, 0x19, 0x00, 0x0d, 0x00,
		0x18, 0x00, 0x16, 0x08, 0x06, 0x06, 0x01, 0x06, 0x03, 0x08, 0x05, 0x05,
		0x01, 0x05, 0x03, 0x08, 0x04, 0x04, 0x01, 0x04, 0x03, 0x02, 0x01, 0x02,
		0x03, 0x00, 0x10, 0x00, 0x0e, 0x00, 0x0c, 0x02, 0x68, 0x32, 0x08, 0x68,
		0x74, 0x74, 0x70, 0x2f, 0x31, 0x2e, 0x31,
	}

	// Invalid data (Postgres client startup)
	pgClientStartup := []byte{
		0x00, 0x00, 0x00, 0x25, 0x00, 0x03, 0x00, 0x00,
	}

	defer func() { Strict = false }()
	for _, strict := range []bool{false, true} {
		Strict = strict

		name := "loose"
		if strict {
			name = "strict"
		}

		t.Run(name, func(t *testing.T) {

			t.Run("SSLv3 Client Hello", func(t *testing.T) {
				p, _, err := Detect(Client, sslV3ClientHello, 0, 0)
				if err != nil {
					t.Fatal(err)
					return
				}
				t.Logf("detected %s version %s", p.Name, p.Version)
				if p.Name != ProtocolSSL {
					t.Fatalf("expected ssl protocol, got %s", p.Name)
					return
				}
			})

			t.Run("TLS 1.1 Client Hello", func(t *testing.T) {
				p, _, err := Detect(Client, tls11ClientHello, 0, 0)
				if err != nil {
					t.Fatal(err)
					return
				}
				t.Logf("detected %s version %s", p.Name, p.Version)
				if p.Name != ProtocolTLS {
					t.Fatalf("expected tls protocol, got %s", p.Name)
					return
				}
			})

			t.Run("TLS 1.1 partial Client Hello", func(t *testing.T) {
				p, _, err := Detect(Client, tls11ClientHelloPartial, 0, 0)
				if strict {
					if !errors.Is(err, ErrUnknown) {
						t.Fatalf("expected unknown format, got error %T: %q", err, err)
					} else {
						t.Logf("error %q, as expected", err)
					}
				} else {
					if err != nil {
						t.Fatal(err)
						return
					}
					t.Logf("detected %s version %s", p.Name, p.Version)
					if p.Name != ProtocolTLS {
						t.Fatalf("expected tls protocol, got %s", p.Name)
						return
					}
				}
			})

			t.Run("TLS 1.2 Client Hello", func(t *testing.T) {
				p, _, err := Detect(Client, tls12ClientHello, 0, 0)
				if err != nil {
					t.Fatal(err)
					return
				}
				t.Logf("detected %s version %s", p.Name, p.Version)
				if p.Name != ProtocolTLS {
					t.Fatalf("expected tls protocol, got %s", p.Name)
					return
				}
			})

			t.Run("TLS 1.3 Client Hello", func(t *testing.T) {
				p, _, err := Detect(Client, tls13ClientHello, 0, 0)
				if err != nil {
					t.Fatal(err)
					return
				}
				t.Logf("detected %s version %s", p.Name, p.Version)
				if p.Name != ProtocolTLS {
					t.Fatalf("expected tls protocol, got %s", p.Name)
					return
				}
			})

			t.Run("Invalid PostgreSQL", func(t *testing.T) {
				_, _, err := Detect(Server, pgClientStartup, 0, 0)
				if !errors.Is(err, ErrUnknown) {
					t.Fatalf("expected unknown format, got error %T: %q", err, err)
				} else {
					t.Logf("error %q, as expected", err)
				}
			})
		})
	}
}

func testDecodeHexString(s string) ([]byte, error) {
	s = strings.TrimSpace(s)
	s = strings.ReplaceAll(s, " ", "")
	s = strings.ReplaceAll(s, "\n", "")
	s = strings.ReplaceAll(s, "\t", "")
	return hex.DecodeString(s)
}

func testMustDecodeHexString(s string) []byte {
	b, err := testDecodeHexString(s)
	if err != nil {
		panic(err)
	}
	return b
}