package firewall

import (
	"bufio"
	"bytes"
	"fmt"
	"net"
	"os/exec"
	"strconv"
	"strings"
)

func ListActiveConnections() ([]ActiveConnection, error) {
	var conns []ActiveConnection
	for _, family := range []string{"ipv4", "ipv6"} {
		out, err := exec.Command("conntrack", "-L", "-f", family, "-o", "extended").CombinedOutput()
		if err != nil {
			return nil, fmt.Errorf("cannot read %s connection tracking: %s: %w", family, out, err)
		}
		scanner := bufio.NewScanner(bytes.NewReader(out))
		scanner.Buffer(make([]byte, 4096), 1024*1024)
		for scanner.Scan() {
			c := parseConntrackLine(scanner.Text())
			if c.Protocol != "" && c.SrcAddress != "" {
				conns = append(conns, c)
			}
		}
		if err := scanner.Err(); err != nil {
			return nil, err
		}
	}
	if conns == nil {
		conns = []ActiveConnection{}
	}
	return conns, nil
}
func parseConntrackLine(line string) ActiveConnection {
	var c ActiveConnection
	fields := strings.Fields(line)
	if len(fields) < 4 {
		return c
	}
	start := 0
	c.IPVersion = 4
	if fields[0] == "ipv4" || fields[0] == "ipv6" {
		if fields[0] == "ipv6" {
			c.IPVersion = 6
		}
		start = 2
	}
	if len(fields) <= start+2 {
		return c
	}
	c.Protocol = strings.ToUpper(fields[start])
	if c.Protocol == "ICMPV6" || c.Protocol == "IPV6-ICMP" {
		c.Protocol = "ICMPV6"
	}
	c.Timeout, _ = strconv.Atoi(fields[start+2])
	var src, dst []string
	var sport, dport []int
	for _, field := range fields[start+3:] {
		kv := strings.SplitN(field, "=", 2)
		if len(kv) == 2 {
			n, _ := strconv.Atoi(kv[1])
			switch kv[0] {
			case "src":
				src = append(src, kv[1])
			case "dst":
				dst = append(dst, kv[1])
			case "sport":
				sport = append(sport, n)
			case "dport":
				dport = append(dport, n)
			case "mark":
				c.ConnectionMark = kv[1]
			case "zone":
				c.Zone = n
			case "bytes":
				v, _ := strconv.ParseUint(kv[1], 10, 64)
				c.Bytes += v
			case "packets":
				v, _ := strconv.ParseUint(kv[1], 10, 64)
				c.Packets += v
			case "id":
				if c.ICMPID == nil {
					value := n
					c.ICMPID = &value
				}
			case "type":
				if c.ICMPType == nil {
					value := n
					c.ICMPType = &value
				}
			case "code":
				if c.ICMPCode == nil {
					value := n
					c.ICMPCode = &value
				}
			}
			continue
		}
		if field == "[ASSURED]" {
			c.Assured = true
		}
		switch field {
		case "ESTABLISHED", "TIME_WAIT", "CLOSE_WAIT", "FIN_WAIT", "SYN_SENT", "SYN_RECV", "CLOSE", "LAST_ACK", "LISTEN":
			c.TCPState = field
		}
	}
	if len(src) > 0 {
		c.SrcAddress = src[0]
	}
	if len(src) > 1 {
		c.ReplySrcAddress = src[1]
	}
	if len(dst) > 0 {
		c.DstAddress = dst[0]
	}
	if len(dst) > 1 {
		c.ReplyDstAddress = dst[1]
	}
	if len(sport) > 0 {
		c.SrcPort = sport[0]
	}
	if len(sport) > 1 {
		c.ReplySrcPort = sport[1]
	}
	if len(dport) > 0 {
		c.DstPort = dport[0]
	}
	if len(dport) > 1 {
		c.ReplyDstPort = dport[1]
	}
	if c.TCPState == "" {
		c.TCPState = c.Protocol
	}
	c.ID = fmt.Sprintf("%d-%d-%s-%s:%d-%s:%d", c.IPVersion, c.Zone, c.Protocol, c.SrcAddress, c.SrcPort, c.DstAddress, c.DstPort)
	if c.ICMPID != nil && c.ICMPType != nil && c.ICMPCode != nil {
		c.ID += fmt.Sprintf("-%d-%d-%d", *c.ICMPID, *c.ICMPType, *c.ICMPCode)
	}
	return c
}

