// Copyright (c) 2020-2021 Tigera, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package commands

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/docopt/docopt-go"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"

	"github.com/projectcalico/calico/felix/bpf"
	"github.com/projectcalico/calico/felix/bpf/conntrack"
	v2 "github.com/projectcalico/calico/felix/bpf/conntrack/v2"
	v4 "github.com/projectcalico/calico/felix/bpf/conntrack/v4"
	"github.com/projectcalico/calico/felix/bpf/maps"
)

func init() {
	conntrackCmd.AddCommand(newConntrackDumpCmd())
	conntrackCmd.AddCommand(newConntrackRemoveCmd())
	conntrackCmd.AddCommand(newConntrackCleanCmd())
	conntrackCmd.AddCommand(newConntrackWriteCmd())
	conntrackCmd.AddCommand(newConntrackCreateCmd())
	conntrackCmd.AddCommand(newConntrackStatsCmd())
	rootCmd.AddCommand(conntrackCmd)
}

// conntrackCmd represents the conntrack command
var (
	conntrackCmd = &cobra.Command{
		Use:   "conntrack",
		Short: "Manipulates connection tracking",
	}

	voidIP4 = net.IPv4(0, 0, 0, 0)
	voidIP6 = net.ParseIP("::")
)

type conntrackOpts struct {
	version int
	ipv6    bool
}

type conntrackDumpCmd struct {
	*cobra.Command
	raw bool
	conntrackOpts
}

