main.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"strconv"

	"log"
	"os"
	"strings"

	"github.com/fastnetmon/fastnetmon-go"
	"golang.org/x/crypto/ssh"
)

var fast_logger = log.New(os.Stderr, fmt.Sprintf(" %d ", os.Getpid()), log.LstdFlags)

type Configuration struct {
	Log_path    string `json:"log_path" fastnetmon_type:"string"`
	SSHHost     string `json:"ssh_host"`
	SSHPort     uint   `json:"ssh_port"`
	SSHUser     string `json:"ssh_user"`
	SSHPassword string `json:"ssh_password"`
	SSHKeyPath  string `json:"ssh_key_path"`
	ACLIdIPv4   uint   `json:"acl_id_ipv4"`
	ACLIdIPv6   uint   `json:"acl_id_ipv6"`
	Sudo        bool   `json:"sudo"`
}

var conf Configuration

// TODO: it's not very reliable and we need to expose protocol as numberf from FastNetMon itself
// to avoid double conversion from string to number
func protocolNameToNumber(name string) (uint, error) {
	switch strings.ToLower(name) {
	case "tcp":
		return 6, nil
	case "udp":
		return 17, nil
	case "icmp":
		return 1, nil
	case "gre":
		return 47, nil
	case "icmpv6":
		return 58, nil
	default:
		return 0, fmt.Errorf("unknown protocol: %s", name)
	}
}

// flowSpecRuleToACLCommand converts a FastNetMon FlowSpecRule to a Bison
// "vif acl rule" command string
// https://docs.bisonrouter.com/runtime_commands.html#vif-acl-rule-add
func flowSpecRuleToACLCommand(ipv6 bool, rule fastnetmon.FlowSpecRule, aclIDIPv4 uint, aclIDIPv6 uint, prio uint) (string, error) {
	// Determine IP protocol version from prefixes; default to IPv4.
	ipProto := "ipv4"

	if ipv6 {
		ipProto = "ipv6"
	}

	if err := validatePrefixIPVersion(rule.SourcePrefix, ipv6, "source"); err != nil {
		return "", err
	}

	if err := validatePrefixIPVersion(rule.DestinationPrefix, ipv6, "destination"); err != nil {
		return "", err
	}

	aclID := aclIDIPv4
	if ipProto == "ipv6" {
		aclID = aclIDIPv6
	}

	cmd := fmt.Sprintf("vif acl rule %s add aclid %d prio %d", ipProto, aclID, prio)

	// Optional: IP protocol number (first entry wins).
	if len(rule.Protocols) > 1 {
		return "", fmt.Errorf("multiple protocols specified, only one is supported")
	}

	if len(rule.Protocols) == 1 {
		protoNum, err := protocolNameToNumber(rule.Protocols[0])
		if err != nil {
			return "", err
		}
		cmd += fmt.Sprintf(" proto %d", protoNum)
	}

	if rule.SourcePrefix != "" {
		cmd += fmt.Sprintf(" src %s", rule.SourcePrefix)
	}

	if rule.DestinationPrefix != "" {
		cmd += fmt.Sprintf(" dst %s", rule.DestinationPrefix)
	}

	// Source ports: single value or min-max range.
	if len(rule.SourcePorts) > 0 {
		cmd += fmt.Sprintf(" sport %d", rule.SourcePorts[0])
	} else if len(rule.SourcePorts) > 1 {
		return "", fmt.Errorf("multiple source ports specified, only one is supported")
	}

	// Destination ports: single value or min-max range.
	if len(rule.DestinationPorts) > 0 {
		cmd += fmt.Sprintf(" dport %d", rule.DestinationPorts[0])
	} else if len(rule.DestinationPorts) > 1 {
		return "", fmt.Errorf("multiple destination ports specified, only one is supported")
	}

	return cmd, nil
}