type ConnectionSelector struct {
	Protocol  string `json:"protocol"`
	Src       string `json:"src"`
	Sport     int    `json:"sport"`
	Dst       string `json:"dst"`
	Dport     int    `json:"dport"`
	IPVersion int    `json:"ip_version"`
	Zone      int    `json:"zone"`
	ICMPID    *int   `json:"icmp_id"`
	ICMPType  *int   `json:"icmp_type"`
	ICMPCode  *int   `json:"icmp_code"`
}

func connectionDeleteArgs(s ConnectionSelector) ([]string, error) {
	src, dst := net.ParseIP(s.Src), net.ParseIP(s.Dst)
	if src == nil || dst == nil || (src.To4() == nil) != (dst.To4() == nil) {
		return nil, fmt.Errorf("a complete source/destination IP tuple is required")
	}
	family := "ipv4"
	version := 4
	if src.To4() == nil {
		family = "ipv6"
		version = 6
	}
	if s.IPVersion != 0 && s.IPVersion != version {
		return nil, fmt.Errorf("IP family does not match the addresses")
	}
	if s.Zone < 0 || s.Zone > 65535 {
		return nil, fmt.Errorf("invalid connection zone")
	}
	proto := strings.ToLower(s.Protocol)
	if proto == "ipv6-icmp" {
		proto = "icmpv6"
	}
	args := []string{"-D", "-f", family, "-p", proto, "-s", src.String(), "-d", dst.String(), "--zone", strconv.Itoa(s.Zone)}
	switch proto {
	case "tcp", "udp", "sctp", "dccp":
		if s.Sport < 1 || s.Sport > 65535 || s.Dport < 1 || s.Dport > 65535 {
			return nil, fmt.Errorf("both source and destination ports are required")
		}
		args = append(args, "--sport", strconv.Itoa(s.Sport), "--dport", strconv.Itoa(s.Dport))
	case "icmp", "icmpv6":
		if (proto == "icmpv6") != (version == 6) {
			return nil, fmt.Errorf("ICMP family mismatch")
		}
		if s.ICMPID == nil || s.ICMPType == nil || s.ICMPCode == nil || *s.ICMPID < 0 || *s.ICMPID > 65535 || *s.ICMPType < 0 || *s.ICMPType > 255 || *s.ICMPCode < 0 || *s.ICMPCode > 255 {
			return nil, fmt.Errorf("ICMP ID, type and code are required")
		}
		prefix := "--" + proto
		args = append(args, prefix+"-id", strconv.Itoa(*s.ICMPID), prefix+"-type", strconv.Itoa(*s.ICMPType), prefix+"-code", strconv.Itoa(*s.ICMPCode))
	case "gre", "esp", "ah":
		return nil, fmt.Errorf("deleting %s requires protocol-specific selectors and is not supported", proto)
	default:
		return nil, fmt.Errorf("unsupported connection protocol")
	}
	return args, nil
}
func KillFlow(s ConnectionSelector) error {
	args, err := connectionDeleteArgs(s)
	if err != nil {
		return err
	}
	out, err := exec.Command("conntrack", args...).CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to terminate connection: %s: %w", out, err)
	}
	return nil
}
func KillConnection(protocol, src string, sport int, dst string, dport int) error {
	return KillFlow(ConnectionSelector{Protocol: protocol, Src: src, Sport: sport, Dst: dst, Dport: dport})
}
func ValidateConnectionSelector(s ConnectionSelector) error {
	_, err := connectionDeleteArgs(s)
	return err
}