package main

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/spf13/cobra"

	"housecarl/internal/api"
)

// Traffic as it happens: connections the kernel tracks, and tracing
// packets through the ruleset.

func connectionsCmd() *cobra.Command {
	var search string
	var limit int
	cmd := &cobra.Command{
		Use: "connections", Aliases: []string{"conns"}, Short: "The connections the firewall tracks, and closing them",
		Long: `Lists the connections the kernel tracks: what the "established" built-in
lets through without checking the rules again. A rule added later doesn't
stop them by itself; close them, and their next packet is checked by the
rules as it stands.

  hc connections --search 192.168.1.50
  hc connections --search ":443 tcp"
  hc connections close --search 192.168.1.50`,
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			cs, err := get[[]api.Connection]("/connections", q("q", search, "limit", strconv.Itoa(limit)))
			if err != nil {
				return err
			}
			return output(cs, func(w io.Writer) {
				fmt.Fprintln(w, "PROTO\tFROM\tTO\tNAT\tSTATE\tEXPIRES")
				for _, c := range cs {
					fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n", c.Protocol, hostPort(c.Source, c.SourcePort), hostPort(c.Destination, c.DestinationPort), natText(c), c.State, fmtSeconds(c.Timeout))
				}
			})
		},
	}
	cmd.Flags().StringVar(&search, "search", "", "only those with all these words: addresses, :port, protocol, state")
	cmd.Flags().IntVar(&limit, "limit", 500, "at most this many")
	var csearch string
	var yes bool
	closeCmd := &cobra.Command{
		Use: "close --search TEXT", Short: "Close connections: the rules check them again with their next packet", Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			if strings.TrimSpace(csearch) == "" {
				return fmt.Errorf("say which: --search 192.168.1.50 (see 'hc connections --search' first)")
			}
			if !confirm("Close every connection matching "+strconvQuote(csearch)+"?", yes) {
				return fmt.Errorf("nothing closed")
			}
			c, err := newClient()
			if err != nil {
				return err
			}
			var out map[string]int
			if err := c.Do(ctx(), "POST", "/connections/close", nil, api.CloseConnections{Search: csearch}, &out); err != nil {
				return err
			}
			if !flagQuiet {
				fmt.Printf("closed %s\n", plural(out["closed"], "connection"))
			}
			return nil
		},
	}
	closeCmd.Flags().StringVar(&csearch, "search", "", "the connections to close, as for 'hc connections --search'")
	closeCmd.Flags().BoolVarP(&yes, "yes", "y", false, "don't ask")
	cmd.AddCommand(closeCmd)
	return cmd
}

func strconvQuote(s string) string { return strconv.Quote(s) }

func hostPort(addr string, port int) string {
	if port == 0 {
		return addr
	}
	if strings.Contains(addr, ":") {
		return fmt.Sprintf("[%s]:%d", addr, port)
	}
	return fmt.Sprintf("%s:%d", addr, port)
}

// natText: "to 192.168.1.10:443" (a port forward), "as 203.0.113.10"
// (masqueraded).
func natText(c api.Connection) string {
	var parts []string
	if c.DNATTo != "" {
		port := c.DNATPort
		if port == 0 {
			port = c.DestinationPort
		}
		parts = append(parts, "to "+hostPort(c.DNATTo, port))
	}
	if c.SNATTo != "" {
		parts = append(parts, "as "+c.SNATTo)
	}
	return strings.Join(parts, ", ")
}