func validatePrefixIPVersion(prefix string, ipv6 bool, prefixType string) error {
	if prefix == "" {
		return nil
	}

	ip, _, err := net.ParseCIDR(prefix)
	if err != nil {
		return fmt.Errorf("invalid %s prefix: %s", prefixType, prefix)
	}

	isIPv4Prefix := ip.To4() != nil

	if ipv6 && isIPv4Prefix {
		return fmt.Errorf("%s prefix %s is IPv4 but ipv6 ACL requested", prefixType, prefix)
	}

	if !ipv6 && !isIPv4Prefix {
		return fmt.Errorf("%s prefix %s is IPv6 but ipv4 ACL requested", prefixType, prefix)
	}

	return nil
}

// ACLRule represents a single rule in an ACL list
type ACLRule struct {
	Prio  uint
	Proto string // "any" or numeric protocol number as string
	Src   string // "any" or CIDR/IP
	Dst   string // "any" or CIDR/IP
	Sport string // "any" or port number as string
	Dport string // "any" or port number as string
}

// ACLList represents the parsed output of an ACL list command
type ACLList struct {
	ACLID    uint
	Type     string
	Action   string
	NumRules uint
	Rules    []ACLRule
}

// largestPriority returns the highest priority value among all rules in the ACL list.
// Returns 0 if there are no rules.
func (a *ACLList) largestPriority() uint {
	var max uint
	for _, r := range a.Rules {
		if r.Prio > max {
			max = r.Prio
		}
	}
	return max
}

// flowSpecRuleToACLRule converts a FlowSpecRule to an ACLRule (without priority).
// Fields not set in the FlowSpecRule are represented as "any".
func flowSpecRuleToACLRule(rule fastnetmon.FlowSpecRule) (ACLRule, error) {
	var aclRule ACLRule

	// Protocol
	if len(rule.Protocols) == 0 {
		aclRule.Proto = "any"
	} else if len(rule.Protocols) == 1 {
		protoNum, err := protocolNameToNumber(rule.Protocols[0])
		if err != nil {
			return aclRule, err
		}
		aclRule.Proto = fmt.Sprintf("%d", protoNum)
	} else {
		return aclRule, fmt.Errorf("multiple protocols specified, only one is supported")
	}

	// Source prefix
	if rule.SourcePrefix != "" {
		aclRule.Src = normalizePrefix(rule.SourcePrefix)
	} else {
		aclRule.Src = "any"
	}

	// Destination prefix
	if rule.DestinationPrefix != "" {
		aclRule.Dst = normalizePrefix(rule.DestinationPrefix)
	} else {
		aclRule.Dst = "any"
	}

	// Source port
	if len(rule.SourcePorts) > 0 {
		aclRule.Sport = fmt.Sprintf("%d", rule.SourcePorts[0])
	} else {
		aclRule.Sport = "any"
	}

	// Destination port
	if len(rule.DestinationPorts) > 0 {
		aclRule.Dport = fmt.Sprintf("%d", rule.DestinationPorts[0])
	} else {
		aclRule.Dport = "any"
	}

	return aclRule, nil
}

// normalizePrefix strips /32 (IPv4) or /128 (IPv6) suffixes from a CIDR prefix,
// returning just the IP address, to match the Bison ACL output format.
func normalizePrefix(prefix string) string {
	prefix = strings.TrimSuffix(prefix, "/32")
	prefix = strings.TrimSuffix(prefix, "/128")
	return prefix
}

// findFlowSpecRulePriority searches the ACLList for a rule matching the given
// FlowSpecRule (ignoring priority) and returns its priority.
// Returns 0 and false if no matching rule is found.
func (a *ACLList) findFlowSpecRulePriority(needle ACLRule) (uint, bool, error) {
	for _, r := range a.Rules {
		if r.Proto == needle.Proto && r.Src == needle.Src && r.Dst == needle.Dst &&
			r.Sport == needle.Sport && r.Dport == needle.Dport {
			return r.Prio, true, nil
		}
	}

	return 0, false, nil
}

