package dns

import (
	"testing"

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

func TestDetectDNS(t *testing.T) {
	// A valid DNS query for "www.google.com" (A record)
	dnsQuery := []byte{
		0x12, 0x34, // Transaction ID
		0x01, 0x00, // Flags (Standard Query, Recursion Desired)
		0x00, 0x01, // Questions: 1
		0x00, 0x00, // Answer RRs: 0
		0x00, 0x00, // Authority RRs: 0
		0x00, 0x00, // Additional RRs: 0
		0x03, 'w', 'w', 'w', // Question section:
		0x06, 'g', 'o', 'o', 'g', 'l', 'e', //
		0x03, 'c', 'o', 'm', //
		0x00,       // Null terminator for the name
		0x00, 0x01, // Type: A (Host Address)
		0x00, 0x01, // Class: IN (Internet)
	}

	// A valid DNS reply for the above query
	dnsReply := []byte{
		0x12, 0x34, // Transaction ID (matches query)
		0x81, 0x80, // Flags (Response, Recursion Desired, Recursion Available)
		0x00, 0x01, // Questions: 1
		0x00, 0x01, // Answer RRs: 1
		0x00, 0x00, // Authority RRs: 0
		0x00, 0x00, // Additional RRs: 0
		0x03, 'w', 'w', 'w', // Question section (same as query)
		0x06, 'g', 'o', 'o', 'g', 'l', 'e', //
		0x03, 'c', 'o', 'm', //
		0x00,       //
		0x00, 0x01, //
		0x00, 0x01, //
		0xc0, 0x0c, // Answer section. Name: Pointer to the question name at offset 12 (0x0c)
		0x00, 0x01, // Type: A
		0x00, 0x01, // Class: IN
		0x00, 0x00, 0x01, 0x2c, // Time to Live (TTL): 300 seconds
		0x00, 0x04, // Data Length: 4 bytes
		0x08, 0x08, 0x08, 0x08, // RDATA: IP Address 8.8.8.8
	}

	// This is just the first 12 bytes of a DNS query with a single question.
	headerWithQuestionPromise := []byte{
		0x12, 0x34, // Transaction ID
		0x01, 0x00, // Flags (Standard Query)
		0x00, 0x01, // Questions: 1  <-- This promises more data!
		0x00, 0x00, // Answer RRs: 0
		0x00, 0x00, // Authority RRs: 0
		0x00, 0x00, // Additional RRs: 0
	}

	// This is a valid, self-contained "empty" DNS message.
	headerWithNoPromise := []byte{
		0x56, 0x78, // Transaction ID
		0x01, 0x00, // Flags (Standard Query)
		0x00, 0x00, // Questions: 0  <-- This promises NO more data.
		0x00, 0x00, // Answer RRs: 0
		0x00, 0x00, // Authority RRs: 0
		0x00, 0x00, // Additional RRs: 0
	}

	t.Run("query", func(t *testing.T) {
		p, c, err := protocol.Detect(protocol.Both, dnsQuery, 1234, 53)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
			return
		}
		t.Logf("detected %s confidence %g%%", p.Name, c*100)
		if p.Name != protocolName {
			t.Errorf("expected %q protocol, got %q", protocolName, p.Name)
		}
	})

	t.Run("answer", func(t *testing.T) {
		p, c, err := protocol.Detect(protocol.Both, dnsReply, 53, 1234)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
			return
		}
		t.Logf("detected %s confidence %g%%", p.Name, c*100)
		if p.Name != protocolName {
			t.Errorf("expected %q protocol, got %q", protocolName, p.Name)
		}
	})

	t.Run("header with no promise", func(t *testing.T) {
		p, c, err := protocol.Detect(protocol.Both, headerWithNoPromise, 1234, 53)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
			return
		}
		t.Logf("detected %s confidence %g%%", p.Name, c*100)
		if p.Name != protocolName {
			t.Errorf("expected %q protocol, got %q", protocolName, p.Name)
		}
	})

	t.Run("header with question promise", func(t *testing.T) {
		p, c, err := protocol.Detect(protocol.Both, headerWithQuestionPromise, 1234, 53)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
			return
		}
		t.Logf("detected %s confidence %g%%", p.Name, c*100)
		if p.Name != protocolName {
			t.Errorf("expected %q protocol, got %q", protocolName, p.Name)
		}
	})
}