func traceCmd() *cobra.Command {
	var in api.TraceInput
	cmd := &cobra.Command{
		Use: "trace", Short: "Follow packets through every table's rules: what each met, and what decided",
		Long: `Traces packets from or to an address, on a port (either way), through the
whole ruleset, Housecarl's and every other table's, for a while: each
rule and policy they meet, and what accepted or dropped them. Ctrl-C
stops it.

  hc trace --address 192.168.1.50 --port 443
  hc trace --port 22 --seconds 120
  hc trace --address 2001:db8:1::/64 --protocol icmpv6`,
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			c, err := newClient()
			if err != nil {
				return err
			}
			var st api.TraceStatus
			if err := c.Do(ctx(), "POST", "/trace", nil, in, &st); err != nil {
				return err
			}
			if !flagQuiet {
				fmt.Fprintf(os.Stderr, "tracing until %s (Ctrl-C stops)\n", st.Ends.Local().Format("15:04:05"))
			}
			sig, stop := signal.NotifyContext(context.Background(), os.Interrupt)
			defer stop()
			shown := 0
			for {
				select {
				case <-sig.Done():
					c.Do(ctx(), "DELETE", "/trace", nil, nil, &st)
				case <-time.After(time.Second):
					if err := c.Do(ctx(), "GET", "/trace", nil, nil, &st); err != nil {
						return err
					}
				}
				for ; shown < len(st.Packets); shown++ {
					printPacket(os.Stdout, st.Packets[shown])
				}
				if !st.Running || sig.Err() != nil {
					break
				}
			}
			if st.Error != "" {
				return fmt.Errorf("the trace failed: %s", st.Error)
			}
			if !flagQuiet {
				fmt.Fprintf(os.Stderr, "%d packets traced\n", len(st.Packets))
			}
			return nil
		},
	}
	f := cmd.Flags()
	f.StringVar(&in.Address, "address", "", "packets from or to this address or prefix")
	f.IntVar(&in.Port, "port", 0, "packets from or to this port")
	f.StringVar(&in.Protocol, "protocol", "", "tcp, udp, sctp, icmp or icmpv6 (default any)")
	f.IntVar(&in.Seconds, "seconds", 60, "how long (5 to 600)")
	return cmd
}

// printPacket writes a traced packet: its headers, each step, its fate.
func printPacket(w io.Writer, p api.TracePacket) {
	what := p.Protocol
	if p.TCPFlags != "" {
		what += " " + p.TCPFlags
	}
	if p.ICMPType != "" {
		what += " " + p.ICMPType
	}
	where := ""
	if p.In != "" {
		where += " in " + p.In
	}
	if p.Out != "" {
		where += " out " + p.Out
	}
	fmt.Fprintf(w, "%s  %s -> %s  %s%s\n", p.Time.Local().Format("15:04:05.000"), hostPort(p.Source, p.SourcePort), hostPort(p.Destination, p.DestinationPort), what, where)
	for _, s := range p.Steps {
		name := s.Label
		if name == "" {
			if s.Kind == "policy" {
				name = s.Chain + " policy"
			} else {
				name = s.Rule
			}
		}
		fmt.Fprintf(w, "    %s %s: %s -> %s\n", s.Table, s.Chain, name, s.Verdict)
	}
	decided := p.DecidedBy.Label
	if decided == "" {
		decided = strings.TrimSpace(p.DecidedBy.Table + " " + p.DecidedBy.Chain + " " + p.DecidedBy.Kind)
	}
	if decided == "" {
		fmt.Fprintf(w, "    => %s\n", strings.ToUpper(p.Verdict))
	} else {
		fmt.Fprintf(w, "    => %s by %s\n", strings.ToUpper(p.Verdict), decided)
	}
	if wd := p.WouldDrop; wd != nil {
		name := wd.Label
		if name == "" {
			name = strings.TrimSpace(wd.Chain + " " + wd.Mark)
		}
		fmt.Fprintf(w, "    => shadow mode: Housecarl would have dropped it, by %s\n", name)
	}
}

func matrixCmd() *cobra.Command {
	return &cobra.Command{
		Use: "matrix [FROM TO]", Short: "What the rules do to new connections between zones",
		Long: `Shows, for each pair of zones (and the firewall itself), what happens to a
new connection: allowed or blocked, decided by a rule or the policy, and
how many rules decide for part of the traffic first (+N: exceptions,
such as some services or addresses). With FROM and TO (a zone, or
"firewall"), that pair's exceptions in order, and what decides the rest.

  hc matrix
  hc matrix guest wan
  hc matrix lan firewall`,
		Args: cobra.RangeArgs(0, 2),
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 1 {
				return fmt.Errorf("give both: FROM and TO")
			}
			m, err := get[api.Matrix]("/matrix", nil)
			if err != nil {
				return err
			}
			if len(args) == 2 {
				for _, c := range m.Cells {
					if strings.EqualFold(c.From, args[0]) && strings.EqualFold(c.To, args[1]) {
						return output(c, func(w io.Writer) { printCell(w, c) })
					}
				}
				return fmt.Errorf("no pair %s → %s (zones: %s, and firewall)", args[0], args[1], strings.Join(m.Zones, ", "))
			}
			return output(m, func(w io.Writer) { printMatrix(w, m) })
		},
	}
}

func printMatrix(w io.Writer, m api.Matrix) {
	writeMatrix(w, m, isTerminal(os.Stdout) && flagOutput == "table")
}

