Added Layer interface, deleted some code

Author: shadowy-pycoder

layers/arp.go

@@ -82,6 +82,10 @@ func (ap *ARPPacket) Parse(data []byte) error {
 	return nil
 }
 
+func (ap *ARPPacket) NextLayer() (string, []byte) {
+	return "", nil
+}
+
 func (ap *ARPPacket) ptype() string {
 	var proto string
 	switch ap.ProtocolType {

layers/dns.go

@@ -3,3 +3,15 @@ package layers
 // https://en.wikipedia.org/wiki/Domain_Name_System
 // port 53
 type DNSMessage struct{}
+
+func (d *DNSMessage) String() string {
+	return ""
+}
+
+func (d *DNSMessage) Parse(data []byte) error {
+	return nil
+}
+
+func (d *DNSMessage) NextLayer() (string, []byte) {
+	return "", nil
+}

layers/ethernet.go

@@ -14,10 +14,11 @@ type EthernetFrame struct {
 	DstMAC    net.HardwareAddr // MAC address of the destination device.
 	SrcMAC    net.HardwareAddr // MAC address of the source device.
 	EtherType uint16           // The protocol of the upper layer.
-	Payload   []byte
+	payload   []byte
 }
 
 func (ef *EthernetFrame) String() string {
+	ethType, _ := ef.NextLayer()
 	return fmt.Sprintf(`Ethernet Frame:
 - DstMAC: %s
 - SrcMAC: %s
@@ -26,11 +27,11 @@ func (ef *EthernetFrame) String() string {
 %s`,
 		ef.DstMAC,
 		ef.SrcMAC,
-		ef.NextLayer(),
+		ethType,
 		ef.EtherType,
-		len(ef.Payload),
-		ef.Payload,
-		hex.Dump(ef.Payload))
+		len(ef.payload),
+		ef.payload,
+		hex.Dump(ef.payload))
 }
 
 // Parse parses the given byte data into an Ethernet frame.
@@ -41,22 +42,22 @@ func (ef *EthernetFrame) Parse(data []byte) error {
 	ef.DstMAC = net.HardwareAddr(data[0:6])
 	ef.SrcMAC = net.HardwareAddr(data[6:12])
 	ef.EtherType = binary.BigEndian.Uint16(data[12:14])
-	ef.Payload = data[ethernetHeaderSize:]
+	ef.payload = data[ethernetHeaderSize:]
 	return nil
 }
 
-// NextLayer returns the name of the next layer protocol based on the EtherType field of the EthernetFrame.
-func (ef *EthernetFrame) NextLayer() string {
-	var ets string
+// NextLayer returns the name and payload of the next layer protocol based on the EtherType field of the EthernetFrame.
+func (ef *EthernetFrame) NextLayer() (string, []byte) {
+	var layer string
 	switch ef.EtherType {
 	case 0x0800:
-		ets = "IPv4"
+		layer = "IPv4"
 	case 0x0806:
-		ets = "ARP"
+		layer = "ARP"
 	case 0x86dd:
-		ets = "IPv6"
+		layer = "IPv6"
 	default:
-		ets = "Unknown"
+		layer = ""
 	}
-	return ets
+	return layer, ef.payload
 }

layers/ftp.go

@@ -3,3 +3,15 @@ package layers
 // https://mavlink.io/zh/services/ftp.html
 // port 21 port 20
 type FTPMessage struct{}
+
+func (f *FTPMessage) String() string {
+	return ""
+}
+
+func (f *FTPMessage) Parse(data []byte) error {
+	return nil
+}
+
+func (f *FTPMessage) NextLayer() (string, []byte) {
+	return "", nil
+}

layers/http.go

@@ -3,3 +3,15 @@ package layers
 // https://developer.mozilla.org/en-US/docs/Web/HTTP/Messages
 // port 80
 type HTTPMessage struct{}
+
+func (h *HTTPMessage) String() string {
+	return ""
+}
+
+func (h *HTTPMessage) Parse(data []byte) error {
+	return nil
+}
+
+func (h *HTTPMessage) NextLayer() (string, []byte) {
+	return "", nil
+}

layers/icmp.go

@@ -58,6 +58,9 @@ func (i *ICMPSegment) Parse(data []byte) error {
 	}
 	return nil
 }
+func (i *ICMPSegment) NextLayer() (string, []byte) {
+	return "", nil
+}
 
 func (i *ICMPSegment) typecode() (string, string) {
 	// https://en.wikipedia.org/wiki/Internet_Control_Message_Protocol

layers/icmpv6.go

@@ -61,6 +61,10 @@ func (i *ICMPv6Segment) Parse(data []byte) error {
 	return nil
 }
 
+func (i *ICMPv6Segment) NextLayer() (string, []byte) {
+	return "", nil
+}
+
 func (i *ICMPv6Segment) typecode() (string, string) {
 	// https://en.wikipedia.org/wiki/ICMPv6
 	var mtype, code string

layers/ipv4.go

@@ -24,10 +24,11 @@ type IPv4Packet struct {
 	SrcIP          netip.Addr // IPv4 address of the sender of the packet.
 	DstIP          netip.Addr // IPv4 address of the receiver of the packet.
 	Options        []byte     // if ihl > 5
-	Payload        []byte
+	payload        []byte
 }
 
 func (p *IPv4Packet) String() string {
+	proto, _ := p.NextLayer()
 	return fmt.Sprintf(`IPv4 Packet:
 - Version: %d
 - IHL: %d
@@ -55,14 +56,14 @@ func (p *IPv4Packet) String() string {
 		p.flags(),
 		p.FragmentOffset,
 		p.TTL,
-		p.NextLayer(),
+		proto,
 		p.Protocol,
 		p.HeaderChecksum,
 		p.SrcIP,
 		p.DstIP,
 		p.Options,
-		len(p.Payload),
-		p.Payload)
+		len(p.payload),
+		p.payload)
 }
 
 // Parse parses the given byte data into an IPv4 packet struct.
@@ -89,27 +90,27 @@ func (p *IPv4Packet) Parse(data []byte) error {
 	if p.IHL > 5 {
 		offset := headerSizeIPv4 + ((p.IHL - 5) << 2)
 		p.Options = data[headerSizeIPv4:offset]
-		p.Payload = data[offset:]
+		p.payload = data[offset:]
 	} else {
-		p.Payload = data[headerSizeIPv4:]
+		p.payload = data[headerSizeIPv4:]
 	}
 	return nil
 }
 
-func (p *IPv4Packet) NextLayer() string {
+func (p *IPv4Packet) NextLayer() (string, []byte) {
 	// https://en.wikipedia.org/wiki/List_of_IP_protocol_numbers
-	var proto string
+	var layer string
 	switch p.Protocol {
 	case 1:
-		proto = "ICMP"
+		layer = "ICMP"
 	case 6:
-		proto = "TCP"
+		layer = "TCP"
 	case 17:
-		proto = "UDP"
+		layer = "UDP"
 	default:
-		proto = "Unknown"
+		layer = ""
 	}
-	return proto
+	return layer, p.payload
 }
 
 func (p *IPv4Packet) dscp() string {

layers/ipv6.go

@@ -22,10 +22,11 @@ type IPv6Packet struct {
 	HopLimit uint8
 	SrcIP    netip.Addr // The unicast IPv6 address of the sending node.
 	DstIP    netip.Addr // The IPv6 unicast or multicast address of the destination node(s).
-	Payload  []byte
+	payload  []byte
 }
 
 func (p *IPv6Packet) String() string {
+	proto, _ := p.NextLayer()
 	return fmt.Sprintf(`IPv6 Packet:
 - Version: %d
 - Traffic Class: %s
@@ -39,13 +40,13 @@ func (p *IPv6Packet) String() string {
 		p.Version,
 		p.trafficClass(),
 		p.PayloadLength,
-		p.NextLayer(),
+		proto,
 		p.NextHeader,
 		p.HopLimit,
 		p.SrcIP,
 		p.DstIP,
-		len(p.Payload),
-		p.Payload,
+		len(p.payload),
+		p.payload,
 	)
 }
 
@@ -63,45 +64,45 @@ func (p *IPv6Packet) Parse(data []byte) error {
 	p.HopLimit = data[7]
 	p.SrcIP, _ = netip.AddrFromSlice(data[8:24])
 	p.DstIP, _ = netip.AddrFromSlice(data[24:headerSizeIPv6])
-	p.Payload = data[headerSizeIPv6:]
+	p.payload = data[headerSizeIPv6:]
 	return nil
 }
 
-// NextLayer returns the type of the next header.
-func (p *IPv6Packet) NextLayer() string {
+// NextLayer returns the type of the next header and payload.
+func (p *IPv6Packet) NextLayer() (string, []byte) {
 	// https://en.wikipedia.org/wiki/List_of_IP_protocol_numbers
-	var proto string
+	var layer string
 	switch p.NextHeader {
 	case 0:
-		proto = "HOPOPT"
+		layer = "HOPOPT"
 	case 6:
-		proto = "TCP"
+		layer = "TCP"
 	case 17:
-		proto = "UDP"
+		layer = "UDP"
 	case 43:
-		proto = "Route"
+		layer = "Route"
 	case 44:
-		proto = "Fragment"
+		layer = "Fragment"
 	case 50:
-		proto = "Encapsulating Security Payload"
+		layer = "Encapsulating Security payload"
 	case 51:
-		proto = "Authentication Header"
+		layer = "Authentication Header"
 	case 58:
-		proto = "ICMPv6"
+		layer = "ICMPv6"
 	case 59:
-		proto = "NoNxt"
+		layer = "NoNxt"
 	case 60:
-		proto = "Opts"
+		layer = "Opts"
 	case 135:
-		proto = "Mobility"
+		layer = "Mobility"
 	case 139:
-		proto = "Host Identity Protocol"
+		layer = "Host Identity Protocol"
 	case 140:
-		proto = "Shim6 Protocol"
+		layer = "Shim6 Protocol"
 	default:
-		proto = "Unknown"
+		layer = ""
 	}
-	return proto
+	return layer, p.payload
 }
 
 func (p *IPv6Packet) trafficClass() string {

layers/layers.go

@@ -0,0 +1,47 @@
+package layers
+
+import "fmt"
+
+var LayerMap = map[string]Layer{
+	"ETH":    &EthernetFrame{},
+	"IPv4":   &IPv4Packet{},
+	"IPv6":   &IPv6Packet{},
+	"ARP":    &ARPPacket{},
+	"TCP":    &TCPSegment{},
+	"UDP":    &UDPSegment{},
+	"ICMP":   &ICMPSegment{},
+	"ICMPv6": &ICMPv6Segment{},
+	"DNS":    &DNSMessage{},
+	"FTP":    &DNSMessage{},
+	"HTTP":   &HTTPMessage{},
+	"SNMP":   &SNMPMessage{},
+	"SSH":    &SSHMessage{},
+	"TLS":    &TLSMessage{},
+}
+
+type Layer interface {
+	fmt.Stringer
+	Parse(data []byte) error
+	NextLayer() (name string, payload []byte)
+}
+
+func nextAppLayer(src, dst uint16) string {
+	var layer string
+	switch {
+	case src == 20 || dst == 20 || src == 21 || dst == 21:
+		layer = "FTP"
+	case src == 22 || dst == 22:
+		layer = "SSH"
+	case src == 53 || dst == 53:
+		layer = "DNS"
+	case src == 80 || dst == 80:
+		layer = "HTTP"
+	case src == 161 || dst == 161 || src == 162 || dst == 162:
+		layer = "SNMP"
+	case src == 443 || dst == 443:
+		layer = "TLS"
+	default:
+		layer = ""
+	}
+	return layer
+}