// bad-actors ranks client IPs by scanner behaviour in raw nginx/ingress spans
// and reports which would be safe to block. It only reads ClickHouse; nothing
// is blocked. Review its output before feeding any blocker.
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/ClickHouse/clickhouse-go/v2"

	"git.tornberg.me/go-otel/internal/honeypot"
	"git.tornberg.me/go-otel/internal/tracectx"
	"git.tornberg.me/go-otel/internal/triage"
	"git.tornberg.me/go-otel/internal/uaclass"
)

type otlp struct {
	ResourceSpans []struct {
		ScopeSpans []struct {
			Spans []struct {
				TraceID      string `json:"traceId"`
				SpanID       string `json:"spanId"`
				ParentSpanID string `json:"parentSpanId"`
				Attributes   []struct {
					Key   string                     `json:"key"`
					Value map[string]json.RawMessage `json:"value"`
				} `json:"attributes"`
			} `json:"spans"`
		} `json:"scopeSpans"`
	} `json:"resourceSpans"`
}

type actor struct {
	IP        string
	Total     int
	Probes    int
	Paths     map[string]struct{}
	Statuses  map[string]int
	UserAgent string
}

func main() {
	addr := flag.String("addr", envOr("OTEL_CLICKHOUSE_ADDR", "127.0.0.1:9000"), "ClickHouse host:port")
	db := flag.String("database", envOr("OTEL_CLICKHOUSE_DB", "otel"), "database")
	user := flag.String("user", envOr("OTEL_CLICKHOUSE_USER", "default"), "user")
	pass := flag.String("password", os.Getenv("OTEL_CLICKHOUSE_PASSWORD"), "password")
	since := flag.String("since", "", "inclusive UTC RFC3339")
	until := flag.String("until", "", "exclusive UTC RFC3339")
	service := flag.String("service", "nginx", "ingress service name")
	minProbes := flag.Int("min-probes", 10, "minimum probe requests to be blockable")
	minPaths := flag.Int("min-paths", 3, "minimum distinct probe paths")
	minShare := flag.Float64("min-share", 0.8, "minimum share of the IP's requests that are probes")
	broadPaths := flag.Int("broad-paths", 50, "distinct probe paths that mark a scanner regardless of share")
	broadProbes := flag.Int("broad-probes", 100, "probe requests required with -broad-paths")
	top := flag.Int("top", 30, "rows to print")
	flag.Parse()
	start, err1 := time.Parse(time.RFC3339, *since)
	end, err2 := time.Parse(time.RFC3339, *until)
	if err1 != nil || err2 != nil || !start.Before(end) {
		log.Fatal("-since/-until must be valid RFC3339 and ordered")
	}
	conn, err := clickhouse.Open(&clickhouse.Options{Addr: []string{*addr}, Auth: clickhouse.Auth{Database: *db, Username: *user, Password: *pass}})
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
	defer cancel()
	rows, err := conn.Query(ctx, fmt.Sprintf(`SELECT data FROM %s.entries WHERE type='traces' AND service=? AND timestamp >= toDateTime64('%s',9,'UTC') AND timestamp < toDateTime64('%s',9,'UTC')`,
		*db, start.UTC().Format("2006-01-02 15:04:05.000000000"), end.UTC().Format("2006-01-02 15:04:05.000000000")), *service)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	// Pass 1 collects every span; pass 2 resolves parents, because the outer
	// nginx span (real client) and the inner loopback span are exported in
	// different envelopes and arrive in either order.
	type record struct {
		key                  tracectx.Key
		ip, path, status, ua string
	}
	var records []record
	var all []tracectx.Span
	for rows.Next() {
		var data string
		if err := rows.Scan(&data); err != nil {
			log.Fatal(err)
		}
		var d otlp
		if json.Unmarshal([]byte(data), &d) != nil {
			continue
		}
		for _, rs := range d.ResourceSpans {
			for _, ss := range rs.ScopeSpans {
				for _, sp := range ss.Spans {
					at := map[string]string{}
					for _, a := range sp.Attributes {
						for _, raw := range a.Value {
							var v any
							_ = json.Unmarshal(raw, &v)
							at[a.Key] = fmt.Sprint(v)
						}
					}
					ip := strings.TrimSpace(strings.SplitN(first(at, "http.request.header.x-real-ip", "http.request.header.x-forwarded-for", "client.address", "net.peer.ip"), ",", 2)[0])
					path := strings.SplitN(first(at, "http.target", "url.path", "http.route"), "?", 2)[0]
					all = append(all, tracectx.Span{TraceID: sp.TraceID, SpanID: sp.SpanID, ParentID: sp.ParentSpanID, Service: *service, PeerIP: ip, Edge: true})
					records = append(records, record{key: tracectx.Key{TraceID: sp.TraceID, SpanID: sp.SpanID}, ip: ip, path: path, status: first(at, "http.status_code"), ua: first(at, "http.user_agent")})
				}
			}
		}
	}
	if err := rows.Err(); err != nil {
		log.Fatal(err)
	}

	resolved := tracectx.Resolve(all)
	actors := map[string]*actor{}
	spans, recovered, recoveredProbes := 0, 0, 0
	for _, r := range records {
		ip := r.ip
		if in, ok := resolved[r.key]; ok && in.ClientIP != "" && in.ClientHops > 0 && internal(ip) {
			ip = in.ClientIP // hop behind a proxy: use the real client from an ancestor
			recovered++
			if triage.IsProbePath(r.path) {
				recoveredProbes++
			}
		}
		if ip == "" {
			continue
		}
		spans++
		a := actors[ip]
		if a == nil {
			a = &actor{IP: ip, Paths: map[string]struct{}{}, Statuses: map[string]int{}}
			actors[ip] = a
		}
		a.Total++
		if a.UserAgent == "" {
			a.UserAgent = r.ua
		}
		if triage.IsProbePath(r.path) {
			a.Probes++
			a.Paths[r.path] = struct{}{}
			a.Statuses[r.status]++
		}
	}
	fmt.Printf("resolved %d hop spans to their real client via the parent span (%d of them probes)\n", recovered, recoveredProbes)

	var list []*actor
	totalProbes, blockableProbes, blockable, skippedInternal, skippedInternalProbes := 0, 0, 0, 0, 0
	for _, a := range actors {
		totalProbes += a.Probes
		if a.Probes == 0 {
			continue
		}
		if internal(a.IP) {
			skippedInternal++
			skippedInternalProbes += a.Probes
			a.UserAgent = "[INTERNAL] " + a.UserAgent
		} else if (a.Probes >= *minProbes && len(a.Paths) >= *minPaths && float64(a.Probes)/float64(a.Total) >= *minShare) ||
			(a.Probes >= *broadProbes && len(a.Paths) >= *broadPaths) {
			blockable++
			blockableProbes += a.Probes
			a.UserAgent = "[BLOCK] " + a.UserAgent
		}
		list = append(list, a)
	}
	sort.Slice(list, func(i, j int) bool { return list[i].Probes > list[j].Probes })
	fmt.Printf("%d spans with an IP, %d distinct IPs, %d probe requests\n", spans, len(actors), totalProbes)
	fmt.Printf("blockable (public; >=%d probes/>=%d paths/>=%.0f%% probes, or >=%d probes over >=%d paths): %d IPs covering %d probes (%.0f%%)\n", *minProbes, *minPaths, *minShare*100, *broadProbes, *broadPaths, blockable, blockableProbes, pct(blockableProbes, totalProbes))
	fmt.Printf("internal/loopback (never blocked): %d IPs, %d probes (%.0f%%) - real client is hidden behind a proxy hop\n\n", skippedInternal, skippedInternalProbes, pct(skippedInternalProbes, totalProbes))
	fmt.Printf("%-40s %6s %6s %6s  %-14s %s\n", "ip", "probes", "total", "paths", "class", "verdict / user-agent / statuses")
	for i, a := range list {
		if i >= *top {
			break
		}
		fmt.Printf("%-40s %6d %6d %6d  %-14s %.60s %v\n", a.IP, a.Probes, a.Total, len(a.Paths), uaclass.Of(strings.TrimPrefix(strings.TrimPrefix(a.UserAgent, "[BLOCK] "), "[INTERNAL] ")), a.UserAgent, a.Statuses)
	}
}

// internal reports addresses that must never be blocked; it is the same
// definition the tracker and blocker use.
func internal(s string) bool { return honeypot.Untrackable(s) }

func first(m map[string]string, keys ...string) string {
	for _, k := range keys {
		if v := m[k]; v != "" && v != "<nil>" {
			return v
		}
	}
	return ""
}
func pct(a, b int) float64 {
	if b == 0 {
		return 0
	}
	return 100 * float64(a) / float64(b)
}
func envOr(k, f string) string {
	if v := os.Getenv(k); v != "" {
		return v
	}
	return f
}