// writeMatrix writes the grid, in colour on a terminal.
func writeMatrix(w io.Writer, m api.Matrix, color bool) {
	cols := append(append([]string{}, m.Zones...), "firewall")
	cells := map[[2]string]api.MatrixCell{}
	for _, c := range m.Cells {
		cells[[2]string{c.From, c.To}] = c
	}
	// The grid is padded by hand: tabwriter would count the colours'
	// escape codes as text, and misalign the columns.
	grid := [][]string{append([]string{"FROM \\ TO"}, cols...)}
	kinds := [][]string{make([]string, len(cols)+1)}
	same := false
	for _, from := range cols {
		row, kind := []string{from}, []string{""}
		for _, to := range cols {
			c, ok := cells[[2]string{from, to}]
			if !ok {
				row, kind = append(row, "-"), append(kind, "")
				continue
			}
			text := map[string]string{"allowed": "allow", "blocked": "block", "partly": "partly"}[c.Summary]
			if c.Allowed > 0 {
				text += fmt.Sprintf(" +%d", c.Allowed)
			}
			if c.Blocked > 0 {
				text += fmt.Sprintf(" -%d", c.Blocked)
			}
			if c.SameZone {
				text += " *"
				same = true
			}
			row, kind = append(row, text), append(kind, c.Summary)
		}
		grid, kinds = append(grid, row), append(kinds, kind)
	}
	width := make([]int, len(cols)+1)
	for _, row := range grid {
		for i, t := range row {
			width[i] = max(width[i], utf8.RuneCountInString(t))
		}
	}
	paint := map[string]string{"allowed": "\x1b[32m", "blocked": "\x1b[31m", "partly": "\x1b[33m"}
	for r, row := range grid {
		var b strings.Builder
		for i, t := range row {
			pad := strings.Repeat(" ", width[i]-utf8.RuneCountInString(t))
			if i == len(row)-1 {
				pad = ""
			} else {
				pad += "  "
			}
			if p := paint[kinds[r][i]]; color && p != "" {
				t = p + t + "\x1b[0m"
			}
			b.WriteString(t + pad)
		}
		fmt.Fprintln(w, b.String())
	}
	// What the partly open pairs let through.
	var open [][2]string
	for _, from := range cols {
		for _, to := range cols {
			c, ok := cells[[2]string{from, to}]
			if !ok || len(c.Open) == 0 || c.SameZone {
				continue
			}
			items := c.Open
			more := ""
			if len(items) > 6 {
				items, more = items[:6], fmt.Sprintf("; %d more", len(c.Open)-6)
			}
			open = append(open, [2]string{from + " → " + to, strings.Join(items, "; ") + more})
		}
	}
	if len(open) > 0 {
		fmt.Fprintln(w, "\nLet through for part of the traffic:")
		for _, o := range open {
			fmt.Fprintf(w, "  %s\t%s\n", o[0], o[1])
		}
	}
	fmt.Fprintln(w, "\nallow, block: every new connection; partly: some are let through, some not.")
	fmt.Fprintln(w, "+N: rules that let part of the traffic through; -N: rules that block part of it. 'hc matrix FROM TO' lists them.")
	if same {
		fmt.Fprintln(w, "*: a zone to itself: only traffic routed between its interfaces passes the firewall.")
	}
}

func printCell(w io.Writer, c api.MatrixCell) {
	summary := map[string]string{"allowed": "all allowed", "blocked": "all blocked", "partly": "partly allowed"}[c.Summary]
	fmt.Fprintf(w, "%s → %s (the %s chain), new connections: %s\n", c.From, c.To, c.Chain, summary)
	if c.Note != "" {
		fmt.Fprintf(w, "  (%s)\n", c.Note)
	}
	hits := func(h *uint64) string {
		if h == nil {
			return "-"
		}
		return fmt.Sprintf("%d", *h)
	}
	fmt.Fprintln(w, "  VERDICT\tWHAT\tBY\tHITS")
	for _, e := range c.Exceptions {
		fmt.Fprintf(w, "  %s\t%s\t%s\t%s\n", e.Verdict, e.What, e.Label, hits(e.Hits))
	}
	fmt.Fprintf(w, "  %s\tthe rest\t%s\t%s\n", c.Verdict, c.DecidedBy, hits(c.Hits))
	fmt.Fprintln(w, "\nHits: each rule's packets since housecarld started, in all its traffic (a rule can cover several pairs).")
}