main.go

package main

import (
	"fmt"
	"log"
	"os/exec"
	"strconv"
	"strings"
)

type Rule struct {
  Name string
  Src string
  Enabled bool
  Src_ip string
  Dest string
  Index int
  Target string
}

type Usage struct {
  Ip string
  Tx_bytes int64
  Rx_bytes int64
}

func parse_usage(usage string) []*Usage {
  var all_usage []*Usage
  for idx, line := range strings.Split(usage, "\n") {
    if idx == 0 {
      // skip header lines
      log.Printf("Header: %s", line)
      continue
    }
    if line == "" {
      continue
    }
    colums := strings.Fields(line)
    if len(colums) < 5 {
      log.Printf("Failed to split line %s got colums %s\n", line, colums)
      continue
    }
    rx_bytes, err := strconv.ParseInt(colums[2], 10, 64)
    if err != nil {
      log.Fatalf("Failed to parse rx_bytes %s", colums[2])
    }
    tx_bytes, err := strconv.ParseInt(colums[4], 10, 64)
    if err != nil {
      log.Fatalf("Failed to parse rx_bytes %s", colums[4])
    }
    usage := &Usage{}
    usage.Tx_bytes = tx_bytes
    usage.Rx_bytes = rx_bytes
    usage.Ip = colums[0]
    all_usage = append(all_usage, usage)
  }
  return all_usage
}

func parse_rules(firewall string) []*Rule {
  var rules []*Rule
  rule := &Rule{}
  for _, line := range(strings.SplitAfter(firewall, "\n")) {
    if line == "" {
      continue
    }
    key_value := strings.SplitN(line, "=", 2)
    key := strings.TrimSpace(key_value[0])
    value := strings.Trim(key_value[1], "' \n")
    if value == "rule" {
      index, err := strconv.ParseInt(key[15:len(key) -1], 10, 31)
      if err != nil {
        log.Panicf("Failed to parse index on line %s", key)
      }
      rule = &Rule{}
      rule.Index = int(index)
      rule.Enabled = true
      rules = append(rules, rule)
      continue
    }
    if strings.HasSuffix(key, "target") {
      rule.Target = value
    } else if strings.HasSuffix(key, "name") {
      rule.Name = value
    } else if strings.HasSuffix(key, "src_ip") {
      rule.Src_ip = value
    } else if strings.HasSuffix(key, "src") {
      rule.Src = value
    } else if strings.HasSuffix(key, "enabled") && value == "0" {
        rule.Enabled = false
    }

  }
  return rules
}

func find_usage_by_ip(all_usage []*Usage, ip string) *Usage {
  for _, usage := range all_usage {
    if usage.Ip == ip {
      return usage
    }
  }
  return nil
}

func toggle_rule(rule *Rule, enable bool) {
  value := 0
  if enable {
    value = 1
  }
  line := fmt.Sprintf("firewall.@rule[%d].enabled=%d", rule.Index, value)
  cmd := exec.Command("uci", "set", line)
  err := cmd.Run()
  if err != nil {
    log.Panicln("Failed to execute uci set %s", line)
  }
  cmd = exec.Command("uci", "commit")
  err = cmd.Run()
  if err != nil {
    log.Panicln("Failed to execute uci commit")
  }
  cmd = exec.Command("fw4", "reload")
  err = cmd.Run()
  if err != nil {
    log.Panicln("Failed to reload fw4")
  }
}

func main() {
  var limit int64 = 512 * 1024 * 1024
  cmd := exec.Command("nlbw", "-c", "csv", "-g", "ip", "-q")
  stdout, err := cmd.Output()
  if err != nil {
    log.Fatalf("Failed to get output for nlbw %s", err)
  }
  cmd = exec.Command("uci", "show", "firewall")
  firewall_out, err := cmd.Output()
  if err != nil {
    log.Fatalf("Failed to get output for uci firewall %s", err)
  }
  output := string(stdout)
  firewall := string(firewall_out)
  rules := parse_rules(firewall)
  all_usage := parse_usage(output)
  for _, rule := range(rules) {
    if rule.Target == "DROP" {
      usage := find_usage_by_ip(all_usage, rule.Src_ip)
      if usage == nil {
        if rule.Enabled {
          log.Printf("Unblocking %s", rule.Src_ip)
          toggle_rule(rule, false)
        }
        continue
      }
      total_bytes := usage.Rx_bytes + usage.Tx_bytes
      log.Printf("Found rule for %s currently using %d", usage.Ip, total_bytes / 1024 / 1024)
      if total_bytes > limit && !rule.Enabled {
        log.Printf("Blocking %s", usage.Ip)
        toggle_rule(rule, true)
      } else if total_bytes < limit && rule.Enabled {
          log.Printf("Unblocking %s", rule.Src_ip)
          toggle_rule(rule, false)
      }
    }
  }
}