package api

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"math/big"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

func TestBrowserOIDCSession(t *testing.T) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	var issuer, challenge string
	exchanges := 0
	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{"issuer": issuer, "jwks_uri": issuer + "/keys", "authorization_endpoint": issuer + "/authorize", "token_endpoint": issuer + "/token"})
	})
	mux.HandleFunc("/keys", func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{"keys": []any{map[string]string{"kty": "RSA", "use": "sig", "alg": "RS256", "kid": "test-key", "n": base64.RawURLEncoding.EncodeToString(key.N.Bytes()), "e": base64.RawURLEncoding.EncodeToString(big.NewInt(int64(key.E)).Bytes())}}})
	})
	mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
		exchanges++
		_ = r.ParseForm()
		digest := sha256.Sum256([]byte(r.Form.Get("code_verifier")))
		if r.Form.Get("code") != "valid-code" || base64.RawURLEncoding.EncodeToString(digest[:]) != challenge || r.Form.Get("redirect_uri") != "https://dns.example/auth/oidc/callback" {
			t.Error("incorrect code, PKCE verifier or redirect")
			http.Error(w, "invalid grant", 400)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		token := signedOIDCToken(t, key, map[string]any{"iss": issuer, "aud": "https://dns.example", "sub": "admin", "token_use": "access", "roles": []string{"admin"}, "exp": time.Now().Add(2 * time.Minute).Unix()})
		_ = json.NewEncoder(w).Encode(map[string]any{"access_token": token, "token_type": "Bearer", "expires_in": 120})
	})
	provider := httptest.NewTLSServer(mux)
	defer provider.Close()
	issuer = provider.URL
	original := http.DefaultTransport
	http.DefaultTransport = provider.Client().Transport
	defer func() { http.DefaultTransport = original }()
	h := newHarness(t, func(c *config.Config) {
		c.API.Auth.Mode = "oidc"
		c.API.Auth.OIDCIssuer = issuer
		c.API.Auth.OIDCAudience = "https://dns.example"
		c.API.Auth.OIDCClientID = "dns"
	})
	start := httptest.NewRecorder()
	h.srv.Handler().ServeHTTP(start, httptest.NewRequest("GET", "https://dns.example/api/auth/oidc/start?next=%2Fzones", nil))
	if start.Code != 302 {
		t.Fatalf("start: %d %s", start.Code, start.Body.String())
	}
	target, err := url.Parse(start.Header().Get("Location"))
	if err != nil {
		t.Fatal(err)
	}
	challenge = target.Query().Get("code_challenge")
	if challenge == "" || target.Query().Get("code_challenge_method") != "S256" {
		t.Fatal("missing PKCE")
	}
	binding := start.Result().Cookies()[0]
	callback := func(origin, state string, cookie *http.Cookie) *httptest.ResponseRecorder {
		b, _ := json.Marshal(map[string]string{"code": "valid-code", "state": state})
		r := httptest.NewRequest("POST", "https://dns.example/api/auth/oidc/callback", strings.NewReader(string(b)))
		r.Header.Set("Origin", origin)
		r.Header.Set("Content-Type", "application/json")
		if cookie != nil {
			r.AddCookie(cookie)
		}
		w := httptest.NewRecorder()
		h.srv.Handler().ServeHTTP(w, r)
		return w
	}
	state := target.Query().Get("state")
	for _, tc := range []struct {
		origin, state string
		cookie        *http.Cookie
	}{{"https://evil.example", state, binding}, {"https://dns.example", "wrong", binding}, {"https://dns.example", state, nil}} {
		if w := callback(tc.origin, tc.state, tc.cookie); w.Code == 200 {
			t.Fatal("accepted unbound or cross-origin callback")
		}
	}
	if exchanges != 0 {
		t.Fatal("invalid callback reached token endpoint")
	}
	w := callback("https://dns.example", state, binding)
	if w.Code != 200 {
		t.Fatalf("callback: %d %s", w.Code, w.Body.String())
	}
	if strings.Contains(w.Body.String(), "access_token") || !strings.Contains(w.Body.String(), "/zones") {
		t.Fatal("callback leaked token or lost destination")
	}
	var sessionCookie *http.Cookie
	for _, c := range w.Result().Cookies() {
		if c.Name == browserSessionCookie {
			sessionCookie = c
		}
	}
	if sessionCookie == nil || !sessionCookie.Secure || !sessionCookie.HttpOnly || sessionCookie.SameSite != http.SameSiteLaxMode || sessionCookie.MaxAge > 120 {
		t.Fatalf("unsafe session cookie: %+v", sessionCookie)
	}
	if w := callback("https://dns.example", state, binding); w.Code == 200 || exchanges != 1 {
		t.Fatal("authorization code replay accepted")
	}
	request := func(method, path, origin string) *httptest.ResponseRecorder {
		r := httptest.NewRequest(method, "https://dns.example"+path, nil)
		r.AddCookie(sessionCookie)
		if origin != "" {
			r.Header.Set("Origin", origin)
		}
		w := httptest.NewRecorder()
		h.srv.Handler().ServeHTTP(w, r)
		return w
	}
	if w := request("GET", "/api/whoami", ""); w.Code != 200 {
		t.Fatalf("session refused: %d %s", w.Code, w.Body.String())
	}
	if w := request("POST", "/api/cache/flush", "https://evil.example"); w.Code == 200 {
		t.Fatal("cross-origin cookie mutation accepted")
	}
	if w := request("POST", "/api/auth/logout", "https://dns.example"); w.Code != 200 {
		t.Fatal("logout failed")
	}
	if w := request("GET", "/api/whoami", ""); w.Code != 401 {
		t.Fatal("revoked session still works")
	}
}

func TestSafeLoginNext(t *testing.T) {
	for _, next := range []string{"//evil.example", "/\\evil.example", "/auth/login", "/\r\nLocation: evil", "https://evil.example"} {
		if safeLoginNext(next) != "/" {
			t.Fatalf("unsafe redirect %q", next)
		}
	}
	if safeLoginNext("/zones?zone=example.net") != "/zones?zone=example.net" {
		t.Fatal("lost local destination")
	}
}