ja4sentinel/internal/tlsparse/parser.go

// Package tlsparse provides TLS ClientHello extraction from captured packets
package tlsparse

import (
	"encoding/binary"
	"fmt"
	"sync"
	"time"

	"ja4sentinel/api"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
)

// ConnectionState represents the state of a TCP connection for TLS parsing
type ConnectionState int

const (
	// NEW: Observed SYN from client on a monitored port
	NEW ConnectionState = iota
	// WAIT_CLIENT_HELLO: Accumulating segments for complete ClientHello
	WAIT_CLIENT_HELLO
	// JA4_DONE: JA4 computed and logged, stop tracking this flow
	JA4_DONE
)

// ConnectionFlow tracks a single TCP flow for TLS handshake extraction
type ConnectionFlow struct {
	State       ConnectionState
	CreatedAt   time.Time
	LastSeen    time.Time
	SrcIP       string
	SrcPort     uint16
	DstIP       string
	DstPort     uint16
	IPMeta      api.IPMeta
	TCPMeta     api.TCPMeta
	HelloBuffer []byte
}

// ParserImpl implements the api.Parser interface for TLS parsing
type ParserImpl struct {
	mu           sync.RWMutex
	flows        map[string]*ConnectionFlow
	flowTimeout  time.Duration
	cleanupDone  chan struct{}
	cleanupClose chan struct{}
}

// NewParser creates a new TLS parser with connection state tracking
func NewParser() *ParserImpl {
	return NewParserWithTimeout(30 * time.Second)
}

// NewParserWithTimeout creates a new TLS parser with a custom flow timeout
func NewParserWithTimeout(timeout time.Duration) *ParserImpl {
	p := &ParserImpl{
		flows:        make(map[string]*ConnectionFlow),
		flowTimeout:  timeout,
		cleanupDone:  make(chan struct{}),
		cleanupClose: make(chan struct{}),
	}
	go p.cleanupLoop()
	return p
}

// flowKey generates a unique key for a TCP flow
func flowKey(srcIP string, srcPort uint16, dstIP string, dstPort uint16) string {
	return fmt.Sprintf("%s:%d->%s:%d", srcIP, srcPort, dstIP, dstPort)
}

// cleanupLoop periodically removes expired flows
func (p *ParserImpl) cleanupLoop() {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			p.cleanupExpiredFlows()
		case <-p.cleanupClose:
			close(p.cleanupDone)
			return
		}
	}
}

// cleanupExpiredFlows removes flows that have timed out or are done
func (p *ParserImpl) cleanupExpiredFlows() {
	p.mu.Lock()
	defer p.mu.Unlock()

	now := time.Now()
	for key, flow := range p.flows {
		if flow.State == JA4_DONE || now.Sub(flow.LastSeen) > p.flowTimeout {
			delete(p.flows, key)
		}
	}
}