// parseACLListOutput parses the output of "rcli vif acl list" command.
// Example format:
//
//	acl id 21, type ipv4_tuple, action deny, num rules 20
//	--
//	prio 100, proto 17, src 10.0.0.1/24, dst any, sport 1234, dport 123
func parseACLListOutput(output string) (ACLList, error) {
	var result ACLList

	lines := strings.Split(strings.TrimSpace(output), "\n")
	if len(lines) < 1 {
		return result, fmt.Errorf("empty output")
	}

	// Parse header line: "acl id 21, type ipv4_tuple, action deny, num rules 20"
	header := lines[0]
	headerParts := strings.Split(header, ", ")
	if len(headerParts) != 4 {
		return result, fmt.Errorf("invalid header format: %q", header)
	}

	// Parse "acl id 21"
	aclIDStr := strings.TrimPrefix(headerParts[0], "acl id ")
	if aclIDStr == headerParts[0] {
		return result, fmt.Errorf("invalid acl id field: %q", headerParts[0])
	}
	aclID, err := strconv.ParseUint(aclIDStr, 10, 64)
	if err != nil {
		return result, fmt.Errorf("invalid acl id value: %q", aclIDStr)
	}
	result.ACLID = uint(aclID)

	// Parse "type ipv4_tuple"
	typeStr := strings.TrimPrefix(headerParts[1], "type ")
	if typeStr == headerParts[1] {
		return result, fmt.Errorf("invalid type field: %q", headerParts[1])
	}
	result.Type = typeStr

	// Parse "action deny"
	actionStr := strings.TrimPrefix(headerParts[2], "action ")
	if actionStr == headerParts[2] {
		return result, fmt.Errorf("invalid action field: %q", headerParts[2])
	}
	result.Action = actionStr

	// Parse "num rules 20"
	numRulesStr := strings.TrimPrefix(headerParts[3], "num rules ")
	if numRulesStr == headerParts[3] {
		return result, fmt.Errorf("invalid num rules field: %q", headerParts[3])
	}

	numRules, err := strconv.ParseUint(numRulesStr, 10, 64)
	if err != nil {
		return result, fmt.Errorf("invalid num rules value: %q", numRulesStr)
	}
	result.NumRules = uint(numRules)

	// Parse rules after the "--" separator
	ruleStartIdx := -1
	for i, line := range lines {
		if strings.TrimSpace(line) == "--" {
			ruleStartIdx = i + 1
			break
		}
	}

	if ruleStartIdx == -1 || ruleStartIdx >= len(lines) {
		return result, nil // No rules
	}

	for _, line := range lines[ruleStartIdx:] {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		rule, err := parseACLRuleLine(line)
		if err != nil {
			return result, fmt.Errorf("error parsing rule %q: %w", line, err)
		}
		result.Rules = append(result.Rules, rule)
	}

	return result, nil
}

// parseACLRuleLine parses a single ACL rule line like:
// "prio 100, proto 17, src 10.0.0.1/24, dst any, sport 1234, dport 123"
func parseACLRuleLine(line string) (ACLRule, error) {
	var rule ACLRule

	parts := strings.Split(line, ", ")
	if len(parts) != 6 {
		return rule, fmt.Errorf("expected 6 fields, got %d", len(parts))
	}

	// Parse "prio 100"
	prioStr := strings.TrimPrefix(parts[0], "prio ")
	if prioStr == parts[0] {
		return rule, fmt.Errorf("invalid prio field: %q", parts[0])
	}
	prio, err := strconv.ParseUint(prioStr, 10, 64)
	if err != nil {
		return rule, fmt.Errorf("invalid prio value: %q", prioStr)
	}
	rule.Prio = uint(prio)

	// Parse "proto 17" or "proto any"
	protoStr := strings.TrimPrefix(parts[1], "proto ")
	if protoStr == parts[1] {
		return rule, fmt.Errorf("invalid proto field: %q", parts[1])
	}
	rule.Proto = protoStr

	// Parse "src 10.0.0.1/24" or "src any"
	srcStr := strings.TrimPrefix(parts[2], "src ")
	if srcStr == parts[2] {
		return rule, fmt.Errorf("invalid src field: %q", parts[2])
	}
	rule.Src = srcStr

	// Parse "dst any" or "dst 10.0.0.1/24"
	dstStr := strings.TrimPrefix(parts[3], "dst ")
	if dstStr == parts[3] {
		return rule, fmt.Errorf("invalid dst field: %q", parts[3])
	}
	rule.Dst = dstStr

	// Parse "sport 1234" or "sport any"
	sportStr := strings.TrimPrefix(parts[4], "sport ")
	if sportStr == parts[4] {
		return rule, fmt.Errorf("invalid sport field: %q", parts[4])
	}
	rule.Sport = sportStr

	// Parse "dport 123" or "dport any"
	dportStr := strings.TrimPrefix(parts[5], "dport ")
	if dportStr == parts[5] {
		return rule, fmt.Errorf("invalid dport field: %q", parts[5])
	}
	rule.Dport = dportStr

	return rule, nil
}