func newConntrackDumpCmd() *cobra.Command {
	cmd := &conntrackDumpCmd{
		Command: &cobra.Command{
			Use:   "dump [--ver=<version>] [--raw]",
			Short: "Dumps connection tracking table",
		},
	}

	cmd.Command.Flags().IntVarP((&cmd.version), "ver", "v", 4, "version to dump from")
	cmd.Command.Flags().BoolVar((&cmd.raw), "raw", false, "dump the raw conntrack table as is. For version < 3 it is always raw")
	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func dumpCtMapV2(ctMap maps.Map) error {
	err := ctMap.Iter(func(k, v []byte) maps.IteratorAction {
		var ctKey v2.Key
		if len(k) != len(ctKey) {
			log.Panic("Key has unexpected length")
		}
		copy(ctKey[:], k[:])

		var ctVal v2.Value
		if len(v) != len(ctVal) {
			log.Panic("Value has unexpected length")
		}
		copy(ctVal[:], v[:])

		fmt.Printf("%v -> %v", ctKey, ctVal)
		dumpExtrav2(ctKey, ctVal)
		fmt.Printf("\n")
		return maps.IterNone
	})
	return err
}

func (cmd *conntrackDumpCmd) Run(c *cobra.Command, _ []string) {
	var (
		ctMap maps.Map
		ctErr error
	)

	cmd.ipv6 = ipv6 != nil && *ipv6
	if cmd.version < 3 && cmd.version != 0 {
		cmd.raw = true
	}

	if ctMap, ctErr = GetCTMap(cmd.version, cmd.ipv6); ctErr != nil {
		log.WithError(ctErr).Fatal("Failed to get ConntrackMap")
	}
	if err := ctMap.Open(); err != nil {
		log.WithError(err).Fatal("Failed to access ConntrackMap")
	}
	defer ctMap.Close()
	if cmd.version == 2 {
		err := dumpCtMapV2(ctMap)
		if err != nil {
			log.WithError(err).Fatal("Failed to iterate over conntrack entries")
		}
		return
	}

	keyFromBytes := conntrack.KeyFromBytes
	valFromBytes := conntrack.ValueFromBytes
	if cmd.ipv6 {
		keyFromBytes = conntrack.KeyV6FromBytes
		valFromBytes = conntrack.ValueV6FromBytes
	}

	if *jsonOutput {
		if err := cmd.dumpJSON(ctMap, keyFromBytes, valFromBytes); err != nil {
			log.WithError(err).Error("Failed to dump conntrack as JSON")
		}
		return
	}

	err := ctMap.Iter(func(k, v []byte) maps.IteratorAction {
		ctKey := keyFromBytes(k)
		ctVal := valFromBytes(v)

		if cmd.raw {
			fmt.Printf("%v -> %v", ctKey, ctVal)
			dumpExtra(ctKey, ctVal)
			fmt.Printf("\n")
		} else {
			cmd.prettyDump(ctKey, ctVal)
		}
		return maps.IterNone
	})
	if err != nil {
		log.WithError(err).Fatal("Failed to iterate over conntrack entries")
	}
}

func protoStr(proto uint8) string {
	switch proto {
	case 6:
		return "TCP"
	case 17:
		return "UDP"
	case 1:
		return "ICMP"
	case 58:
		return "ICMP6"
	}

	return "UNKNOWN"
}

func (cmd *conntrackDumpCmd) prettyDump(k conntrack.KeyInterface, v conntrack.ValueInterface) {
	d := v.Data()

	switch v.Type() {
	case conntrack.TypeNormal:
		if v.Flags()&v4.FlagSrcDstBA != 0 {
			cmd.Printf("%s %s:%d -> %s:%d ", protoStr(k.Proto()), k.AddrB(), k.PortB(), k.AddrA(), k.PortA())
		} else {
			cmd.Printf("%s %s:%d -> %s:%d ", protoStr(k.Proto()), k.AddrA(), k.PortA(), k.AddrB(), k.PortB())
		}
	case conntrack.TypeNATForward:
		return
	case conntrack.TypeNATReverse:
		if v.Flags()&v4.FlagSrcDstBA != 0 {
			cmd.Printf("%s %s:%d -> %s:%d -> %s:%d ",
				protoStr(k.Proto()), k.AddrB(), k.PortB(), d.OrigDst, d.OrigPort, k.AddrA(), k.PortA())
		} else {
			cmd.Printf("%s %s:%d -> %s:%d -> %s:%d ",
				protoStr(k.Proto()), k.AddrA(), k.PortA(), d.OrigDst, d.OrigPort, k.AddrB(), k.PortB())
		}

		if (cmd.ipv6 && !d.TunIP.Equal(voidIP6)) || (!cmd.ipv6 && !d.TunIP.Equal(voidIP4)) {
			cmd.Printf("external client, service forwarded to/from %s ", d.TunIP)
		}
	}

	if v.Flags()&v4.FlagHostPSNAT != 0 {
		cmd.Printf("source port changed from %d ", d.OrigSPort)
	}

	now := bpf.KTimeNanos()
	cmd.Printf(" Active ago %s", time.Duration(now-v.LastSeen()))

	if k.Proto() == 6 {
		if (v.IsForwardDSR() && d.FINsSeenDSR()) || d.FINsSeen() {
			cmd.Printf(" CLOSED")
		} else if d.RSTSeen() {
			cmd.Printf(" RESET")
		} else if d.Established() {
			cmd.Printf(" ESTABLISHED")
		} else {
			cmd.Printf(" SYN-SENT")
		}
	}

	cmd.Printf("\n")
}

// ctEntryJSON is a conntrack entry in raw JSON mode — one object per map entry.
type ctEntryJSON struct {
	Proto     string   `json:"proto"`
	AddrA     string   `json:"addr_a"`
	PortA     uint16   `json:"port_a"`
	AddrB     string   `json:"addr_b"`
	PortB     uint16   `json:"port_b"`
	Type      string   `json:"type"`
	Flags     []string `json:"flags"`
	ActiveAgo string   `json:"active_ago"`
	TCPState  string   `json:"tcp_state,omitempty"`
	// For nat-forward entries: the key of the reverse entry this points to.
	RevProto string `json:"rev_proto,omitempty"`
	RevAddrA string `json:"rev_addr_a,omitempty"`
	RevPortA uint16 `json:"rev_port_a,omitempty"`
	RevAddrB string `json:"rev_addr_b,omitempty"`
	RevPortB uint16 `json:"rev_port_b,omitempty"`
}

// ctConnectionJSON is a logical connection in pretty JSON mode.
// For NAT connections, forward and reverse entries are grouped.
type ctConnectionJSON struct {
	Type        string   `json:"type"`
	Proto       string   `json:"proto"`
	Src         string   `json:"src"`
	Dst         string   `json:"dst"`
	OrigDst     string   `json:"orig_dst,omitempty"`
	TunnelIP    string   `json:"tunnel_ip,omitempty"`
	OrigSrcPort uint16   `json:"orig_src_port,omitempty"`
	Flags       []string `json:"flags"`
	ActiveAgo   string   `json:"active_ago"`
	TCPState    string   `json:"tcp_state,omitempty"`
}

func ctTypeStr(t uint8) string {
	switch t {
	case conntrack.TypeNormal:
		return "normal"
	case conntrack.TypeNATForward:
		return "nat-forward"
	case conntrack.TypeNATReverse:
		return "nat-reverse"
	}
	return fmt.Sprintf("unknown(%d)", t)
}

func tcpStateStr(k conntrack.KeyInterface, v conntrack.ValueInterface) string {
	if k.Proto() != 6 {
		return ""
	}
	d := v.Data()
	if (v.IsForwardDSR() && d.FINsSeenDSR()) || d.FINsSeen() {
		return "CLOSED"
	}
	if d.RSTSeen() {
		return "RESET"
	}
	if d.Established() {
		return "ESTABLISHED"
	}
	return "SYN-SENT"
}

func (cmd *conntrackDumpCmd) dumpJSON(
	ctMap maps.Map,
	keyFromBytes func([]byte) conntrack.KeyInterface,
	valFromBytes func([]byte) conntrack.ValueInterface,
) error {
	now := bpf.KTimeNanos()

	if cmd.raw {
		return cmd.dumpRawJSON(ctMap, keyFromBytes, valFromBytes, now)
	}
	return cmd.dumpPrettyJSON(ctMap, keyFromBytes, valFromBytes, now)
}

func (cmd *conntrackDumpCmd) dumpRawJSON(
	ctMap maps.Map,
	keyFromBytes func([]byte) conntrack.KeyInterface,
	valFromBytes func([]byte) conntrack.ValueInterface,
	now int64,
) error {
	var entries []ctEntryJSON
	err := ctMap.Iter(func(k, v []byte) maps.IteratorAction {
		ctKey := keyFromBytes(k)
		ctVal := valFromBytes(v)
		entry := ctEntryJSON{
			Proto:     protoStr(ctKey.Proto()),
			AddrA:     ctKey.AddrA().String(),
			PortA:     ctKey.PortA(),
			AddrB:     ctKey.AddrB().String(),
			PortB:     ctKey.PortB(),
			Type:      ctTypeStr(ctVal.Type()),
			Flags:     v4.FlagNames(ctVal.Flags()),
			ActiveAgo: time.Duration(now - ctVal.LastSeen()).String(),
			TCPState:  tcpStateStr(ctKey, ctVal),
		}
		if ctVal.Type() == conntrack.TypeNATForward {
			rk := ctVal.ReverseNATKey()
			entry.RevProto = protoStr(rk.Proto())
			entry.RevAddrA = rk.AddrA().String()
			entry.RevPortA = rk.PortA()
			entry.RevAddrB = rk.AddrB().String()
			entry.RevPortB = rk.PortB()
		}
		entries = append(entries, entry)
		return maps.IterNone
	})
	if err != nil {
		return fmt.Errorf("failed to iterate over conntrack entries: %w", err)
	}

	enc := json.NewEncoder(cmd.OutOrStdout())
	enc.SetIndent("", "  ")
	return enc.Encode(entries)
}

func (cmd *conntrackDumpCmd) dumpPrettyJSON(
	ctMap maps.Map,
	keyFromBytes func([]byte) conntrack.KeyInterface,
	valFromBytes func([]byte) conntrack.ValueInterface,
	now int64,
) error {
	type ctEntry struct {
		key conntrack.KeyInterface
		val conntrack.ValueInterface
	}

	// Collect all entries; skip NAT forward entries since the reverse
	// entry already contains the full connection information.
	var normals []ctEntry
	var natReverses []ctEntry

	err := ctMap.Iter(func(k, v []byte) maps.IteratorAction {
		ctKey := keyFromBytes(k)
		ctVal := valFromBytes(v)
		e := ctEntry{key: ctKey, val: ctVal}

		switch ctVal.Type() {
		case conntrack.TypeNormal:
			normals = append(normals, e)
		case conntrack.TypeNATReverse:
			natReverses = append(natReverses, e)
		}
		return maps.IterNone
	})
	if err != nil {
		return fmt.Errorf("failed to iterate over conntrack entries: %w", err)
	}

	var connections []ctConnectionJSON

	// Emit normal connections.
	for _, e := range normals {
		k, v := e.key, e.val
		src, dst := orientedAddrs(k, v)
		connections = append(connections, ctConnectionJSON{
			Type:      "normal",
			Proto:     protoStr(k.Proto()),
			Src:       src,
			Dst:       dst,
			Flags:     v4.FlagNames(v.Flags()),
			ActiveAgo: time.Duration(now - v.LastSeen()).String(),
			TCPState:  tcpStateStr(k, v),
		})
	}

	// Emit NAT connections from reverse entries.
	for _, e := range natReverses {
		k, v := e.key, e.val
		d := v.Data()
		src, dst := orientedAddrs(k, v)
		origDst := net.JoinHostPort(d.OrigDst.String(), fmt.Sprint(d.OrigPort))

		conn := ctConnectionJSON{
			Type:      "nat",
			Proto:     protoStr(k.Proto()),
			Src:       src,
			OrigDst:   origDst,
			Dst:       dst,
			Flags:     v4.FlagNames(v.Flags()),
			ActiveAgo: time.Duration(now - v.LastSeen()).String(),
			TCPState:  tcpStateStr(k, v),
		}

		if (cmd.ipv6 && !d.TunIP.Equal(voidIP6)) || (!cmd.ipv6 && !d.TunIP.Equal(voidIP4)) {
			conn.TunnelIP = d.TunIP.String()
		}
		if v.Flags()&v4.FlagHostPSNAT != 0 {
			conn.OrigSrcPort = d.OrigSPort
		}

		connections = append(connections, conn)
	}

	enc := json.NewEncoder(cmd.OutOrStdout())
	enc.SetIndent("", "  ")
	return enc.Encode(connections)
}

// orientedAddrs returns src and dst as "ip:port" strings, respecting
// the FlagSrcDstBA direction flag.
func orientedAddrs(k conntrack.KeyInterface, v conntrack.ValueInterface) (src, dst string) {
	if v.Flags()&v4.FlagSrcDstBA != 0 {
		src = net.JoinHostPort(k.AddrB().String(), fmt.Sprint(k.PortB()))
		dst = net.JoinHostPort(k.AddrA().String(), fmt.Sprint(k.PortA()))
	} else {
		src = net.JoinHostPort(k.AddrA().String(), fmt.Sprint(k.PortA()))
		dst = net.JoinHostPort(k.AddrB().String(), fmt.Sprint(k.PortB()))
	}
	return
}

func dumpExtrav2(k v2.Key, v v2.Value) {
	now := bpf.KTimeNanos()

	fmt.Printf(" Active ago %s", time.Duration(now-v.LastSeen()))

	if k.Proto() != conntrack.ProtoTCP {
		return
	}

	if v.Type() == conntrack.TypeNATForward {
		return
	}

	data := v.Data()

	if (v.IsForwardDSR() && data.FINsSeenDSR()) || data.FINsSeen() {
		fmt.Printf(" CLOSED")
		return
	}

	if data.RSTSeen() {
		fmt.Printf(" RESET")
		return
	}

	if data.Established() {
		fmt.Printf(" ESTABLISHED")
		return
	}

	fmt.Printf(" SYN-SENT")
}

func dumpExtra(k conntrack.KeyInterface, v conntrack.ValueInterface) {
	now := bpf.KTimeNanos()

	fmt.Printf(" Active ago %s", time.Duration(now-v.LastSeen()))

	if k.Proto() != conntrack.ProtoTCP {
		return
	}

	if v.Type() == conntrack.TypeNATForward {
		return
	}

	data := v.Data()

	if (v.IsForwardDSR() && data.FINsSeenDSR()) || data.FINsSeen() {
		fmt.Printf(" CLOSED")
		return
	}

	if data.RSTSeen() {
		fmt.Printf(" RESET")
		return
	}

	if data.Established() {
		fmt.Printf(" ESTABLISHED")
		return
	}

	fmt.Printf(" SYN-SENT")
}

type conntrackRemoveCmd struct {
	*cobra.Command

	Proto string `docopt:"<proto>"`
	IP1   string `docopt:"<ip1>"`
	IP2   string `docopt:"<ip2>"`

	proto uint8
	ip1   net.IP
	ip2   net.IP

	conntrackOpts
}

func newConntrackRemoveCmd() *cobra.Command {
	cmd := &conntrackRemoveCmd{
		Command: &cobra.Command{
			Use:   "remove <proto> <ip1> <ip2>",
			Short: "Removes connection tracking",
		},
	}

	cmd.Command.Flags().IntVarP((&cmd.version), "ver", "v", 4, "version to remove from")
	cmd.Command.Args = cmd.Args
	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func (cmd *conntrackRemoveCmd) Args(c *cobra.Command, args []string) error {
	a, err := docopt.ParseArgs(makeDocUsage(c), args, "")
	if err != nil {
		return err
	}

	err = a.Bind(cmd)
	if err != nil {
		return err
	}

	switch proto := strings.ToLower(args[0]); proto {
	case "udp":
		cmd.proto = 17
	case "tcp":
		cmd.proto = 6
	default:
		return fmt.Errorf("unknown protocol %s", proto)
	}

	cmd.ip1 = net.ParseIP(cmd.IP1)
	if cmd.ip1 == nil {
		return fmt.Errorf("ip1: %q is not an ip", cmd.IP1)
	}

	cmd.ip2 = net.ParseIP(cmd.IP2)
	if cmd.ip2 == nil {
		return fmt.Errorf("ip2: %q is not an ip", cmd.IP2)
	}

	return nil
}

func (cmd *conntrackRemoveCmd) Run(c *cobra.Command, _ []string) {
	var (
		ctMap maps.Map
		ctErr error
	)

	cmd.ipv6 = ipv6 != nil && *ipv6

	if ctMap, ctErr = GetCTMap(cmd.version, cmd.ipv6); ctErr != nil {
		log.WithError(ctErr).Fatal("Failed to get ConntrackMap")
	}
	if err := ctMap.Open(); err != nil {
		log.WithError(err).Error("Failed to access ConntrackMap")
	}
	err := ctMap.Iter(func(k, v []byte) maps.IteratorAction {
		var ctKey conntrack.Key
		if len(k) != len(ctKey) {
			log.Panic("Key has unexpected length")
		}
		copy(ctKey[:], k[:])

		log.Infof("Examining conntrack key: %v", ctKey)

		if ctKey.Proto() != cmd.proto {
			return maps.IterNone
		}

		if ctKey.AddrA().Equal(cmd.ip1) && ctKey.AddrB().Equal(cmd.ip2) {
			log.Info("Match")
			return maps.IterDelete
		} else if ctKey.AddrB().Equal(cmd.ip1) && ctKey.AddrA().Equal(cmd.ip2) {
			log.Info("Match")
			return maps.IterDelete
		}
		return maps.IterNone
	})
	if err != nil {
		log.WithError(err).Fatal("Failed to iterate over conntrack entries")
	}
}

type conntrackCleanCmd struct {
	*cobra.Command
	conntrackOpts
}

func newConntrackCleanCmd() *cobra.Command {
	cmd := &conntrackCleanCmd{
		Command: &cobra.Command{
			Use:   "clean [--ver=<version>]",
			Short: "Cleans all conntrack entries",
		},
	}

	cmd.Command.Flags().IntVarP((&cmd.version), "ver", "v", 4, "conntrack version to clean")
	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func (cmd *conntrackCleanCmd) Run(c *cobra.Command, _ []string) {
	var (
		ctMap maps.Map
		ctErr error
	)

	cmd.ipv6 = ipv6 != nil && *ipv6

	if ctMap, ctErr = GetCTMap(cmd.version, cmd.ipv6); ctErr != nil {
		log.WithError(ctErr).Fatal("Failed to get ConntrackMap")
	}
	if err := ctMap.Open(); err != nil {
		log.WithError(err).Error("Failed to access ConntrackMap")
	}

	// Disable debug if set while deleting
	loglevel := log.GetLevel()
	log.SetLevel(log.WarnLevel)
	err := ctMap.Iter(func(k, v []byte) maps.IteratorAction {
		return maps.IterDelete
	})

	log.SetLevel(loglevel)
	if err != nil {
		log.WithError(err).Fatal("Failed to iterate over conntrack entries")
	}
}

type conntrackCreateCmd struct {
	*cobra.Command
	conntrackOpts
}

func newConntrackCreateCmd() *cobra.Command {
	cmd := &conntrackCreateCmd{
		Command: &cobra.Command{
			Use:   "create [--ver=<version>]",
			Short: "Creates a conntrack map of specified version",
		},
	}

	cmd.Command.Flags().IntVarP((&cmd.version), "ver", "v", 4, "conntrack version to create")
	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func (cmd *conntrackCreateCmd) Run(c *cobra.Command, _ []string) {
	cmd.ipv6 = ipv6 != nil && *ipv6
	ctMap := getCTMap(cmd.version, cmd.ipv6)

	if err := ctMap.EnsureExists(); err != nil {
		log.WithError(err).Errorf("Failed to create conntrackMap version %d", cmd.version)
	}
}

type conntrackWriteCmd struct {
	*cobra.Command

	Key   string `docopt:"<key>"`
	Value string `docopt:"<value>"`

	key []byte
	val []byte
	conntrackOpts
}

func newConntrackWriteCmd() *cobra.Command {
	cmd := &conntrackWriteCmd{
		Command: &cobra.Command{
			Use:   "write <key> <value>",
			Short: "Writes a key-value pair, each encoded in base64",
		},
	}

	cmd.Command.Flags().IntVarP((&cmd.version), "ver", "v", 4, "conntrack map version")
	cmd.Command.Args = cmd.Args
	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func (cmd *conntrackWriteCmd) Args(c *cobra.Command, args []string) error {
	a, err := docopt.ParseArgs(makeDocUsage(c), args, "")
	if err != nil {
		return err
	}

	err = a.Bind(cmd)
	if err != nil {
		return err
	}

	cmd.ipv6 = ipv6 != nil && *ipv6

	cmd.key, err = base64.StdEncoding.DecodeString(cmd.Key)
	if err != nil {
		switch cmd.version {
		case 2:
			if len(cmd.key) != len(v2.Key{}) {
				return fmt.Errorf("failed to decode key: %s", err)
			}
		default:
			if cmd.ipv6 && len(cmd.key) != len(conntrack.KeyV6{}) {
				return fmt.Errorf("failed to decode key: %s", err)
			} else if !cmd.ipv6 && len(cmd.key) != len(conntrack.Key{}) {
				return fmt.Errorf("failed to decode key: %s", err)
			}
		}
	}

	cmd.val, err = base64.StdEncoding.DecodeString(cmd.Value)
	if err != nil {
		switch cmd.version {
		case 2:
			if len(cmd.val) != len(v2.Value{}) {
				return fmt.Errorf("failed to decode val: %s", err)
			}
		default:
			if cmd.ipv6 && len(cmd.val) != len(conntrack.ValueV6{}) {
				return fmt.Errorf("failed to decode val: %s", err)
			} else if !cmd.ipv6 && len(cmd.val) != len(conntrack.Value{}) {
				return fmt.Errorf("failed to decode val: %s", err)
			}
		}
	}
	return nil
}

func (cmd *conntrackWriteCmd) Run(c *cobra.Command, _ []string) {
	var (
		ctMap maps.Map
		ctErr error
	)

	cmd.ipv6 = ipv6 != nil && *ipv6

	if ctMap, ctErr = GetCTMap(cmd.version, cmd.ipv6); ctErr != nil {
		log.WithError(ctErr).Fatal("Failed to get ConntrackMap")
	}

	if err := ctMap.Open(); err != nil {
		log.WithError(err).Error("Failed to access ConntrackMap")
	}

	if err := ctMap.Update(cmd.key, cmd.val); err != nil {
		log.WithError(err).Error("Failed to update ConntrackMap")
	}
}

type conntrackStatsCmd struct {
	*cobra.Command

	established int
	reset       int
	closed      int
	synSent     int
	total       int
	nat         int

	protos map[int]int

	conntrackOpts
}

func newConntrackStatsCmd() *cobra.Command {
	cmd := &conntrackStatsCmd{
		Command: &cobra.Command{
			Use:   "stats [--ver=<version>]",
			Short: "Prints conntrack statistics",
		},
		protos: make(map[int]int),
	}

	cmd.Command.Flags().IntVarP((&cmd.version), "ver", "v", 4, "conntrack map version")

	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func (cmd *conntrackStatsCmd) Run(c *cobra.Command, _ []string) {
	var (
		ctMap maps.Map
		ctErr error
	)

	cmd.ipv6 = ipv6 != nil && *ipv6

	if ctMap, ctErr = GetCTMap(cmd.version, cmd.ipv6); ctErr != nil {
		log.WithError(ctErr).Fatal("Failed to get ConntrackMap")
	}

	if err := ctMap.Open(); err != nil {
		log.WithError(err).Fatal("Failed to access ConntrackMap")
	}

	keyFromBytes := conntrack.KeyFromBytes
	valFromBytes := conntrack.ValueFromBytes
	if cmd.ipv6 {
		keyFromBytes = conntrack.KeyV6FromBytes
		valFromBytes = conntrack.ValueV6FromBytes
	}

	err := ctMap.Iter(func(k, v []byte) maps.IteratorAction {
		ctKey := keyFromBytes(k)
		ctVal := valFromBytes(v)
		d := ctVal.Data()

		if ctVal.Type() == conntrack.TypeNATForward {
			cmd.nat++
			return maps.IterNone
		}

		if ctKey.Proto() == 6 {
			if (ctVal.IsForwardDSR() && d.FINsSeenDSR()) || d.FINsSeen() {
				cmd.closed++
			} else if d.RSTSeen() {
				cmd.reset++
			} else if d.Established() {
				cmd.established++
			} else {
				cmd.synSent++
			}
		}

		cmd.total++
		cmd.protos[int(ctKey.Proto())]++

		return maps.IterNone
	})

	if *jsonOutput {
		type ctStatsJSON struct {
			MapSize        int `json:"map_size"`
			TotalConns     int `json:"total_connections"`
			TotalEntries   int `json:"total_entries"`
			NATConns       int `json:"nat_connections"`
			TCP            int `json:"tcp"`
			UDP            int `json:"udp"`
			Others         int `json:"others"`
			TCPEstablished int `json:"tcp_established"`
			TCPClosed      int `json:"tcp_closed"`
			TCPReset       int `json:"tcp_reset"`
			TCPSynSent     int `json:"tcp_syn_sent"`
		}
		stats := ctStatsJSON{
			MapSize:        maps.Size(ctMap.GetName()),
			TotalConns:     cmd.total,
			TotalEntries:   cmd.total + cmd.nat,
			NATConns:       cmd.nat,
			TCP:            cmd.protos[6],
			UDP:            cmd.protos[17],
			Others:         cmd.total - cmd.protos[6] - cmd.protos[17],
			TCPEstablished: cmd.established,
			TCPClosed:      cmd.closed,
			TCPReset:       cmd.reset,
			TCPSynSent:     cmd.synSent,
		}
		enc := json.NewEncoder(cmd.OutOrStdout())
		enc.SetIndent("", "  ")
		if encErr := enc.Encode(stats); encErr != nil {
			log.WithError(encErr).Error("Failed to encode JSON")
		}
	} else {
		cmd.Printf("Conntrack map size: %d\n", maps.Size(ctMap.GetName()))

		cmd.Printf("Total connections: %d\n", cmd.total)
		cmd.Printf("Total entries: %d\n", cmd.total+cmd.nat)
		cmd.Printf("NAT connections: %d\n\n", cmd.nat)

		cmd.Printf("TCP : %d\n", cmd.protos[6])
		cmd.Printf("UDP : %d\n", cmd.protos[17])
		cmd.Printf("Others : %d\n\n", cmd.total-cmd.protos[6]-cmd.protos[17])

		cmd.Printf("TCP Established: %d\n", cmd.established)
		cmd.Printf("TCP Closed: %d\n", cmd.closed)
		cmd.Printf("TCP Reset: %d\n", cmd.reset)
		cmd.Printf("TCP Syn-sent: %d\n", cmd.synSent)
	}

	if err != nil {
		log.WithError(err).Fatal("Failed to iterate over conntrack entries")
	}
}

func GetCTMap(version int, ipv6 bool) (maps.Map, error) {
	// Set the map size based on the actual max entries obtained from the map info.
	if err := setCTMapSize(version, ipv6); err != nil {
		return nil, err
	}
	return getCTMap(version, ipv6), nil
}

func getCTMap(version int, ipv6 bool) maps.Map {
	var ctMap maps.Map
	switch version {
	case 2:
		ctMap = conntrack.MapV2()
	default:
		if ipv6 {
			ctMap = conntrack.MapV6()
		} else {
			ctMap = conntrack.Map()
		}
	}
	return ctMap
}

func setCTMapSize(version int, ipv6 bool) error {
	ctMap := getCTMap(version, ipv6)

	if err := ctMap.Open(); err != nil {
		return errors.New("failed to access conntrack Map")
	}
	defer ctMap.Close()

	if mapInfo, err := maps.GetMapInfo(ctMap.MapFD()); err != nil {
		return errors.New("failed to get map info")
	} else {
		// Set the map size based on the actual max entries obtained from the map info.
		maps.SetSize(ctMap.GetName(), mapInfo.MaxEntries)
	}

	return nil
}