// Process extracts TLS ClientHello from a raw packet
func (p *ParserImpl) Process(pkt api.RawPacket) (*api.TLSClientHello, error) {
	if len(pkt.Data) == 0 {
		return nil, fmt.Errorf("empty packet data")
	}

	// Parse packet layers
	packet := gopacket.NewPacket(pkt.Data, layers.LinkTypeEthernet, gopacket.Default)

	// Get IP layer
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		ipLayer = packet.Layer(layers.LayerTypeIPv6)
	}
	if ipLayer == nil {
		return nil, nil // Not an IP packet
	}

	// Get TCP layer
	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	if tcpLayer == nil {
		return nil, nil // Not a TCP packet
	}

	ip, ok := ipLayer.(gopacket.Layer)
	if !ok {
		return nil, fmt.Errorf("failed to cast IP layer")
	}

	tcp, ok := tcpLayer.(*layers.TCP)
	if !ok {
		return nil, fmt.Errorf("failed to cast TCP layer")
	}

	// Extract IP metadata
	ipMeta := extractIPMeta(ip)

	// Extract TCP metadata
	tcpMeta := extractTCPMeta(tcp)

	// Get source/destination info
	var srcIP, dstIP string
	var srcPort, dstPort uint16

	switch v := ip.(type) {
	case *layers.IPv4:
		srcIP = v.SrcIP.String()
		dstIP = v.DstIP.String()
	case *layers.IPv6:
		srcIP = v.SrcIP.String()
		dstIP = v.DstIP.String()
	}

	srcPort = uint16(tcp.SrcPort)
	dstPort = uint16(tcp.DstPort)

	// Get TCP payload (TLS data)
	payload := tcp.Payload
	if len(payload) == 0 {
		return nil, nil // No payload
	}

	// Get or create connection flow
	key := flowKey(srcIP, srcPort, dstIP, dstPort)
	flow := p.getOrCreateFlow(key, srcIP, srcPort, dstIP, dstPort, ipMeta, tcpMeta)

	// Check if flow is already done
	p.mu.RLock()
	isDone := flow.State == JA4_DONE
	p.mu.RUnlock()
	if isDone {
		return nil, nil // Already processed this flow
	}

	// Check if this is a TLS ClientHello
	clientHello, err := parseClientHello(payload)
	if err != nil {
		return nil, err
	}

	if clientHello != nil {
		// Found ClientHello, mark flow as done
		p.mu.Lock()
		flow.State = JA4_DONE
		flow.HelloBuffer = clientHello
		p.mu.Unlock()

		return &api.TLSClientHello{
			SrcIP:   srcIP,
			SrcPort: srcPort,
			DstIP:   dstIP,
			DstPort: dstPort,
			Payload: clientHello,
			IPMeta:  ipMeta,
			TCPMeta: tcpMeta,
		}, nil
	}

	// Check for fragmented ClientHello (accumulate segments)
	if flow.State == WAIT_CLIENT_HELLO || flow.State == NEW {
		p.mu.Lock()
		flow.State = WAIT_CLIENT_HELLO
		flow.HelloBuffer = append(flow.HelloBuffer, payload...)
		bufferCopy := make([]byte, len(flow.HelloBuffer))
		copy(bufferCopy, flow.HelloBuffer)
		p.mu.Unlock()

		// Try to parse accumulated buffer
		clientHello, err := parseClientHello(bufferCopy)
		if err != nil {
			return nil, err
		}
		if clientHello != nil {
			// Complete ClientHello found
			p.mu.Lock()
			flow.State = JA4_DONE
			p.mu.Unlock()

			return &api.TLSClientHello{
				SrcIP:   srcIP,
				SrcPort: srcPort,
				DstIP:   dstIP,
				DstPort: dstPort,
				Payload: clientHello,
				IPMeta:  ipMeta,
				TCPMeta: tcpMeta,
			}, nil
		}
	}

	return nil, nil // No ClientHello found yet
}

// getOrCreateFlow gets existing flow or creates a new one
func (p *ParserImpl) getOrCreateFlow(key string, srcIP string, srcPort uint16, dstIP string, dstPort uint16, ipMeta api.IPMeta, tcpMeta api.TCPMeta) *ConnectionFlow {
	p.mu.Lock()
	defer p.mu.Unlock()

	if flow, exists := p.flows[key]; exists {
		flow.LastSeen = time.Now()
		return flow
	}

	flow := &ConnectionFlow{
		State:       NEW,
		CreatedAt:   time.Now(),
		LastSeen:    time.Now(),
		SrcIP:       srcIP,
		SrcPort:     srcPort,
		DstIP:       dstIP,
		DstPort:     dstPort,
		IPMeta:      ipMeta,
		TCPMeta:     tcpMeta,
		HelloBuffer: make([]byte, 0),
	}
	p.flows[key] = flow
	return flow
}

// Close cleans up the parser and stops background goroutines
func (p *ParserImpl) Close() error {
	close(p.cleanupClose)
	<-p.cleanupDone
	return nil
}

