package protocol

import (
	"slices"
	"strings"

	"golang.org/x/crypto/cryptobyte"

	"git.maze.io/go/dpi"
)

func init() {
	registerTLS()
}

func registerTLS() {
	Register(Both, "\x16\x03\x00", detectTLS) // SSL 3.0
	Register(Both, "\x16\x03\x01", detectTLS) // TLS 1.0
	Register(Both, "\x16\x03\x02", detectTLS) // TLS 1.1
	Register(Both, "\x16\x03\x03", detectTLS) // TLS 1.2
}

const (
	tlsRecordTypeHandshake uint8 = 22
	tlsTypeClientHello     uint8 = 1
	tlsTypeServerHello     uint8 = 2
)

func detectTLS(dir Direction, data []byte, _, _ int) (proto *Protocol, confidence float64) {
	stream := cryptobyte.String(data)

	// A TLS packet always has a content type (1 byte), version (2 bytes) and length (2 bytes).
	if len(stream) < 5 {
		return nil, 0
	}

	// Check for TLS Handshake (type 22)
	var header struct {
		Type    uint8
		Version uint16
		Length  uint32
	}
	if !stream.ReadUint8(&header.Type) || header.Type != 0x16 {
		return nil, 0
	}
	if !stream.ReadUint16(&header.Version) {
		return nil, 0
	}
	if !stream.ReadUint24(&header.Length) {
		return nil, 0
	}

	// Initial confidence
	confidence = 0.5

	// Detected SSL/TLS tlsVersion
	var tlsVersion dpi.TLSVersion

	// Attempt to decode the full TLS Client Hello handshake
	if tlsVersion == 0 {
		if record, err := dpi.DecodeTLSRecord(data); err == nil && record.Type == tlsRecordTypeHandshake && len(record.Data) > 1 {
			tlsVersion = record.Version
			confidence += .15

			switch record.Data[0] {
			case tlsTypeClientHello: // TLS ClientHello
				if hello, err := dpi.DecodeTLSClientHello(record.Data); err == nil {
					tlsVersion = hello.Version
					confidence += .3

					slices.SortStableFunc(hello.ALPNProtocols, func(a, b string) int {
						return strings.Compare(b, a)
					})

					for _, id := range hello.ALPNProtocols {
						if proto = ALPNProtocol[id]; proto != nil {
							return
						}
					}
				}

			case tlsTypeServerHello: // TLS ServerHello
				if hello, err := dpi.DecodeTLSServerHello(record.Data); err == nil {
					tlsVersion = hello.Version
					confidence += .3

					slices.SortStableFunc(hello.ALPNProtocols, func(a, b string) int {
						return strings.Compare(b, a)
					})

					for _, id := range hello.ALPNProtocols {
						if proto = ALPNProtocol[id]; proto != nil {
							return
						}
					}
				}
			}
		}
	}

	// Attempt to decode at least the handshake protocol and version.
	if tlsVersion == 0 && !Strict {
		var handshakeType uint8
		if stream.ReadUint8(&handshakeType) && (handshakeType == 1 || handshakeType == 2) {
			var (
				length      uint32
				versionWord uint16
			)
			if stream.ReadUint24(&length) && stream.ReadUint16(&versionWord) {
				tlsVersion = dpi.TLSVersion(versionWord)
				confidence += .25
			}
		}
	}

	// Fall back to the version in the TLS record header, this is less accurate
	if tlsVersion == 0 && !Strict {
		tlsVersion = dpi.TLSVersion(header.Version)
	}

	// We're "multi protocol", in that SSL is its own protocol
	if tlsVersion == dpi.VersionSSL30 {
		return &Protocol{
			Type:    TypeSSL,
			Version: Version{Major: 3, Minor: 0, Patch: -1},
		}, confidence
	} else if tlsVersion >= dpi.VersionTLS10 && tlsVersion <= dpi.VersionTLS13 {
		return &Protocol{
			Type:    TypeTLS,
			Version: Version{Major: 1, Minor: int(uint8(tlsVersion) - 1), Patch: -1},
		}, confidence
	} else if tlsVersion >= dpi.VersionTLS13Draft && tlsVersion <= dpi.VersionTLS13Draft23 {
		return &Protocol{
			Type:    TypeTLS,
			Version: Version{Major: 1, Minor: 3, Patch: -1},
		}, confidence
	}
	return nil, 0
}

