// Package jwt scans responses for JWTs, decodes them, and flags
// security-relevant attributes: alg=none, weak HMAC secret (dictionary
// crack against common passwords), excessive expiration, missing claims.
//
// The brute-force list is intentionally tiny (~20 common secrets) — the
// goal is to surface obviously-weak keys, not to run offline hashcat. A
// proper cracker belongs in Fase 2's planned "auth" agent.
package jwt

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"encoding/json"
	"hash"
	"io"
	"net/http"
	"regexp"
	"strings"
	"sync"
	"time"

	"god-eye/internal/eventbus"
	gohttp "god-eye/internal/http"
	"god-eye/internal/module"
	"god-eye/internal/store"
)

const ModuleName = "vuln.jwt"

type jwtModule struct{}

func Register() { module.Register(&jwtModule{}) }

func (*jwtModule) Name() string                   { return ModuleName }
func (*jwtModule) Phase() module.Phase            { return module.PhaseAnalysis }
func (*jwtModule) Consumes() []eventbus.EventType { return []eventbus.EventType{eventbus.EventHTTPProbed} }
func (*jwtModule) Produces() []eventbus.EventType {
	return []eventbus.EventType{eventbus.EventVulnerability, eventbus.EventSecret}
}
func (*jwtModule) DefaultEnabled() bool { return true }

// jwtRegex matches the standard three-part base64url JWT shape.
var jwtRegex = regexp.MustCompile(`eyJ[A-Za-z0-9_-]+\.eyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*`)

var weakSecrets = []string{
	"secret", "password", "123456", "admin", "jwt", "jwtsecret",
	"changeme", "default", "test", "dev", "secret_key", "mysecret",
	"your-256-bit-secret", "your-secret-key", "super-secret",
	"supersecret", "helloworld", "qwerty", "abc123", "letmein",
}

func (*jwtModule) Run(mctx module.Context) error {
	timeout := mctx.Config.Int("timeout", 10)
	client := gohttp.GetSharedClient(timeout)

	processed := make(map[string]struct{})
	var mu sync.Mutex
	shouldProcess := func(host string) bool {
		mu.Lock()
		defer mu.Unlock()
		if _, ok := processed[host]; ok {
			return false
		}
		processed[host] = struct{}{}
		return true
	}

	var wg sync.WaitGroup

	for _, h := range mctx.Store.All(mctx.Ctx) {
		if h == nil || h.StatusCode == 0 {
			continue
		}
		if !shouldProcess(h.Subdomain) {
			continue
		}
		host := h.Subdomain
		wg.Add(1)
		go func() { defer wg.Done(); scanHost(mctx, client, host) }()
	}

	sub := mctx.Bus.Subscribe(eventbus.EventHTTPProbed, func(_ context.Context, e eventbus.Event) {
		ev, ok := e.(eventbus.HTTPProbed)
		if !ok || ev.StatusCode == 0 {
			return
		}
		host := ev.Meta().Target
		if !shouldProcess(host) {
			return
		}
		wg.Add(1)
		go func() { defer wg.Done(); scanHost(mctx, client, host) }()
	})
	defer sub.Unsubscribe()

	select {
	case <-time.After(500 * time.Millisecond):
	case <-mctx.Ctx.Done():
	}

	wg.Wait()
	return nil
}

func scanHost(mctx module.Context, client *http.Client, host string) {
	for _, scheme := range []string{"https://", "http://"} {
		if mctx.Ctx.Err() != nil {
			return
		}
		url := scheme + host
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			continue
		}
		req.Header.Set("User-Agent", "god-eye-v2")
		resp, err := client.Do(req)
		if err != nil {
			continue
		}
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 256*1024))
		resp.Body.Close()

		text := string(body)
		// Also check Authorization + Set-Cookie response headers.
		for _, h := range resp.Header.Values("Set-Cookie") {
			text += "\n" + h
		}
		if auth := resp.Header.Get("Authorization"); auth != "" {
			text += "\n" + auth
		}

		matches := jwtRegex.FindAllString(text, -1)
		for _, tok := range uniqueStrings(matches) {
			analyzeJWT(mctx, host, url, tok)
		}
		// One scheme is enough; avoid duplicate noise.
		if len(matches) > 0 {
			return
		}
	}
}

