// SPDX-License-Identifier: BSD-3-Clause

// Based on Go's crypto/tls library:
// https://github.com/golang/go/blob/master/src/crypto/tls/handshake_messages.go#L69
// https://cs.opensource.google/go/go/+/refs/tags/go1.18:src/crypto/tls/handshake_messages.go;l=69

package main

import (
	"crypto/md5"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"golang.org/x/crypto/cryptobyte"
	"log"
	"strconv"
	"strings"
)

// TLS versions
const (
	versionSSL30 = 0x0300 // 768
	versionTLS10 = 0x0301 // 769
	versionTLS11 = 0x0302 // 770
	versionTLS12 = 0x0303 // 771
	versionTLS13 = 0x0304 // 772
)

// TLS compression types
const (
	compressionNone uint8 = 0
)

// TLS extension numbers
const (
	extensionServerName              uint16 = 0
	extensionStatusRequest           uint16 = 5
	extensionSupportedGroups         uint16 = 10 // elliptic_curves in TLS versions prior to 1.3, see RFC 8446, Section 4.2.7
	extensionSupportedPointFormats   uint16 = 11
	extensionSignatureAlgorithms     uint16 = 13
	extensionALPN                    uint16 = 16
	extensionSCT                     uint16 = 18
	extensionPadding                 uint16 = 21
	extensionEncryptThenMac          uint16 = 22
	extensionExtendedMasterSecret    uint16 = 23
	extensionCompressCertificate     uint16 = 27
	extensionSessionTicket           uint16 = 35
	extensionPreSharedKey            uint16 = 41
	extensionEarlyData               uint16 = 42
	extensionSupportedVersions       uint16 = 43
	extensionCookie                  uint16 = 44
	extensionPSKModes                uint16 = 45
	extensionCertificateAuthorities  uint16 = 47
	extensionPostHandshakeAuth       uint16 = 49
	extensionSignatureAlgorithmsCert uint16 = 50
	extensionKeyShare                uint16 = 51
	extensionNextProtoNeg            uint16 = 13172 // not IANA assigned
	extensionApplicationSettings     uint16 = 17513 // not IANA assigned
	extensionRenegotiationInfo       uint16 = 0xff01
)

// TLS signaling cipher suite values
const (
	scsvRenegotiation uint16 = 0x00ff
)

// TLS CertificateStatusType (RFC 3546)
const (
	statusTypeOCSP uint8 = 1
)

// curveID is the type of a TLS identifier for an elliptic curve. See
// https://www.iana.org/assignments/tls-parameters/tls-parameters.xml#tls-parameters-8.
//
// In TLS 1.3, this type is called NamedGroup, but at this time this library
// only supports Elliptic Curve based groups. See RFC 8446, Section 4.2.7.
type curveID uint16

const (
	curveP256 curveID = 23
	curveP384 curveID = 24
	curveP521 curveID = 25
	x25519    curveID = 29
)

// TLS 1.3 Key Share. See RFC 8446, Section 4.2.8.
type keyShare struct {
	Group curveID   `json:"group"`
	Data  byteSlice `json:"data"`
}

// TLS 1.3 PSK Identity. Can be a Session Ticket, or a reference to a saved
// session. See RFC 8446, Section 4.2.11.
type pskIdentity struct {
	Identity            byteSlice `json:"identity"`
	ObfuscatedTicketAge uint32    `json:"obfuscated_ticket_age"`
}

// signatureScheme identifies a signature algorithm supported by TLS. See
// RFC 8446, Section 4.2.3.
type signatureScheme uint16

const (
	// RSASSA-PKCS1-v1_5 algorithms.
	PKCS1WithSHA256 signatureScheme = 0x0401
	PKCS1WithSHA384 signatureScheme = 0x0501
	PKCS1WithSHA512 signatureScheme = 0x0601

	// RSASSA-PSS algorithms with public key OID rsaEncryption.
	PSSWithSHA256 signatureScheme = 0x0804
	PSSWithSHA384 signatureScheme = 0x0805
	PSSWithSHA512 signatureScheme = 0x0806

	// ECDSA algorithms. Only constrained to a specific curve in TLS 1.3.
	ECDSAWithP256AndSHA256 signatureScheme = 0x0403
	ECDSAWithP384AndSHA384 signatureScheme = 0x0503
	ECDSAWithP521AndSHA512 signatureScheme = 0x0603

	// Legacy signature and hash algorithms for TLS 1.2.
	PKCS1WithSHA1 signatureScheme = 0x0201
	ECDSAWithSHA1 signatureScheme = 0x0203
)

