package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"time"
)

// Bounded command output and deadline keep large conntrack tables from exhausting the agent.
type networkBuffer struct {
	bytes.Buffer
	limit int
}

func (b *networkBuffer) Write(p []byte) (int, error) {
	if len(p) > b.limit-b.Len() {
		return 0, errors.New("output limit exceeded")
	}
	return b.Buffer.Write(p)
}
func networkCommand(ctx context.Context, name string, args ...string) ([]byte, error) {
	ctx, cancel := context.WithTimeout(ctx, 4*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, name, args...)
	out := &networkBuffer{limit: 8 << 20}
	stderr := &networkBuffer{limit: 4096}
	cmd.Stdout = out
	cmd.Stderr = stderr
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("%s: %v %s", name, err, trafficText(stderr.String(), 240))
	}
	return out.Bytes(), nil
}
func parseNetworkFlow(line string) (NetworkFlow, bool) {
	var f NetworkFlow
	fields := strings.Fields(line)
	protoIndex := -1
	for i, v := range fields {
		if v == "tcp" || v == "udp" || v == "icmp" || v == "icmpv6" {
			f.Protocol = v
			protoIndex = i
			break
		}
	}
	if protoIndex < 0 {
		return f, false
	}
	values := map[string][]string{}
	for _, v := range fields[protoIndex+1:] {
		if k, val, ok := strings.Cut(v, "="); ok {
			values[k] = append(values[k], val)
		}
	}
	get := func(k string, i int) string {
		if len(values[k]) > i {
			return values[k][i]
		}
		return ""
	}
	integer := func(k string, i int) int { n, _ := strconv.Atoi(get(k, i)); return n }
	number := func(k string, i int) uint64 { n, _ := strconv.ParseUint(get(k, i), 10, 64); return n }
	f.Source = get("src", 0)
	f.Destination = get("dst", 0)
	f.ReplySource = get("src", 1)
	f.ReplyDestination = get("dst", 1)
	if net.ParseIP(f.Source) == nil || net.ParseIP(f.Destination) == nil || net.ParseIP(f.ReplySource) == nil || net.ParseIP(f.ReplyDestination) == nil {
		return f, false
	}
	f.SourcePort = integer("sport", 0)
	f.DestinationPort = integer("dport", 0)
	f.ReplySourcePort = integer("sport", 1)
	f.ReplyDestinationPort = integer("dport", 1)
	f.State = "TRACKED"
	if f.Protocol == "tcp" && len(fields) > protoIndex+3 && !strings.Contains(fields[protoIndex+3], "=") {
		f.State = fields[protoIndex+3]
	}
	if strings.Contains(line, "[UNREPLIED]") {
		f.State = "UNREPLIED"
		if f.Protocol == "tcp" && len(fields) > protoIndex+3 {
			f.State = fields[protoIndex+3]
		}
	}
	if f.Protocol == "udp" && strings.Contains(line, "[ASSURED]") {
		f.State = "ASSURED"
	}
	f.Accounting = len(values["bytes"]) == 2 && len(values["packets"]) == 2
	f.Bytes = number("bytes", 0)
	f.ReplyBytes = number("bytes", 1)
	f.Packets = number("packets", 0)
	f.ReplyPackets = number("packets", 1)
	f.Zone = get("zone", 0)
	return f, true
}
func appendNetworkFlows(snapshot *NetworkSnapshot, raw []byte) {
	for _, line := range strings.Split(string(raw), "\n") {
		if strings.TrimSpace(line) == "" {
			continue
		}
		f, ok := parseNetworkFlow(line)
		if !ok {
			if strings.Contains(line, "src=") {
				snapshot.ParseErrors++
			}
			continue
		}
		if len(snapshot.Flows) >= 5000 {
			snapshot.Limited = true
			continue
		}
		snapshot.Flows = append(snapshot.Flows, f)
	}
}

