yoloyolo/classify/classifier.go

package classify

import (
	"github.com/gopacket/gopacket"
	"github.com/gopacket/gopacket/layers"
)

// Classifier classifies packets against a rule table.
type Classifier struct {
	rules []Rule
}

// NewClassifier creates a classifier with the given rules.
func NewClassifier(rules []Rule) *Classifier {
	return &Classifier{rules: rules}
}

// Classify inspects a gopacket.Packet and returns a ClassifiedPacket.
// It checks ICMP first, then extracts TCP/UDP port info and matches against rules.
// Per D-03: returns ClassUnknown if no rule matches.
func (c *Classifier) Classify(pkt gopacket.Packet) ClassifiedPacket {
	result := ClassifiedPacket{
		Class:  ClassUnknown,
		Length: len(pkt.Data()),
	}

	// Check ICMP
	if pkt.Layer(layers.LayerTypeICMPv4) != nil || pkt.Layer(layers.LayerTypeICMPv6) != nil {
		result.Protocol = "icmp"
		// Match against rules
		for _, rule := range c.rules {
			if rule.Protocol == "icmp" {
				result.Class = rule.Class
				return result
			}
		}
		return result
	}

	// Check TCP
	if tcpLayer := pkt.Layer(layers.LayerTypeTCP); tcpLayer != nil {
		tcp := tcpLayer.(*layers.TCP)
		result.Protocol = "tcp"
		result.SrcPort = uint16(tcp.SrcPort)
		result.DstPort = uint16(tcp.DstPort)
		for _, rule := range c.rules {
			if rule.Protocol == "tcp" && (rule.DstPort == 0 || rule.DstPort == result.DstPort) {
				result.Class = rule.Class
				return result
			}
		}
		return result
	}

	// Check UDP
	if udpLayer := pkt.Layer(layers.LayerTypeUDP); udpLayer != nil {
		udp := udpLayer.(*layers.UDP)
		result.Protocol = "udp"
		result.SrcPort = uint16(udp.SrcPort)
		result.DstPort = uint16(udp.DstPort)
		for _, rule := range c.rules {
			if rule.Protocol == "udp" && (rule.DstPort == 0 || rule.DstPort == result.DstPort) {
				result.Class = rule.Class
				return result
			}
		}
		return result
	}

	// No recognized transport layer -> ClassUnknown per D-03
	return result
}