func executeSSHCommand(host string, port uint, user string, password string, key string, command string) (string, error) {
	var auth []ssh.AuthMethod

	if key != "" {
		signer, err := ssh.ParsePrivateKey([]byte(key))
		if err != nil {
			return "", err
		}
		auth = append(auth, ssh.PublicKeys(signer))
	} else if password != "" {
		auth = append(auth, ssh.Password(password))
	} else {
		return "", fmt.Errorf("no authentication method provided")
	}

	config := &ssh.ClientConfig{
		User:            user,
		Auth:            auth,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	addr := fmt.Sprintf("%s:%d", host, port)

	client, err := ssh.Dial("tcp", addr, config)

	if err != nil {
		return "", err
	}
	defer client.Close()

	session, err := client.NewSession()

	if err != nil {
		return "", err
	}

	defer session.Close()

	output, err := session.CombinedOutput(command)

	if err != nil {
		return string(output), fmt.Errorf("SSH command failed: %v, output: %s", err, output)
	}

	// fast_logger.Printf("SSH command output: %s", output)
	return string(output), nil
}

func main() {
	conf.SSHUser = "fastnetmon"
	conf.SSHPort = 22
	conf.Log_path = "/var/log/fastnetmon/fastnetmon_bison_flow_spec_acl.log"

	file_as_array, err := os.ReadFile("/etc/fastnetmon_bison_flow_spec_acl.json")

	if err != nil {
		log.Fatalf("Could not read configuration file with error: %v", err)
	}

	// This command will override our default configuration
	err = json.Unmarshal(file_as_array, &conf)

	if err != nil {
		log.Fatalf("Could not read json configuration: %v", err)
	}

	log_file, err := os.OpenFile(conf.Log_path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)

	if err != nil {
		fast_logger.Fatalf("Cannot open log file: %v", err)
	}

	defer log_file.Close()

	multi_writer := io.MultiWriter(os.Stdout, log_file)

	fast_logger.SetOutput(multi_writer)

	fast_logger.Printf("Prepared to read data from stdin")
	stdin_data, err := io.ReadAll(os.Stdin)

	if err != nil {
		fast_logger.Fatal("Cannot read data from stdin")
	}

	callback_data := fastnetmon.CallbackDetails{}

	fast_logger.Printf("Callback raw data: %s", stdin_data)

	err = json.Unmarshal([]byte(stdin_data), &callback_data)

	if err != nil {
		fast_logger.Printf("Raw data: %s", stdin_data)
		fast_logger.Fatalf("Cannot unmarshal data: %v", err)
	}

	// Scope can be per host or total hostgroups
	alert_scope := callback_data.AlertScope

	if alert_scope != "host" {
		fast_logger.Fatalf("Unknown scope: %s Only host scope is supported", alert_scope)
	}

	var ssh_key_data []byte

	if conf.SSHKeyPath != "" {
		ssh_key_data, err = os.ReadFile(conf.SSHKeyPath)

		if err != nil {
			fast_logger.Fatalf("Cannot read SSH key file: %v", err)
		}
	}

	if callback_data.Action == "connection_check" {
		fast_logger.Printf("Connection check for host: %s", conf.SSHHost)

		output, err := executeSSHCommand(conf.SSHHost, conf.SSHPort, conf.SSHUser, conf.SSHPassword, string(ssh_key_data), "id")

		if err != nil {
			fast_logger.Fatalf("Error executing SSH command: %v", err)
		}

		fast_logger.Printf("SSH command output: %s", output)
		os.Exit(0)
	}

	fast_logger.Printf("Flow Spec rules: %+v", callback_data.FlowSpecRules)

	if len(callback_data.FlowSpecRules) == 0 {
		fast_logger.Fatalf("Empty list of Flow Spec rules")
	}

	// We have at least one rule. Our core does not return more then one
	flow_spec_rule := callback_data.FlowSpecRules[0]

	rcli_command := "/usr/sbin/rcli"

	if conf.Sudo {
		rcli_command = "sudo " + rcli_command
	}

	// Before we can add or remove ACL rule we need to load cirrent list of ACL rules
	outputAclList, err := executeSSHCommand(conf.SSHHost, conf.SSHPort, conf.SSHUser, conf.SSHPassword, string(ssh_key_data), fmt.Sprintf("%s sh vif acl rules aclid %d", rcli_command, conf.ACLIdIPv4))

	if err != nil {
		fast_logger.Fatalf("Error executing SSH command to get ACL list: %v", err)
	}

	aclList, err := parseACLListOutput(outputAclList)

	if err != nil {
		fast_logger.Fatalf("Error parsing ACL list output: %v", err)
	}

	fast_logger.Printf("Current ACL list: %+v", aclList)

	// We need largest priority to add new rule with unique priority and avoid conflicts with existing rules
	largestPriority := aclList.largestPriority()

	fast_logger.Printf("Largest existing ACL rule priority: %d", largestPriority)

	ipv6 := false

	if callback_data.Action == "partial_block" {

		// Use largestPriority+1 as priority to uniquely identify our rule and avoid conflicts with existing rules
		aclString, err := flowSpecRuleToACLCommand(ipv6, flow_spec_rule, conf.ACLIdIPv4, conf.ACLIdIPv6, largestPriority+1)

		if err != nil {
			fast_logger.Fatalf("Error generating ACL command: %v", err)
		}

		sshCommand := fmt.Sprintf("%s %s", rcli_command, aclString)

		fast_logger.Printf("Executing SSH command: %s", sshCommand)

		output, err := executeSSHCommand(conf.SSHHost, conf.SSHPort, conf.SSHUser, conf.SSHPassword, string(ssh_key_data), sshCommand)

		if err != nil {
			fast_logger.Fatalf("Error executing SSH command: %v", err)
		}

		fast_logger.Printf("SSH command output: %s", output)
		fast_logger.Printf("Successfully added ACL rule with priority %d", largestPriority+1)

	} else if callback_data.Action == "partial_unblock" {
		needle, err := flowSpecRuleToACLRule(flow_spec_rule)
		if err != nil {
			fast_logger.Fatalf("Error converting FlowSpecRule to ACLRule: %v", err)
		}

		fast_logger.Printf("Looking for ACL rule matching: %+v", needle)

		// To remove ACL rule we need to find its priority first
		prio, found, err := aclList.findFlowSpecRulePriority(needle)

		if err != nil {
			fast_logger.Fatalf("Error finding ACL rule priority: %v", err)
		}

		if !found {
			fast_logger.Printf("ACL rule not found, nothing to remove")
			os.Exit(0)
		}

		fast_logger.Printf("Found ACL rule with priority %d, removing it", prio)

		aclRemoveCommand := fmt.Sprintf("%s vif acl rule ipv4 del aclid %d prio %d", rcli_command, conf.ACLIdIPv4, prio)

		fast_logger.Printf("Running command: %s", aclRemoveCommand)

		output, err := executeSSHCommand(conf.SSHHost, conf.SSHPort, conf.SSHUser, conf.SSHPassword, string(ssh_key_data), aclRemoveCommand)

		if err != nil {
			fast_logger.Fatalf("Error executing SSH command: %v", err)
		}

		fast_logger.Printf("SSH command output: %s", output)
		fast_logger.Printf("Successfully removed ACL rule with priority %d", prio)
	} else {
		fast_logger.Fatalf("Unknown action: %s", callback_data.Action)
	}
}