// ALPNProtocol is a map of TLS Application-Layer Protocol Negotiation (ALPN) Protocol identifier to [Protocol].
//
// See https://www.iana.org/assignments/tls-extensiontype-values/tls-extensiontype-values.xhtml#alpn-protocol-ids\
var ALPNProtocol = map[string]*Protocol{
	"acme-tls/1":         {Type: TypeACME, Encapsulation: TypeTLS, Version: Version{Major: 1}},                        // ACME-TLS/1
	"co":                 {Type: TypeCoAP, Encapsulation: TypeTLS},                                                    // CoAP (over DTLS)
	"coap":               {Type: TypeCoAP, Encapsulation: TypeTLS},                                                    // CoAP (over TLS)
	"c-webrtc":           {Type: TypeWebRTC, Encapsulation: TypeTLS},                                                  // Confidential WebRTC Media and Data
	"dot":                {Type: TypeDNS, Encapsulation: TypeTLS},                                                     // DNS-over-TLS
	"ftp":                {Type: TypeFTP, Encapsulation: TypeTLS},                                                     // FTP
	"http/0.9":           {Type: TypeHTTP, Encapsulation: TypeTLS, Version: Version{Major: 0, Minor: 9, Patch: -1}},   // HTTP/0.9
	"http/1.0":           {Type: TypeHTTP, Encapsulation: TypeTLS, Version: Version{Major: 1, Minor: 0, Patch: -1}},   // HTTP/1.0
	"http/1.1":           {Type: TypeHTTP, Encapsulation: TypeTLS, Version: Version{Major: 1, Minor: 1, Patch: -1}},   // HTTP/1.1
	"h2":                 {Type: TypeHTTP, Encapsulation: TypeTLS, Version: Version{Major: 2, Minor: -1, Patch: -1}},  // HTTP/2 (over TLS)
	"h2c":                {Type: TypeHTTP, Encapsulation: TypeTLS, Version: Version{Major: 2, Minor: -1, Patch: -1}},  // HTTP/2 (over TCP)
	"h3":                 {Type: TypeHTTP, Encapsulation: TypeTLS, Version: Version{Major: 3, Minor: -1, Patch: -1}},  // HTTP/3
	"irc":                {Type: TypeIRC, Encapsulation: TypeTLS},                                                     // IRC
	"imap":               {Type: TypeIMAP, Encapsulation: TypeTLS},                                                    // IMAP
	"managesieve":        {Type: TypeManageSieve, Encapsulation: TypeTLS},                                             // ManageSieve
	"mqtt":               {Type: TypeMQTT, Encapsulation: TypeTLS},                                                    // MQTT
	"nntp":               {Type: TypeNNTP, Encapsulation: TypeTLS},                                                    // NNTP (reading)
	"nnsp":               {Type: TypeNNTP, Encapsulation: TypeTLS},                                                    // NNTP (transit)
	"postgresql":         {Type: TypePostgreSQL, Encapsulation: TypeTLS},                                              // PostgreSQL
	"pop3":               {Type: TypePOP3, Encapsulation: TypeTLS},                                                    // POP3
	"radius/1.0":         {Type: TypeRADIUS, Encapsulation: TypeTLS, Version: Version{Major: 1, Minor: 0, Patch: -1}}, // RADIUS/1.0
	"radius/1.1":         {Type: TypeRADIUS, Encapsulation: TypeTLS, Version: Version{Major: 1, Minor: 1, Patch: -1}}, // RADIUS/1.1
	"smb":                {Type: TypeSMB, Encapsulation: TypeTLS, Version: Version{Major: 2, Minor: -1, Patch: -1}},   // SMB2
	"stun.nat-discovery": {Type: TypeSTUN, Encapsulation: TypeTLS},                                                    // NAT discovery using Session Traversal Utilities for NAT (STUN)
	"stun.turn":          {Type: TypeSTUN, Encapsulation: TypeTLS},                                                    // Traversal Using Relays around NAT (TURN)
	"sunrpc":             {Type: TypeSunRPC, Encapsulation: TypeTLS},                                                  // SunRPC
	"webrtc":             {Type: TypeWebRTC, Encapsulation: TypeTLS},                                                  // WebRTC Media and Data
	"xmpp-client":        {Type: TypeXMPP, Encapsulation: TypeTLS},                                                    // XMPP jabber:client namespace
	"xmpp-server":        {Type: TypeXMPP, Encapsulation: TypeTLS},                                                    // XMPP jabber:server namespace
}