package probe

import (
	"context"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"strings"
	"testing"

	"github.com/itworx/pulse/internal/audit"
)

type resolverFunc func(context.Context, string) ([]netip.Addr, error)

func (f resolverFunc) LookupIP(ctx context.Context, host string) ([]netip.Addr, error) {
	return f(ctx, host)
}

func privatePolicy() NetworkPolicy {
	return NetworkPolicy{Revision: 1, AllowedNetworks: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}}
}

func TestNetworkPolicyBlocksSSRFAndAllowsExplicitLAN(t *testing.T) {
	policy := privatePolicy()
	resolver := resolverFunc(func(_ context.Context, host string) ([]netip.Addr, error) {
		if host == "lan.internal" {
			return []netip.Addr{netip.MustParseAddr("10.10.0.8")}, nil
		}
		return []netip.Addr{netip.MustParseAddr("169.254.169.254")}, nil
	})
	if _, err := policy.ValidateTarget(context.Background(), resolver, "https", "lan.internal", 443); err != nil {
		t.Fatal(err)
	}
	if _, err := policy.ValidateTarget(context.Background(), resolver, "https", "metadata.internal", 443); err == nil {
		t.Fatal("metadata target was allowed")
	}
	if _, err := policy.ValidateTarget(context.Background(), resolver, "https", resolverHost("127.0.0.1"), 80); err == nil {
		t.Fatal("loopback target was allowed")
	}
}

func TestNetworkPolicyRejectsDNSRebindingAndRedirectEscape(t *testing.T) {
	policy := privatePolicy()
	addresses := []netip.Addr{netip.MustParseAddr("10.0.0.2")}
	resolver := resolverFunc(func(_ context.Context, _ string) ([]netip.Addr, error) {
		current := addresses
		addresses = []netip.Addr{netip.MustParseAddr("169.254.169.254")}
		return current, nil
	})
	if _, err := policy.ValidateTarget(context.Background(), resolver, "https", "service.internal", 443); err != nil {
		t.Fatal(err)
	}
	if _, err := policy.ValidateTarget(context.Background(), resolver, "https", "service.internal", 443); err == nil {
		t.Fatal("DNS rebinding to a blocked address was allowed")
	}
	resolver = resolverFunc(func(_ context.Context, _ string) ([]netip.Addr, error) {
		return []netip.Addr{netip.MustParseAddr("10.0.0.2"), netip.MustParseAddr("169.254.169.254")}, nil
	})
	if _, err := policy.ValidateTarget(context.Background(), resolver, "https", "service.internal", 443); err == nil {
		t.Fatal("mixed DNS answer was allowed")
	}
	client, err := NewSafeClient(policy, resolverFunc(func(_ context.Context, host string) ([]netip.Addr, error) {
		if host == "safe.internal" {
			return []netip.Addr{netip.MustParseAddr("10.0.0.2")}, nil
		}
		return []netip.Addr{netip.MustParseAddr("169.254.169.254")}, nil
	}))
	if err != nil {
		t.Fatal(err)
	}
	redirect := httptest.NewRequest(http.MethodGet, "https://metadata.internal/", nil)
	if err := client.HTTP.CheckRedirect(redirect, nil); err == nil {
		t.Fatal("redirect escape was allowed")
	}
}

func TestNetworkPolicyBoundsRequestsResponsesAndAudits(t *testing.T) {
	policy := privatePolicy()
	request := httptest.NewRequest(http.MethodPost, "https://example.com", strings.NewReader("body"))
	if err := policy.ValidateRequest(request); err == nil {
		t.Fatal("POST/body request was allowed")
	}
	request = httptest.NewRequest(http.MethodGet, "https://example.com", nil)
	request.Header.Set("Authorization", "secret")
	if err := policy.ValidateRequest(request); err == nil {
		t.Fatal("authorization header was allowed")
	}
	response := &http.Response{Body: ioNopCloser{Reader: strings.NewReader("12345")}}
	if _, err := ReadLimitedBody(response, 4); err == nil {
		t.Fatal("oversized response was allowed")
	}
	store := &audit.MemoryStore{}
	if err := RecordPolicyChange(context.Background(), store, "user", "policy", "corr", NetworkPolicy{Revision: 1}, policy); err != nil {
		t.Fatal(err)
	}
	if len(store.Events) != 1 || store.Events[0].Action != "probe.network_policy.update" || store.Events[0].After["maxResponseBytes"] == nil {
		t.Fatalf("audit=%+v", store.Events)
	}
	if err := RecordPolicyChange(context.Background(), nil, "user", "policy", "corr", policy, policy); err == nil {
		t.Fatal("nil audit store was accepted")
	}
}

type ioNopCloser struct{ *strings.Reader }

func (ioNopCloser) Close() error { return nil }

func resolverHost(host string) string { return host }