/
iptables.go
107 lines (96 loc) · 2.76 KB
/
iptables.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
package utils
import (
"bytes"
"fmt"
"os/exec"
"strings"
"k8s.io/klog/v2"
)
var hasWait bool
//nolint:gochecknoinits // This is actually a good usage of the init() function
func init() {
path, err := exec.LookPath("iptables-restore")
if err != nil {
return
}
args := []string{"iptables-restore", "--help"}
cmd := exec.Cmd{
Path: path,
Args: args,
}
cmdOutput, err := cmd.CombinedOutput()
if err != nil {
return
}
hasWait = strings.Contains(string(cmdOutput), "wait")
}
// SaveInto calls `iptables-save` for given table and stores result in a given buffer.
func SaveInto(table string, buffer *bytes.Buffer) error {
path, err := exec.LookPath("iptables-save")
if err != nil {
return err
}
stderrBuffer := bytes.NewBuffer(nil)
args := []string{"iptables-save", "-t", table}
klog.V(9).Infof("running iptables command: path=`%s` args=%+v", path, args)
cmd := exec.Cmd{
Path: path,
Args: args,
Stdout: buffer,
Stderr: stderrBuffer,
}
if err := cmd.Run(); err != nil {
return fmt.Errorf("%v (%s)", err, stderrBuffer)
}
return nil
}
// Restore runs `iptables-restore` passing data through []byte.
func Restore(table string, data []byte) error {
path, err := exec.LookPath("iptables-restore")
if err != nil {
return err
}
var args []string
if hasWait {
args = []string{"iptables-restore", "--wait", "-T", table}
} else {
args = []string{"iptables-restore", "-T", table}
}
klog.V(9).Infof("running iptables command: path=`%s` args=%+v", path, args)
cmd := exec.Cmd{
Path: path,
Args: args,
Stdin: bytes.NewBuffer(data),
}
b, err := cmd.CombinedOutput()
if err != nil {
return fmt.Errorf("%v (%s)", err, b)
}
return nil
}
// AppendUnique ensures that rule is in chain only once in the buffer and that the occurrence is at the end of the
// buffer
func AppendUnique(buffer bytes.Buffer, chain string, rule []string) bytes.Buffer {
var desiredBuffer bytes.Buffer
// First we need to remove any previous instances of the rule that exist, so that we can be sure that our version
// is unique and appended to the very end of the buffer
rules := strings.Split(buffer.String(), "\n")
if len(rules) > 0 && rules[len(rules)-1] == "" {
rules = rules[:len(rules)-1]
}
for _, foundRule := range rules {
if strings.Contains(foundRule, chain) && strings.Contains(foundRule, strings.Join(rule, " ")) {
continue
}
desiredBuffer.WriteString(foundRule + "\n")
}
// Now append the rule that we wanted to be unique
desiredBuffer = Append(desiredBuffer, chain, rule)
return desiredBuffer
}
// Append appends rule to chain at the end of buffer
func Append(buffer bytes.Buffer, chain string, rule []string) bytes.Buffer {
ruleStr := strings.Join(append([]string{"-A", chain}, rule...), " ")
buffer.WriteString(ruleStr + "\n")
return buffer
}