// Check if this is a GREASE value (RFC 8701)
func isGREASE(value uint16) bool {
	// Values for: cipher suites, ALPN,
	// extensions, named groups, signature algorithms and TLS versions:
	greaseValues := [16]uint16{
		0x0a0a, 0x1a1a, 0x2a2a, 0x3a3a, 0x4a4a, 0x5a5a, 0x6a6a, 0x7a7a,
		0x8a8a, 0x9a9a, 0xaaaa, 0xbaba, 0xcaca, 0xdada, 0xeaea, 0xfafa,
	}
	// Values for PSK key exchange modes:
	//greasePSKModes := [8]uint8{
	//	0x0b, 0x2a, 0x49, 0x68, 0x87, 0xa6, 0xc5, 0xe4,
	//}
	// Run the check
	for _, greaseValue := range greaseValues {
		if value == greaseValue {
			return true
		}
	}
	return false
}

// readUint8LengthPrefixed acts like s.ReadUint8LengthPrefixed, but targets a
// []byte instead of a cryptobyte.String.
func readUint8LengthPrefixed(s *cryptobyte.String, out *[]byte) bool {
	return s.ReadUint8LengthPrefixed((*cryptobyte.String)(out))
}

// readUint16LengthPrefixed acts like s.ReadUint16LengthPrefixed, but targets a
// []byte instead of a cryptobyte.String.
func readUint16LengthPrefixed(s *cryptobyte.String, out *[]byte) bool {
	return s.ReadUint16LengthPrefixed((*cryptobyte.String)(out))
}

type byteSlice []byte

func (bs byteSlice) MarshalText() ([]byte, error) {
	// Used in json.Marshal() to convert []byte to hex string.
	return []byte(hex.EncodeToString(bs)), nil
}

type uint8Slice []uint8

func (numbers uint8Slice) MarshalJSON() ([]byte, error) {
	// Used in json.Marshal() to convert []uint8 to array.
	newSlice := make([]uint, len(numbers))
	for i, n := range numbers {
		newSlice[i] = uint(n)
	}
	return json.Marshal(newSlice)
}

type extensionData struct {
	Raw                          byteSlice         `json:"raw"`
	ServerName                   string            `json:"server_name,omitempty"`
	StatusType                   uint8             `json:"status_type,omitempty"`
	SupportedGroups              []curveID         `json:"supported_groups,omitempty"`
	SupportedPointFormats        uint8Slice        `json:"supported_point_formats,omitempty"`
	SupportedSignatureAlgorithms []signatureScheme `json:"supported_signature_algorithms,omitempty"`
	RenegotiationInfo            []byte            `json:"renegotiation_info,omitempty"`
	AlpnProtocols                []string          `json:"alpn_protocols,omitempty"`
	SupportedVersions            []uint16          `json:"supported_tls_versions,omitempty"`
	Cookie                       byteSlice         `json:"cookie,omitempty"`
	KeyShares                    []keyShare        `json:"key_shares,omitempty"`
	PskModes                     uint8Slice        `json:"psk_modes,omitempty"`
	PskIdentities                []pskIdentity     `json:"psk_identities,omitempty"`
	PskBinders                   []byteSlice       `json:"psk_binders,omitempty"`
	Length                       uint16            `json:"length,omitempty"` // padding
}

type extension struct {
	Code uint16        `json:"code"`
	Name string        `json:"name"`
	Data extensionData `json:"data"`
}

type highlights struct {
	//SupportedTLSVersions     []uint16
	GmtUnixTime                uint32 `json:"gmt_unix_time"` // first 4 bytes of client random
	SecureRenegotiationSupport bool   `json:"secure_renegotiation_support"`
	OcspStaplingSupport        bool   `json:"ocsp_stapling_support"`
	SctSupport                 bool   `json:"sct_support"`
	// Go's crypto/tls server does not support early data.
	EarlyData bool      `json:"-"` // don't include in JSON
	JA3       string    `json:"ja3"`
	JA3MD5    byteSlice `json:"ja3_md5"`
}

type clientHelloMsg struct {
	Raw                    byteSlice   `json:"raw"`
	RecordHeaderTLSVersion uint16      `json:"record_header_tls_version"` // TLSv1.0 (769)
	TLSVersion             uint16      `json:"client_tls_version"`        // TLSv1.2 (771)
	Random                 byteSlice   `json:"random"`
	SessionID              byteSlice   `json:"session_id"`
	CipherSuites           []uint16    `json:"cipher_suites"`
	CompressionMethods     uint8Slice  `json:"compression_methods"`
	Extensions             []extension `json:"extensions"`
	Highlights             highlights  `json:"highlights"`
}

