Merge pull request #56 from asilvr/structured-stats

Added StructuredStats method to returned structured stat data

iptables/iptables.go

@@ -73,6 +73,20 @@ type IPTables struct {
 	mode           string // the underlying iptables operating mode, e.g. nf_tables
 }
 
+// Stat represents a structured statistic entry.
+type Stat struct {
+	Packets     uint64     `json:"pkts"`
+	Bytes       uint64     `json:"bytes"`
+	Target      string     `json:"target"`
+	Protocol    string     `json:"prot"`
+	Opt         string     `json:"opt"`
+	Input       string     `json:"in"`
+	Output      string     `json:"out"`
+	Source      *net.IPNet `json:"source"`
+	Destination *net.IPNet `json:"destination"`
+	Options     string     `json:"options"`
+}
+
 // New creates a new IPTables.
 // For backwards compatibility, this always uses IPv4, i.e. "iptables".
 func New() (*IPTables, error) {
@@ -264,6 +278,63 @@ func (ipt *IPTables) Stats(table, chain string) ([][]string, error) {
 	return rows, nil
 }
 
+// ParseStat parses a single statistic row into a Stat struct. The input should
+// be a string slice that is returned from calling the Stat method.
+func (ipt *IPTables) ParseStat(stat []string) (parsed Stat, err error) {
+	// For forward-compatibility, expect at least 10 fields in the stat
+	if len(stat) < 10 {
+		return parsed, fmt.Errorf("stat contained fewer fields than expected")
+	}
+
+	// Convert the fields that are not plain strings
+	parsed.Packets, err = strconv.ParseUint(stat[0], 0, 64)
+	if err != nil {
+		return parsed, fmt.Errorf(err.Error(), "could not parse packets")
+	}
+	parsed.Bytes, err = strconv.ParseUint(stat[1], 0, 64)
+	if err != nil {
+		return parsed, fmt.Errorf(err.Error(), "could not parse bytes")
+	}
+	_, parsed.Source, err = net.ParseCIDR(stat[7])
+	if err != nil {
+		return parsed, fmt.Errorf(err.Error(), "could not parse source")
+	}
+	_, parsed.Destination, err = net.ParseCIDR(stat[8])
+	if err != nil {
+		return parsed, fmt.Errorf(err.Error(), "could not parse destination")
+	}
+
+	// Put the fields that are strings
+	parsed.Target = stat[2]
+	parsed.Protocol = stat[3]
+	parsed.Opt = stat[4]
+	parsed.Input = stat[5]
+	parsed.Output = stat[6]
+	parsed.Options = stat[9]
+
+	return parsed, nil
+}
+
+// StructuredStats returns statistics as structured data which may be further
+// parsed and marshaled.
+func (ipt *IPTables) StructuredStats(table, chain string) ([]Stat, error) {
+	rawStats, err := ipt.Stats(table, chain)
+	if err != nil {
+		return nil, err
+	}
+
+	structStats := []Stat{}
+	for _, rawStat := range rawStats {
+		stat, err := ipt.ParseStat(rawStat)
+		if err != nil {
+			return nil, err
+		}
+		structStats = append(structStats, stat)
+	}
+
+	return structStats, nil
+}
+
 func (ipt *IPTables) executeList(args []string) ([]string, error) {
 	var stdout bytes.Buffer
 	if err := ipt.runWithOutput(args, &stdout); err != nil {

iptables/iptables_test.go

@@ -18,6 +18,7 @@ import (
 	"crypto/rand"
 	"fmt"
 	"math/big"
+	"net"
 	"os"
 	"reflect"
 	"testing"
@@ -307,6 +308,41 @@ func runRulesTests(t *testing.T, ipt *IPTables) {
 		t.Fatalf("Stats mismatch: \ngot  %#v \nneed %#v", stats, expectedStats)
 	}
 
+	structStats, err := ipt.StructuredStats("filter", chain)
+	if err != nil {
+		t.Fatalf("StructuredStats failed: %v", err)
+	}
+
+	// It's okay to not check the following errors as they will be evaluated
+	// in the subsequent usage
+	_, address1CIDR, _ := net.ParseCIDR(address1)
+	_, address2CIDR, _ := net.ParseCIDR(address2)
+	_, subnet1CIDR, _ := net.ParseCIDR(subnet1)
+	_, subnet2CIDR, _ := net.ParseCIDR(subnet2)
+
+	expectedStructStats := []Stat{
+		{0, 0, "ACCEPT", "all", opt, "*", "*", subnet1CIDR, address1CIDR, ""},
+		{0, 0, "ACCEPT", "all", opt, "*", "*", subnet2CIDR, address2CIDR, ""},
+		{0, 0, "ACCEPT", "all", opt, "*", "*", subnet2CIDR, address1CIDR, ""},
+		{0, 0, "ACCEPT", "all", opt, "*", "*", address1CIDR, subnet2CIDR, ""},
+	}
+
+	if !reflect.DeepEqual(structStats, expectedStructStats) {
+		t.Fatalf("StructuredStats mismatch: \ngot  %#v \nneed %#v",
+			structStats, expectedStructStats)
+	}
+
+	for i, stat := range expectedStats {
+		stat, err := ipt.ParseStat(stat)
+		if err != nil {
+			t.Fatalf("ParseStat failed: %v", err)
+		}
+		if !reflect.DeepEqual(stat, expectedStructStats[i]) {
+			t.Fatalf("ParseStat mismatch: \ngot  %#v \nneed %#v",
+				stat, expectedStructStats[i])
+		}
+	}
+
 	// Clear the chain that was created.
 	err = ipt.ClearChain("filter", chain)
 	if err != nil {