// Only counters immediately before a terminal drop/reject in a CrowdSec table/chain
// are reported as blocked packets. General chain and set counters are not summed.
func parseFirewallCounters(raw []byte) ([]FirewallCounter, error) {
	var doc struct {
		NFTables []struct {
			Rule *struct {
				Family string                       `json:"family"`
				Table  string                       `json:"table"`
				Chain  string                       `json:"chain"`
				Handle int                          `json:"handle"`
				Expr   []map[string]json.RawMessage `json:"expr"`
			} `json:"rule"`
		} `json:"nftables"`
	}
	if err := json.Unmarshal(raw, &doc); err != nil {
		return nil, err
	}
	result := []FirewallCounter{}
	for _, item := range doc.NFTables {
		r := item.Rule
		if r == nil || !strings.Contains(strings.ToLower(r.Table+" "+r.Chain), "crowdsec") {
			continue
		}
		if len(r.Expr) < 2 {
			continue
		}
		last := r.Expr[len(r.Expr)-1]
		_, drop := last["drop"]
		_, reject := last["reject"]
		if !drop && !reject {
			continue
		}
		counter, ok := r.Expr[len(r.Expr)-2]["counter"]
		if !ok {
			continue
		}
		var c struct {
			Packets uint64 `json:"packets"`
			Bytes   uint64 `json:"bytes"`
		}
		if json.Unmarshal(counter, &c) != nil || len(counter) == 0 || counter[0] != '{' {
			continue
		}
		result = append(result, FirewallCounter{Name: fmt.Sprintf("%s/%s/%s #%d", r.Family, r.Table, r.Chain, r.Handle), Packets: c.Packets, Bytes: c.Bytes})
		if len(result) == 128 {
			break
		}
	}
	return result, nil
}
func collectNetwork(ctx context.Context) NetworkSnapshot {
	s := NetworkSnapshot{At: time.Now().Unix(), Flows: []NetworkFlow{}}
	if runtime.GOOS != "linux" {
		s.Error = "Сбор conntrack поддерживается только на Linux"
		return s
	}
	// Proc contains both address families when available. Modern kernels often expose only netlink.
	for _, path := range []string{"/proc/net/nf_conntrack", "/proc/net/ip_conntrack"} {
		f, err := os.Open(path)
		if err != nil {
			continue
		}
		raw, err := io.ReadAll(io.LimitReader(f, (8<<20)+1))
		f.Close()
		if err != nil {
			s.Error = err.Error()
			return s
		}
		if len(raw) > 8<<20 {
			s.Error = "Таблица conntrack превышает лимит чтения 8 MiB"
			return s
		}
		appendNetworkFlows(&s, raw)
		return s
	}
	for _, family := range []string{"ipv4", "ipv6"} {
		raw, err := networkCommand(ctx, "conntrack", "-L", "-f", family, "-o", "extended")
		if err != nil {
			s.Error += "Conntrack " + family + ": нужен пакет conntrack и права root/CAP_NET_ADMIN. " + err.Error() + "; "
			continue
		}
		appendNetworkFlows(&s, raw)
	}
	return s
}
func startAgentNetwork(config agentConfig) func() {
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() {
		defer close(done)
		client := &http.Client{Timeout: 10 * time.Second}
		ticker := time.NewTicker(5 * time.Second)
		defer ticker.Stop()
		var firewall []FirewallCounter
		var firewallAt int64
		var firewallErr string
		lastFirewall := time.Time{}
		for {
			if ctx.Err() != nil {
				return
			}
			snapshot := NetworkSnapshot{At: time.Now().Unix(), Flows: []NetworkFlow{}}
			if message := ensureNetworkDependency(); message != "" {
				snapshot.Error = message
			} else {
				snapshot = collectNetwork(ctx)
			}
			if time.Since(lastFirewall) >= 30*time.Second {
				lastFirewall = time.Now()
				raw, err := networkCommand(ctx, "nft", "-j", "list", "ruleset")
				firewall = nil
				firewallAt = 0
				firewallErr = ""
				if err == nil {
					firewall, err = parseFirewallCounters(raw)
				}
				if err != nil {
					firewallErr = "Счётчики nftables недоступны: " + err.Error()
				} else if len(firewall) == 0 {
					firewallErr = "В правилах CrowdSec нет поддерживаемых счётчиков перед drop/reject. Это не означает отсутствие блокировок."
				} else {
					firewallAt = time.Now().Unix()
				}
			}
			snapshot.Firewall = firewall
			snapshot.FirewallAt = firewallAt
			snapshot.FirewallError = firewallErr
			data, _ := json.Marshal(snapshot)
			req, err := http.NewRequestWithContext(ctx, "POST", config.Server+"/api/agent/network", bytes.NewReader(data))
			if err != nil {
				return
			}
			req.Header.Set("Authorization", "Bearer "+config.ID+"."+config.Secret)
			req.Header.Set("Content-Type", "application/json")
			if res, err := client.Do(req); err == nil {
				io.Copy(io.Discard, io.LimitReader(res.Body, 4096))
				res.Body.Close()
			}
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
			}
		}
	}()
	return func() { cancel(); <-done }
}