func (m *clientHelloMsg) unmarshal(data []byte) bool {
	*m = clientHelloMsg{Raw: data}
	s := cryptobyte.String(data)

	var random []byte
	var sessionID []byte
	if !s.Skip(1) || !s.ReadUint16(&m.RecordHeaderTLSVersion) || !s.Skip(2) ||
		!s.Skip(4) || // message type and uint24 length field
		!s.ReadUint16(&m.TLSVersion) || !s.ReadBytes(&random, 32) ||
		!readUint8LengthPrefixed(&s, &sessionID) {
		return false
	}
	m.Random = random
	m.Highlights.GmtUnixTime = binary.BigEndian.Uint32(random[0:4])

	m.SessionID = sessionID

	var cipherSuites cryptobyte.String
	if !s.ReadUint16LengthPrefixed(&cipherSuites) {
		return false
	}
	m.CipherSuites = []uint16{}
	m.Highlights.SecureRenegotiationSupport = false
	for !cipherSuites.Empty() {
		var suite uint16
		if !cipherSuites.ReadUint16(&suite) {
			return false
		}
		if suite == scsvRenegotiation {
			m.Highlights.SecureRenegotiationSupport = true
		}
		m.CipherSuites = append(m.CipherSuites, suite)
	}

	var compressionMethods []uint8
	if !readUint8LengthPrefixed(&s, &compressionMethods) {
		return false
	}
	m.CompressionMethods = compressionMethods

	if s.Empty() {
		// ClientHello is optionally followed by extension data
		return true
	}

	var extensions cryptobyte.String
	if !s.ReadUint16LengthPrefixed(&extensions) || !s.Empty() {
		return false
	}

	for !extensions.Empty() {
		var extension extension
		var extData cryptobyte.String
		if !extensions.ReadUint16(&extension.Code) ||
			!extensions.ReadUint16LengthPrefixed(&extData) {
			return false
		}

		extension.Data.Raw = []byte(extData)

		switch extension.Code {
		case extensionServerName:
			// RFC 6066, Section 3
			extension.Name = "server_name" // Server Name Indication
			var nameList cryptobyte.String
			if !extData.ReadUint16LengthPrefixed(&nameList) || nameList.Empty() {
				return false
			}
			for !nameList.Empty() {
				var nameType uint8
				var serverName cryptobyte.String
				if !nameList.ReadUint8(&nameType) ||
					!nameList.ReadUint16LengthPrefixed(&serverName) ||
					serverName.Empty() {
					return false
				}
				if nameType != 0 {
					continue
				}
				if len(extension.Data.ServerName) != 0 {
					// Multiple names of the same name_type are prohibited.
					return false
				}
				sn := string(serverName)
				extension.Data.ServerName = sn
				// An SNI value may not include a trailing dot.
				if strings.HasSuffix(sn, ".") {
					return false
				}
			}
		case extensionNextProtoNeg:
			// draft-agl-tls-nextprotoneg-04
			extension.Name = "next_protocol_negotiation" // Next Protocol Negotiation
		case extensionStatusRequest:
			// RFC 4366, Section 3.6
			extension.Name = "status_request" // Certificate Status Request
			var ignored cryptobyte.String
			if !extData.ReadUint8(&extension.Data.StatusType) ||
				!extData.ReadUint16LengthPrefixed(&ignored) || // responder_id_list
				!extData.ReadUint16LengthPrefixed(&ignored) { // request_extensions
				return false
			}
			m.Highlights.OcspStaplingSupport = extension.Data.StatusType == statusTypeOCSP
		case extensionSupportedGroups:
			// RFC 4492, sections 5.1.1 and RFC 8446, Section 4.2.7
			extension.Name = "supported_groups" // Supported Groups
			//extension.Name = "elliptic_curves" // old name
			var curves cryptobyte.String
			if !extData.ReadUint16LengthPrefixed(&curves) || curves.Empty() {
				return false
			}
			extension.Data.SupportedGroups = make([]curveID, 0)
			for !curves.Empty() {
				var curve uint16
				if !curves.ReadUint16(&curve) {
					return false
				}
				extension.Data.SupportedGroups = append(
					extension.Data.SupportedGroups, curveID(curve))
			}
		case extensionSupportedPointFormats:
			// RFC 4492, Section 5.1.2
			extension.Name = "ec_point_formats" // Supported Point Formats
			var supportedPointFormats []uint8
			if !readUint8LengthPrefixed(&extData, &supportedPointFormats) ||
				len(supportedPointFormats) == 0 {
				return false
			}
			extension.Data.SupportedPointFormats = supportedPointFormats
		case extensionSessionTicket:
			// RFC 5077, Section 3.2
			extension.Name = "session_ticket" // Session Ticket
		case extensionSignatureAlgorithms:
			// RFC 5246, Section 7.4.1.4.1
			extension.Name = "signature_algorithms" // Signature Algorithms
			var sigAndAlgs cryptobyte.String
			if !extData.ReadUint16LengthPrefixed(&sigAndAlgs) || sigAndAlgs.Empty() {
				return false
			}
			for !sigAndAlgs.Empty() {
				var sigAndAlg uint16
				if !sigAndAlgs.ReadUint16(&sigAndAlg) {
					return false
				}
				extension.Data.SupportedSignatureAlgorithms = append(
					extension.Data.SupportedSignatureAlgorithms,
					signatureScheme(sigAndAlg))
			}
		case extensionSignatureAlgorithmsCert:
			// RFC 8446, Section 4.2.3
			extension.Name = "signature_algorithms_cert" // Signature Algorithms Cert
			var sigAndAlgs cryptobyte.String
			if !extData.ReadUint16LengthPrefixed(&sigAndAlgs) || sigAndAlgs.Empty() {
				return false
			}
			for !sigAndAlgs.Empty() {
				var sigAndAlg uint16
				if !sigAndAlgs.ReadUint16(&sigAndAlg) {
					return false
				}
				extension.Data.SupportedSignatureAlgorithms = append(
					extension.Data.SupportedSignatureAlgorithms,
					signatureScheme(sigAndAlg))
			}
		case extensionRenegotiationInfo:
			// RFC 5746, Section 3.2
			extension.Name = "renegotiation_info" // Renegotiation Indication
			if !readUint8LengthPrefixed(
				&extData,
				&extension.Data.RenegotiationInfo) {
				return false
			}
			m.Highlights.SecureRenegotiationSupport = true
		case extensionALPN:
			// RFC 7301, Section 3.1
			extension.Name = "application_layer_protocol_negotiation" // Application-Layer Protocol Negotiation
			var protoList cryptobyte.String
			if !extData.ReadUint16LengthPrefixed(&protoList) || protoList.Empty() {
				return false
			}
			for !protoList.Empty() {
				var proto cryptobyte.String
				if !protoList.ReadUint8LengthPrefixed(&proto) || proto.Empty() {
					return false
				}
				extension.Data.AlpnProtocols = append(extension.Data.AlpnProtocols, string(proto))
			}
		case extensionSCT:
			// RFC 6962, Section 3.3.1
			extension.Name = "signed_certificate_timestamp" // Signed Certificate Timestamp
			m.Highlights.SctSupport = true
		case extensionSupportedVersions:
			// RFC 8446, Section 4.2.1
			extension.Name = "supported_versions" // Supported Versions
			var versList cryptobyte.String
			if !extData.ReadUint8LengthPrefixed(&versList) || versList.Empty() {
				return false
			}
			for !versList.Empty() {
				var vers uint16
				if !versList.ReadUint16(&vers) {
					return false
				}
				extension.Data.SupportedVersions = append(extension.Data.SupportedVersions, vers)
			}
		case extensionCookie:
			// RFC 8446, Section 4.2.2
			extension.Name = "cookie" // Cookie
			var cookie []byte
			if !readUint16LengthPrefixed(&extData, &cookie) ||
				len(extension.Data.Cookie) == 0 {
				return false
			}
			extension.Data.Cookie = cookie
		case extensionKeyShare:
			// RFC 8446, Section 4.2.8
			extension.Name = "key_share" // Key Share
			var clientShares cryptobyte.String
			if !extData.ReadUint16LengthPrefixed(&clientShares) {
				return false
			}
			for !clientShares.Empty() {
				var ks keyShare
				var ksData []byte
				if !clientShares.ReadUint16((*uint16)(&ks.Group)) ||
					!readUint16LengthPrefixed(&clientShares, &ksData) ||
					len(ksData) == 0 {
					return false
				}
				ks.Data = ksData
				extension.Data.KeyShares = append(extension.Data.KeyShares, ks)
			}
		case extensionEarlyData:
			// RFC 8446, Section 4.2.10
			extension.Name = "early_data" // Early Data Indication
			m.Highlights.EarlyData = true
		case extensionPSKModes:
			// RFC 8446, Section 4.2.9
			extension.Name = "psk_key_exchange_modes" // Pre-Shared Key Exchange Modes
			var pskModes []uint8
			if !readUint8LengthPrefixed(&extData, &pskModes) {
				return false
			}
			extension.Data.PskModes = pskModes
		case extensionPreSharedKey:
			// RFC 8446, Section 4.2.11
			extension.Name = "pre_shared_key" // Pre-Shared Key
			if !extensions.Empty() {
				return false // pre_shared_key must be the last extension
			}
			var identities cryptobyte.String
			if !extData.ReadUint16LengthPrefixed(&identities) || identities.Empty() {
				return false
			}
			for !identities.Empty() {
				var psk pskIdentity
				var identity []byte
				if !readUint16LengthPrefixed(&identities, &identity) ||
					!identities.ReadUint32(&psk.ObfuscatedTicketAge) ||
					len(identity) == 0 {
					return false
				}
				psk.Identity = identity
				extension.Data.PskIdentities = append(extension.Data.PskIdentities, psk)
			}
			var binders cryptobyte.String
			if !extData.ReadUint16LengthPrefixed(&binders) || binders.Empty() {
				return false
			}
			for !binders.Empty() {
				var binder []byte
				if !readUint8LengthPrefixed(&binders, &binder) ||
					len(binder) == 0 {
					return false
				}
				extension.Data.PskBinders = append(extension.Data.PskBinders, binder)
			}
		case extensionPadding:
			// RFC 7685
			extension.Name = "padding" // Padding
			extension.Data.Length = uint16(len(extData))
		case extensionEncryptThenMac:
			// RFC 7366
			extension.Name = "encrypt_then_mac" // Encrypt-then-MAC"
		case extensionExtendedMasterSecret:
			// RFC 7627
			extension.Name = "extended_master_secret" // Extended Master Secret
		case extensionPostHandshakeAuth:
			// RFC 8446, Section 4.2.6
			extension.Name = "post_handshake_auth" // Post-Handshake Client Authentication
		case extensionCompressCertificate:
			// RFC 8879
			extension.Name = "compress_certificate" // Certificate Compression
			// TODO: parse - https://www.rfc-editor.org/rfc/rfc8879.html
		case extensionApplicationSettings:
			// draft-vvv-tls-alps
			// Application-Layer Protocol Settings
			extension.Name = " application_settings"
			// TODO: parse - https://datatracker.ietf.org/doc/html/draft-vvv-tls-alps
		default:
			// Check if this is a GREASE extension (RFC 8701)
			if isGREASE(extension.Code) {
				extension.Name = "GREASE"
			} else {
				log.Println("Unknown extension:", extension.Code)
			}
		}

		m.Extensions = append(m.Extensions, extension)

		if !extData.Empty() &&
			extension.Name != "GREASE" &&
			extension.Code != extensionPadding {
			log.Printf("Extension %d data not read.\n", extension.Code)
		}
	}

	// JA3
	var ja3 strings.Builder
	var supportedGroups []curveID
	var supportedPointFormats uint8Slice
	ja3.WriteString(strconv.FormatUint(uint64(m.TLSVersion), 10) + ",")
	for i, cs := range m.CipherSuites {
		if !isGREASE(cs) { // ignore GREASE values
			ja3.WriteString(strconv.FormatUint(uint64(cs), 10))
			if i+1 != len(m.CipherSuites) {
				ja3.WriteString("-")
			}
		}
	}
	ja3.WriteString(",")
	for i, e := range m.Extensions {
		if !isGREASE(e.Code) { // ignore GREASE values
			ja3.WriteString(strconv.FormatUint(uint64(e.Code), 10))
			if i+1 != len(m.Extensions) {
				ja3.WriteString("-")
			}
			if e.Code == extensionSupportedGroups {
				supportedGroups = e.Data.SupportedGroups
			} else if e.Code == extensionSupportedPointFormats {
				supportedPointFormats = e.Data.SupportedPointFormats
			}
		}
	}
	ja3.WriteString(",")
	for i, g := range supportedGroups {
		if !isGREASE(uint16(g)) { // ignore GREASE values
			ja3.WriteString(strconv.FormatUint(uint64(g), 10))
			if i+1 != len(supportedGroups) {
				ja3.WriteString("-")
			}
		}
	}
	ja3.WriteString(",")
	for i, pf := range supportedPointFormats {
		if !isGREASE(uint16(pf)) { // ignore GREASE values
			ja3.WriteString(strconv.FormatUint(uint64(pf), 10))
			if i+1 != len(supportedPointFormats) {
				ja3.WriteString("-")
			}
		}
	}
	m.Highlights.JA3 = ja3.String()
	hash := md5.Sum([]byte(m.Highlights.JA3))
	m.Highlights.JA3MD5 = hash[:]

	return true
}