// extractIPMeta extracts IP metadata from the IP layer
func extractIPMeta(ipLayer gopacket.Layer) api.IPMeta {
	meta := api.IPMeta{}

	switch v := ipLayer.(type) {
	case *layers.IPv4:
		meta.TTL = v.TTL
		meta.TotalLength = v.Length
		meta.IPID = v.Id
		meta.DF = v.Flags&layers.IPv4DontFragment != 0
	case *layers.IPv6:
		meta.TTL = v.HopLimit
		meta.TotalLength = uint16(v.Length)
		meta.IPID = 0  // IPv6 doesn't have IP ID
		meta.DF = true // IPv6 doesn't fragment at source
	}

	return meta
}

// extractTCPMeta extracts TCP metadata from the TCP layer
func extractTCPMeta(tcp *layers.TCP) api.TCPMeta {
	meta := api.TCPMeta{
		WindowSize: tcp.Window,
		Options:    make([]string, 0),
	}

	// Parse TCP options
	for _, opt := range tcp.Options {
		switch opt.OptionType {
		case layers.TCPOptionKindMSS:
			meta.MSS = binary.BigEndian.Uint16(opt.OptionData)
			meta.Options = append(meta.Options, "MSS")
		case layers.TCPOptionKindWindowScale:
			if len(opt.OptionData) > 0 {
				meta.WindowScale = opt.OptionData[0]
			}
			meta.Options = append(meta.Options, "WS")
		case layers.TCPOptionKindSACKPermitted:
			meta.Options = append(meta.Options, "SACK")
		case layers.TCPOptionKindTimestamps:
			meta.Options = append(meta.Options, "TS")
		default:
			meta.Options = append(meta.Options, fmt.Sprintf("OPT%d", opt.OptionType))
		}
	}

	return meta
}

// parseClientHello checks if the payload contains a TLS ClientHello and returns it
func parseClientHello(payload []byte) ([]byte, error) {
	if len(payload) < 5 {
		return nil, nil // Too short for TLS record
	}

	// TLS record layer: Content Type (1 byte), Version (2 bytes), Length (2 bytes)
	contentType := payload[0]

	// Check for TLS handshake (content type 22)
	if contentType != 22 {
		return nil, nil // Not a TLS handshake
	}

	// Check TLS version (TLS 1.0 = 0x0301, TLS 1.1 = 0x0302, TLS 1.2 = 0x0303, TLS 1.3 = 0x0304)
	version := binary.BigEndian.Uint16(payload[1:3])
	if version < 0x0301 || version > 0x0304 {
		return nil, nil // Unknown TLS version
	}

	recordLength := int(binary.BigEndian.Uint16(payload[3:5]))
	if len(payload) < 5+recordLength {
		return nil, nil // Incomplete TLS record
	}

	// Parse handshake protocol
	handshakePayload := payload[5 : 5+recordLength]
	if len(handshakePayload) < 1 {
		return nil, nil // Too short for handshake type
	}

	handshakeType := handshakePayload[0]

	// Check for ClientHello (handshake type 1)
	if handshakeType != 1 {
		return nil, nil // Not a ClientHello
	}

	// Return the full TLS record (header + payload) for fingerprinting
	return payload[:5+recordLength], nil
}

// IsClientHello checks if a payload contains a TLS ClientHello
func IsClientHello(payload []byte) bool {
	if len(payload) < 6 {
		return false
	}

	// TLS handshake record
	if payload[0] != 22 {
		return false
	}

	// Check version
	version := binary.BigEndian.Uint16(payload[1:3])
	if version < 0x0301 || version > 0x0304 {
		return false
	}

	recordLength := int(binary.BigEndian.Uint16(payload[3:5]))
	if len(payload) < 5+recordLength {
		return false
	}

	handshakePayload := payload[5 : 5+recordLength]
	if len(handshakePayload) < 1 {
		return false
	}

	// ClientHello type
	return handshakePayload[0] == 1
}