func analyzeJWT(mctx module.Context, host, url, token string) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return
	}
	header, err := base64Decode(parts[0])
	if err != nil {
		return
	}
	payload, err := base64Decode(parts[1])
	if err != nil {
		return
	}

	var h struct {
		Alg string `json:"alg"`
		Kid string `json:"kid"`
		Typ string `json:"typ"`
	}
	if err := json.Unmarshal(header, &h); err != nil {
		return
	}

	severity := eventbus.SeverityInfo
	findings := []string{"JWT detected"}

	if strings.EqualFold(h.Alg, "none") {
		severity = eventbus.SeverityCritical
		findings = append(findings, "alg=none accepted — no signature verification")
	}
	if strings.HasPrefix(strings.ToUpper(h.Alg), "HS") {
		if cracked := tryWeakSecret(token, h.Alg, parts); cracked != "" {
			severity = eventbus.SeverityCritical
			findings = append(findings, "weak HMAC secret cracked: "+cracked)
		}
	}
	if h.Kid != "" && looksInjectable(h.Kid) {
		severity = maxSeverity(severity, eventbus.SeverityMedium)
		findings = append(findings, "kid header may be injectable: "+h.Kid)
	}

	// Inspect payload for excessive expiry.
	var claims map[string]interface{}
	_ = json.Unmarshal(payload, &claims)
	if exp, ok := claims["exp"].(float64); ok {
		expAt := time.Unix(int64(exp), 0)
		if time.Until(expAt) > 365*24*time.Hour {
			severity = maxSeverity(severity, eventbus.SeverityLow)
			findings = append(findings, "exp >1 year")
		}
	}

	redacted := token
	if len(redacted) > 40 {
		redacted = redacted[:20] + "…" + redacted[len(redacted)-10:]
	}

	_ = mctx.Store.Upsert(mctx.Ctx, host, func(sh *store.Host) {
		sh.Secrets = append(sh.Secrets, store.Secret{
			Kind:        "jwt",
			Match:       redacted,
			Location:    url,
			Severity:    string(severity),
			Description: strings.Join(findings, "; "),
			FoundAt:     time.Now(),
		})
	})

	mctx.Bus.Publish(mctx.Ctx, eventbus.SecretFound{
		EventMeta:   eventbus.EventMeta{At: time.Now(), Source: ModuleName, Target: host},
		Kind:        "jwt",
		Match:       redacted,
		Location:    url,
		Severity:    severity,
		Description: strings.Join(findings, "; "),
	})

	if severity == eventbus.SeverityCritical || severity == eventbus.SeverityHigh {
		mctx.Bus.Publish(mctx.Ctx, eventbus.VulnerabilityFound{
			EventMeta:   eventbus.EventMeta{At: time.Now(), Source: ModuleName, Target: host},
			ID:          "jwt-weak",
			Title:       "JWT Weakness",
			Description: strings.Join(findings, "; "),
			Severity:    severity,
			URL:         url,
			Evidence:    redacted,
			Remediation: "Use strong signing keys (256+ bits of entropy), refuse alg=none, rotate keys on compromise, short expiry.",
			OWASP:       "A02:2021-Cryptographic Failures",
		})
	}
}

func tryWeakSecret(token, alg string, parts []string) string {
	signingInput := parts[0] + "." + parts[1]
	sig, err := base64Decode(parts[2])
	if err != nil {
		return ""
	}

	var hashFn func() hash.Hash
	switch strings.ToUpper(alg) {
	case "HS256":
		hashFn = sha256.New
	case "HS384":
		hashFn = func() hash.Hash { return sha512.New384() }
	case "HS512":
		hashFn = sha512.New
	default:
		return ""
	}

	for _, s := range weakSecrets {
		mac := hmac.New(hashFn, []byte(s))
		mac.Write([]byte(signingInput))
		if hmac.Equal(mac.Sum(nil), sig) {
			return s
		}
	}
	return ""
}

// base64Decode unpads and decodes a JWT segment (URL-safe, no padding).
func base64Decode(s string) ([]byte, error) {
	// Add padding if missing.
	if m := len(s) % 4; m != 0 {
		s += strings.Repeat("=", 4-m)
	}
	return base64.URLEncoding.DecodeString(s)
}

func looksInjectable(kid string) bool {
	// kids that include path separators, SQL wildcards, or NUL-like
	// sequences are worth flagging for manual review.
	return strings.ContainsAny(kid, "/\\;'\"$`|")
}

func maxSeverity(a, b eventbus.Severity) eventbus.Severity {
	rank := map[eventbus.Severity]int{
		eventbus.SeverityInfo: 0, eventbus.SeverityLow: 1,
		eventbus.SeverityMedium: 2, eventbus.SeverityHigh: 3, eventbus.SeverityCritical: 4,
	}
	if rank[a] >= rank[b] {
		return a
	}
	return b
}

func uniqueStrings(in []string) []string {
	seen := make(map[string]struct{})
	out := make([]string, 0, len(in))
	for _, s := range in {
		if _, dup := seen[s]; dup {
			continue
		}
		seen[s] = struct{}{}
		out = append(out, s)
	}
	return out
}