From 093425d4d1326941f987920cf2adaaddc5ecf64a Mon Sep 17 00:00:00 2001
From: Guillermo Gaston <gaslor@amazon.com>
Date: Thu, 23 Feb 2023 23:04:37 +0000
Subject: [PATCH 1/2] Add tests for distroless

---
 Makefile                                  |  16 +-
 internal/iptables/detect.go               |  14 +-
 test/Dockerfile.test-alpine               |   2 +-
 test/Dockerfile.test-debian               |   2 +-
 test/Dockerfile.test-distroless           |  60 ++++++
 test/Dockerfile.test-fedora               |   2 +-
 test/build/clean-distroless.sh            |  43 ++++
 test/build/package-utils.sh               |  48 +++++
 test/build/stage-binaries-from-package.sh |  65 ++++++
 test/build/stage-binary-and-deps.sh       |  72 +++++++
 test/run-test.sh                          |  26 ++-
 test/test.sh                              |  85 --------
 test/wrapper_test.go                      | 229 ++++++++++++++++++++++
 13 files changed, 558 insertions(+), 106 deletions(-)
 create mode 100644 test/Dockerfile.test-distroless
 create mode 100755 test/build/clean-distroless.sh
 create mode 100644 test/build/package-utils.sh
 create mode 100755 test/build/stage-binaries-from-package.sh
 create mode 100755 test/build/stage-binary-and-deps.sh
 delete mode 100755 test/test.sh
 create mode 100644 test/wrapper_test.go

diff --git a/Makefile b/Makefile
index 807441e..95a11da 100644
--- a/Makefile
+++ b/Makefile
@@ -18,19 +18,25 @@ fmt: ## Check formatting
     	exit 1; \
 	fi
 
+build-tests: $(BIN_DIR)
+	$(GO) test ./test -c -o $(BIN_DIR)/tests
+
 check: check-debian check-debian-nosanity check-debian-backports check-fedora check-alpine
 
-check-debian: build
+check-debian: build build-tests
 	./test/run-test.sh --build-fail debian
 
-check-debian-nosanity: build
+check-debian-nosanity: build build-tests
 	./test/run-test.sh --build-arg="INSTALL_ARGS=--no-sanity-check" --nft-fail debian-nosanity
 
-check-debian-backports: build
+check-debian-backports: build build-tests
 	./test/run-test.sh --build-arg="REPO=buster-backports" debian-backports
 
-check-fedora: build
+check-fedora: build build-tests
 	./test/run-test.sh fedora
 
-check-alpine: build
+check-alpine: build build-tests
 	./test/run-test.sh alpine
+
+check-distroless: build build-tests
+	./test/run-test.sh distroless
diff --git a/internal/iptables/detect.go b/internal/iptables/detect.go
index a87641b..67e2d60 100644
--- a/internal/iptables/detect.go
+++ b/internal/iptables/detect.go
@@ -39,8 +39,8 @@ func DetectBinaryDir() (string, error) {
 type Mode string
 
 const (
-	legacy Mode = "legacy"
-	nft    Mode = "nft"
+	Legacy Mode = "legacy"
+	NFT    Mode = "nft"
 )
 
 // DetectMode inspects the current iptables entries and tries to
@@ -59,12 +59,12 @@ func DetectMode(ctx context.Context, iptables Installation) Mode {
 	rulesOutput := &bytes.Buffer{}
 	_ = iptables.NFTSave(ctx, rulesOutput, "-t", "mangle")
 	if hasKubeletChains(rulesOutput.Bytes()) {
-		return nft
+		return NFT
 	}
 	rulesOutput.Reset()
 	_ = iptables.NFTSaveIP6(ctx, rulesOutput, "-t", "mangle")
 	if hasKubeletChains(rulesOutput.Bytes()) {
-		return nft
+		return NFT
 	}
 	rulesOutput.Reset()
 
@@ -74,14 +74,14 @@ func DetectMode(ctx context.Context, iptables Installation) Mode {
 	// exist, which we don't want. So we have to grab all the rules.
 	_ = iptables.LegacySave(ctx, rulesOutput)
 	if hasKubeletChains(rulesOutput.Bytes()) {
-		return legacy
+		return Legacy
 	}
 	rulesOutput.Reset()
 	_ = iptables.LegacySaveIP6(ctx, rulesOutput)
 	if hasKubeletChains(rulesOutput.Bytes()) {
-		return legacy
+		return Legacy
 	}
 
 	// If we can't detect any of the 2 patterns, default to nft.
-	return nft
+	return NFT
 }
diff --git a/test/Dockerfile.test-alpine b/test/Dockerfile.test-alpine
index 792190a..b7cc734 100644
--- a/test/Dockerfile.test-alpine
+++ b/test/Dockerfile.test-alpine
@@ -20,4 +20,4 @@ RUN apk add --no-cache iptables
 COPY iptables-wrapper-installer.sh /
 COPY bin/iptables-wrapper /
 RUN /iptables-wrapper-installer.sh
-COPY test/test.sh /
+COPY bin/tests /
diff --git a/test/Dockerfile.test-debian b/test/Dockerfile.test-debian
index de22e6b..ad86f8e 100644
--- a/test/Dockerfile.test-debian
+++ b/test/Dockerfile.test-debian
@@ -26,4 +26,4 @@ RUN echo deb http://deb.debian.org/debian buster-backports main >> /etc/apt/sour
 COPY iptables-wrapper-installer.sh /
 COPY bin/iptables-wrapper /
 RUN /iptables-wrapper-installer.sh ${INSTALL_ARGS}
-COPY test/test.sh /
+COPY bin/tests /
diff --git a/test/Dockerfile.test-distroless b/test/Dockerfile.test-distroless
new file mode 100644
index 0000000..bd3ffec
--- /dev/null
+++ b/test/Dockerfile.test-distroless
@@ -0,0 +1,60 @@
+# Copyright 2023 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+### Dockerfile for building a distroless image for testing
+# Modified version of https://github.com/kubernetes/release/blob/master/images/build/distroless-iptables/distroless/Dockerfile
+
+ARG STAGE_DIR="/opt/stage"
+
+FROM debian:bullseye-slim as build
+ARG STAGE_DIR
+
+COPY test/build/stage-binaries-from-package.sh /
+COPY test/build/package-utils.sh /
+COPY test/build/stage-binary-and-deps.sh /
+
+RUN mkdir -p "${STAGE_DIR}" && \
+    /stage-binaries-from-package.sh "${STAGE_DIR}" conntrack \
+    ebtables    \
+    ipset       \
+    iptables    \
+    kmod		&& \
+	`# below binaries and dash are used by iptables-wrapper-installer.sh` \
+    /stage-binary-and-deps.sh "${STAGE_DIR}" /bin/dash \
+    /bin/mv \
+    /bin/chmod \
+    /bin/grep \
+    /bin/ln  \
+    /bin/rm \
+    /bin/sleep \
+    /usr/bin/wc
+
+RUN ln -sf /bin/dash "${STAGE_DIR}"/bin/sh
+
+FROM gcr.io/distroless/static as intermediate
+ARG STAGE_DIR
+
+COPY test/build/clean-distroless.sh /clean-distroless.sh
+COPY --from=build "${STAGE_DIR}" /
+COPY iptables-wrapper-installer.sh /
+COPY bin/iptables-wrapper /
+# iptables-wrapper-installer needs to know that iptables exists before doing all its magic
+RUN echo "" >  /usr/sbin/iptables && \
+	/iptables-wrapper-installer.sh && \
+    /clean-distroless.sh
+
+FROM scratch
+
+COPY --from=intermediate / /
+COPY bin/tests /
diff --git a/test/Dockerfile.test-fedora b/test/Dockerfile.test-fedora
index dace63e..e00b44d 100644
--- a/test/Dockerfile.test-fedora
+++ b/test/Dockerfile.test-fedora
@@ -20,4 +20,4 @@ RUN dnf install -y iptables iptables-legacy iptables-nft
 COPY iptables-wrapper-installer.sh /
 COPY bin/iptables-wrapper /
 RUN /iptables-wrapper-installer.sh
-COPY test/test.sh /
+COPY bin/tests /
diff --git a/test/build/clean-distroless.sh b/test/build/clean-distroless.sh
new file mode 100755
index 0000000..b46844c
--- /dev/null
+++ b/test/build/clean-distroless.sh
@@ -0,0 +1,43 @@
+#!/bin/sh
+
+# Copyright 2022 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# USAGE: clean-distroless.sh
+
+# Modified version of https://github.com/kubernetes/release/blob/master/images/build/distroless-iptables/distroless/clean-distroless.sh
+
+REMOVE="/usr/share/base-files
+/usr/share/man
+/usr/lib/*-linux-gnu/gconv/
+/usr/bin/c_rehash
+/usr/bin/openssl
+/bin/mv
+/bin/chmod
+/bin/grep
+/bin/ln
+/bin/sleep
+/usr/bin/wc
+/iptables-wrapper-installer.sh
+/bin/sh
+/bin/dash
+/clean-distroless.sh
+/bin/rm"
+
+IFS="
+"
+
+for item in ${REMOVE}; do
+    rm -rf "${item}"
+done
diff --git a/test/build/package-utils.sh b/test/build/package-utils.sh
new file mode 100644
index 0000000..e87cfd9
--- /dev/null
+++ b/test/build/package-utils.sh
@@ -0,0 +1,48 @@
+#!/bin/bash
+
+# Copied from https://github.com/kubernetes/release/blob/master/images/build/distroless-iptables/distroless/package-utils.sh
+
+# Copyright 2022 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# file_to_package identifies the debian package that provided the file $1
+file_to_package() {
+    # `dpkg-query --search $file-pattern` outputs lines with the format: "$package: $file-path"
+    # where $file-path belongs to $package
+    # https://manpages.debian.org/jessie/dpkg/dpkg-query.1.en.html
+    dpkg-query --search "$(realpath "${1}")" | cut -d':' -f1
+}
+
+# package_to_copyright gives the path to the copyright file for the package $1
+package_to_copyright() {
+    echo "/usr/share/doc/${1}/copyright"
+}
+
+# stage_file stages the filepath $1 to $2, following symlinks
+# and staging copyrights
+stage_file() {
+    cp -a --parents "${1}" "${2}"
+    # recursively follow symlinks
+    if [[ -L "${1}" ]]; then
+        stage_file "$(cd "$(dirname "${1}")" || exit; realpath -s "$(readlink "${1}")")" "${2}"
+    fi
+    # get the package so we can stage package metadata as well
+    package="$(file_to_package "${1}")"
+    # stage the copyright for the file
+    cp -a --parents "$(package_to_copyright "${package}")" "${2}"
+    # stage the package status mimicking bazel
+    # https://github.com/bazelbuild/rules_docker/commit/f5432b813e0a11491cf2bf83ff1a923706b36420
+    # instead of parsing the control file, we can just get the actual package status with dpkg
+    dpkg -s "${package}" > "${2}/var/lib/dpkg/status.d/${package}"
+}
diff --git a/test/build/stage-binaries-from-package.sh b/test/build/stage-binaries-from-package.sh
new file mode 100755
index 0000000..b08bbaa
--- /dev/null
+++ b/test/build/stage-binaries-from-package.sh
@@ -0,0 +1,65 @@
+#!/bin/bash
+
+# Copied from https://github.com/kubernetes/release/blob/master/images/build/distroless-iptables/distroless/stage-binaries-from-package.sh
+
+# Copyright 2022 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# USAGE: stage-binaries-from-package.sh /opt/stage package1 package2
+#
+# Stages all the packages and its dependencies (+ libraries and copyrights) to $1
+#
+# This is intended to be used in a multi-stage docker build with a distroless/base
+# or distroless/cc image.
+set -e
+
+. package-utils.sh
+
+stage_file_list() {
+    IFS="
+    "
+    REQUIRED_FILES="$(dpkg -L "${1}" | grep -vE '(/\.|/s?bin/|/usr/share/(man|doc|.*-completion))' | sed 's/\n/ /g')"
+    for file in $REQUIRED_FILES; do 
+        if [ -f "$file" ]; then
+            stage_file "${file}" "${STAGE_DIR}"
+        fi
+    done
+
+    BIN_LIST="$(dpkg -L "${1}" | grep -E '/s?bin/' |sed 's/\n/ /g')"
+    for binary in $BIN_LIST; do
+        /stage-binary-and-deps.sh "${2}" "${binary}"
+    done
+}
+
+get_dependent_packages() {
+    apt-cache depends "${1}" |grep Depends|awk -F '.*Depends:[[:space:]]?' '{print $2}'
+}
+
+main() {
+    STAGE_DIR="${1}/"
+    mkdir -p "${STAGE_DIR}"/var/lib/dpkg/status.d/
+    apt -y update
+    shift
+    while (( "$#" )); do        # While there are arguments still to be shifted
+        PACKAGE="${1}"
+        apt -y install "${PACKAGE}"
+        stage_file_list "${PACKAGE}" "$STAGE_DIR"
+        while IFS= read -r c_dep; do
+            stage_file_list "${c_dep}" "${STAGE_DIR}"
+        done < <(get_dependent_packages "${PACKAGE}")
+        shift
+    done
+}
+
+main "$@"
diff --git a/test/build/stage-binary-and-deps.sh b/test/build/stage-binary-and-deps.sh
new file mode 100755
index 0000000..c88ecc7
--- /dev/null
+++ b/test/build/stage-binary-and-deps.sh
@@ -0,0 +1,72 @@
+#!/bin/bash
+
+# Copied from https://github.com/kubernetes/release/blob/master/images/build/distroless-iptables/distroless/stage-binary-and-deps.sh
+
+# Copyright 2021 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# USAGE: stage-binary-and-deps.sh haproxy /opt/stage
+#
+# Stages $1 and its dependencies + their copyright files to $2
+#
+# This is intended to be used in a multi-stage docker build with a distroless/base
+# or distroless/cc image.
+# This script was originally created by KinD maintainers and can be found at:
+#   https://github.com/kubernetes-sigs/kind/blob/v0.14.0/images/haproxy/stage-binary-and-deps.sh
+
+set -o errexit
+set -o nounset
+set -o pipefail
+
+
+. package-utils.sh
+
+# binary_to_libraries identifies the library files needed by the binary $1 with ldd
+binary_to_libraries() {
+    # see: https://man7.org/linux/man-pages/man1/ldd.1.html
+    ldd "${1}" \
+    `# strip the leading '${name} => ' if any so only '/lib-foo.so (0xf00)' remains` \
+    | sed -E 's#.* => /#/#' \
+    `# we want only the path remaining, not the (0x${LOCATION})` \
+    | awk '{print $1}' \
+    `# linux-vdso.so.1 is a special virtual shared object from the kernel` \
+    `# see: http://man7.org/linux/man-pages/man7/vdso.7.html` \
+    | grep -v 'linux-vdso.so.1'
+}
+
+# main script logic
+main(){
+    local STAGE_DIR="${1}/"
+    shift
+    while (( "$#" )); do
+        BINARY="${1}"
+        # locate the path to the binary
+        local binary_path
+        binary_path="$(which "${BINARY}")"
+    
+        # ensure package metadata dir
+        mkdir -p "${STAGE_DIR}"/var/lib/dpkg/status.d/
+    
+        # stage the binary itself
+        stage_file "${binary_path}" "${STAGE_DIR}"
+    
+        # stage the dependencies of the binary
+        while IFS= read -r c_dep; do
+            stage_file "${c_dep}" "${STAGE_DIR}"
+        done < <(binary_to_libraries "${binary_path}")
+        shift
+    done
+}
+
+main "$@"
diff --git a/test/run-test.sh b/test/run-test.sh
index a43a7b1..b08a6df 100755
--- a/test/run-test.sh
+++ b/test/run-test.sh
@@ -20,7 +20,6 @@ set -o pipefail
 
 if [[ -n "${DEBUG:-}" ]]; then
     set -x
-    dash_x="-x"
 fi
 
 build_arg=""
@@ -81,11 +80,19 @@ function docker() {
 }
 
 function build() {
+	if [[ -z "${DEBUG:-}" ]]; then
+		quiet="-q"
+	fi
+
+	if [[ -z "${CACHE_BUILDS:-}" ]]; then
+		no_cache="--no-cache"
+	fi
+
     build_tag=iptables-wrapper-test-$1
     dockerfile=Dockerfile.test-${1%%-*}
     shift
 
-    docker build --no-cache -q -t ${build_tag} -f test/${dockerfile} "$@" .
+    docker build ${no_cache:-} ${quiet:-} -t ${build_tag} -f test/${dockerfile} "$@" .
 }
 
 function PASS() {
@@ -105,11 +112,18 @@ if ! build "${tag}" ${build_arg}; then
     FAIL "build failed unexpectedly"
 fi
 
-if ! docker run --privileged "iptables-wrapper-test-${tag}" /bin/sh ${dash_x:-} /test.sh legacy; then
-    FAIL "failed legacy iptables / new rules test"
+if ! docker run --privileged "iptables-wrapper-test-${tag}" /tests -test.v -test.run "^TestIPTablesWrapperLegacy$" ; then
+    FAIL "failed legacy iptables"
+fi
+if ! docker run --privileged "iptables-wrapper-test-${tag}" /tests -test.v -test.run "^TestIPTablesWrapperNFT$" ; then
+    FAIL "failed nft iptables"
+fi
+
+if ! docker run --privileged "iptables-wrapper-test-${tag}" /tests -test.v -test.run "^TestIP6TablesWrapperLegacy$" ; then
+    FAIL "failed legacy ip6tables"
 fi
-if ! docker run --privileged "iptables-wrapper-test-${tag}" /bin/sh ${dash_x:-} /test.sh nft; then
-    FAIL "failed nft iptables / new rules test"
+if ! docker run --privileged "iptables-wrapper-test-${tag}" /tests -test.v -test.run "^TestIP6TablesWrapperNFT$" ; then
+    FAIL "failed nft ip6tables"
 fi
 
 PASS "success"
diff --git a/test/test.sh b/test/test.sh
deleted file mode 100755
index 0fe4ed5..0000000
--- a/test/test.sh
+++ /dev/null
@@ -1,85 +0,0 @@
-#!/bin/sh
-#
-# Copyright 2020 The Kubernetes Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-set -eu
-
-mode=$1
-
-case "${mode}" in
-    legacy)
-        wrongmode=nft
-        ;;
-    nft)
-        wrongmode=legacy
-        ;;
-    *)
-        echo "ERROR: bad mode '${mode}'" 1>&2
-        exit 1
-        ;;
-esac
-
-if [ -d /usr/sbin -a -e /usr/sbin/iptables ]; then
-    sbin="/usr/sbin"
-elif [ -d /sbin -a -e /sbin/iptables ]; then
-    sbin="/sbin"
-else
-    echo "ERROR: iptables is not present in either /usr/sbin or /sbin" 1>&2
-    exit 1
-fi
-
-ensure_iptables_undecided() {
-    iptables=$(realpath "${sbin}/iptables")
-    if [ "${iptables}" != "${sbin}/iptables-wrapper" ]; then
-	echo "iptables link was resolved prematurely! (${iptables})" 1>&2
-	exit 1
-    fi
-}
-
-ensure_iptables_resolved() {
-    expected=$1
-    iptables=$(realpath "${sbin}/iptables")
-    if [ "${iptables}" = "${sbin}/iptables-wrapper" ]; then
-	echo "iptables link is not yet resolved!" 1>&2
-	exit 1
-    fi
-    version=$(iptables -V | sed -e 's/.*(\(.*\)).*/\1/')
-    case "${version}/${expected}" in
-	legacy/legacy|nf_tables/nft)
-	    return
-	    ;;
-	*)
-	    echo "iptables link resolved incorrectly (expected ${expected}, got ${version})" 1>&2
-	    exit 1
-	    ;;
-    esac
-}
-
-ensure_iptables_undecided
-
-# Initialize the chosen iptables mode with just a hint chain
-iptables-${mode} -t mangle -N KUBE-IPTABLES-HINT
-
-# Put some junk in the other iptables system
-iptables-${wrongmode} -t filter -N BAD-1
-iptables-${wrongmode} -t filter -A BAD-1 -j ACCEPT
-iptables-${wrongmode} -t filter -N BAD-2
-iptables-${wrongmode} -t filter -A BAD-2 -j DROP
-
-ensure_iptables_undecided
-
-iptables -L > /dev/null
-
-ensure_iptables_resolved ${mode}
diff --git a/test/wrapper_test.go b/test/wrapper_test.go
new file mode 100644
index 0000000..2af20f1
--- /dev/null
+++ b/test/wrapper_test.go
@@ -0,0 +1,229 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package test
+
+import (
+	"bytes"
+	"context"
+	"os/exec"
+	"path/filepath"
+	"regexp"
+	"testing"
+
+	"github.com/kubernetes-sigs/iptables-wrappers/internal/commands"
+	"github.com/kubernetes-sigs/iptables-wrappers/internal/iptables"
+)
+
+// iptablesVersion denotes the IP version for a iptables command, V4 or V6
+type iptablesIPVersion string
+
+const (
+	v4 iptablesIPVersion = "iptables"
+	v6 iptablesIPVersion = "ip6tables"
+)
+
+// iptablesMode represents a iptables mode.
+type iptablesMode struct {
+	original, wrongMode iptables.Mode
+	// expectedIPTablesVStr is the subtring expected in betwen brakets when
+	// running `iptables -V` for this particular mode
+	// ex. for nft -> `iptables v1.8.7 (nf_tables)`
+	expectedIPTablesVStr string
+}
+
+var legacy = iptablesMode{
+	original:             iptables.Legacy,
+	wrongMode:            iptables.NFT,
+	expectedIPTablesVStr: "legacy",
+}
+
+var nft = iptablesMode{
+	original:             iptables.NFT,
+	wrongMode:            iptables.Legacy,
+	expectedIPTablesVStr: "nf_tables",
+}
+
+func TestIPTablesWrapperLegacy(t *testing.T) {
+	tt := newIPTablesWrapperTest(t, v4, legacy)
+	runTest(t, tt)
+}
+
+func TestIPTablesWrapperNFT(t *testing.T) {
+	tt := newIPTablesWrapperTest(t, v4, nft)
+	runTest(t, tt)
+}
+
+func TestIP6TablesWrapperLegacy(t *testing.T) {
+	tt := newIPTablesWrapperTest(t, v6, legacy)
+	runTest(t, tt)
+}
+
+func TestIP6TablesWrapperNFT(t *testing.T) {
+	tt := newIPTablesWrapperTest(t, v6, nft)
+	runTest(t, tt)
+}
+
+func runTest(tb testing.TB, test iptablesWrapperTest) {
+	ctx := context.Background()
+	test.assertIPTablesUndecided(tb)
+
+	tb.Log("Inserting chains")
+	// Initialize the chosen iptables mode with just a hint chain
+	test.iptables.runAndAssertSuccess(ctx, tb, "-t", "mangle", "-N", "KUBE-IPTABLES-HINT")
+
+	// Put some junk in the other iptables system
+	test.wrongModeIPTables.runAndAssertSuccess(ctx, tb, "-t", "filter", "-N", "BAD-1")
+	test.wrongModeIPTables.runAndAssertSuccess(ctx, tb, "-t", "filter", "-A", "BAD-1", "-j", "ACCEPT")
+	test.wrongModeIPTables.runAndAssertSuccess(ctx, tb, "-t", "filter", "-N", "BAD-2")
+	test.wrongModeIPTables.runAndAssertSuccess(ctx, tb, "-t", "filter", "-A", "BAD-2", "-j", "DROP")
+
+	test.assertIPTablesUndecided(tb)
+
+	// This should run the iptables-wrapper
+	tb.Log("Running `iptables -L` command")
+	c := exec.CommandContext(ctx, "iptables", "-L")
+	assertSuccess(tb, commands.RunAndReadError(c))
+
+	test.assertIPTablesResolved(ctx, tb)
+}
+
+type iptablesWrapperTest struct {
+	mode                        iptablesMode
+	iptables, wrongModeIPTables ipTablesRunner
+	sbinPath                    string
+	wrapperPath                 string
+	iptablesPath, ip6tablesPath string
+}
+
+// newIPTablesWrapperTest creates a new test setup for a particular IP version of iptables (iptables or ip6tables)
+// and a particular mode (legacy or nft)
+func newIPTablesWrapperTest(tb testing.TB, ipV iptablesIPVersion, mode iptablesMode) iptablesWrapperTest {
+	sbinPath, err := iptables.DetectBinaryDir()
+	assertSuccess(tb, err)
+
+	return iptablesWrapperTest{
+		mode:              mode,
+		iptables:          newIPTablesRunner(ipV, mode.original),
+		wrongModeIPTables: newIPTablesRunner(ipV, mode.wrongMode),
+		sbinPath:          sbinPath,
+		wrapperPath:       filepath.Join(sbinPath, "iptables-wrapper"),
+		iptablesPath:      filepath.Join(sbinPath, "iptables"),
+		ip6tablesPath:     filepath.Join(sbinPath, "ip6tables"),
+	}
+}
+
+func (tt iptablesWrapperTest) assertIPTablesUndecided(tb testing.TB) {
+	tb.Log("Checking the iptables mode hasn't been decided yet")
+	iptablesRealPath := tt.iptablesRealPath(tb)
+	if !tt.isIPTablesWrapper(iptablesRealPath) {
+		tb.Fatalf("iptables link was resolved prematurely, got [%s]", iptablesRealPath)
+	}
+	tb.Logf("iptables points to %s", iptablesRealPath)
+
+	ip6tablesRealPath := tt.ip6tablesRealPath(tb)
+	if !tt.isIPTablesWrapper(ip6tablesRealPath) {
+		tb.Fatalf("ip6tables link was resolved prematurely, got [%s]", ip6tablesRealPath)
+	}
+	tb.Logf("ip6tables points to %s", ip6tablesRealPath)
+}
+
+func (tt iptablesWrapperTest) assertIPTablesResolved(ctx context.Context, tb testing.TB) {
+	tb.Logf("Checking the iptables mode has been resolved to %s", tt.mode.original)
+	iptablesRealPath := tt.iptablesRealPath(tb)
+	if tt.isIPTablesWrapper(iptablesRealPath) {
+		tb.Fatal("iptables link is not yet resolved")
+	}
+
+	ip6tablesRealPath := tt.iptablesRealPath(tb)
+	if tt.isIPTablesWrapper(ip6tablesRealPath) {
+		tb.Fatal("ip6tables link is not yet resolved")
+	}
+
+	mode := readIPTablesMode(ctx, tb, "iptables")
+	if mode != tt.mode.expectedIPTablesVStr {
+		tb.Fatalf("iptables link resolved incorrectly: expected %s, got %s", tt.mode.expectedIPTablesVStr, mode)
+	}
+
+	mode = readIPTablesMode(ctx, tb, "ip6tables")
+	if mode != tt.mode.expectedIPTablesVStr {
+		tb.Fatalf("ip6tables link resolved incorrectly: expected %s, got %s", tt.mode.expectedIPTablesVStr, mode)
+	}
+}
+
+func (tt iptablesWrapperTest) isIPTablesWrapper(binaryRealPath string) bool {
+	return binaryRealPath == tt.wrapperPath
+}
+
+func (tt iptablesWrapperTest) iptablesRealPath(tb testing.TB) string {
+	return binaryRealPath(tb, tt.iptablesPath)
+}
+
+func (tt iptablesWrapperTest) ip6tablesRealPath(tb testing.TB) string {
+	return binaryRealPath(tb, tt.ip6tablesPath)
+}
+
+func binaryRealPath(tb testing.TB, binary string) string {
+	realPath, err := filepath.EvalSymlinks(binary)
+	assertSuccess(tb, err)
+
+	return realPath
+}
+
+func newIPTablesRunner(ipV iptablesIPVersion, mode iptables.Mode) ipTablesRunner {
+	return ipTablesRunner{
+		binary: string(ipV) + "-" + string(mode),
+	}
+}
+
+type ipTablesRunner struct {
+	binary string
+}
+
+func (r ipTablesRunner) runAndAssertSuccess(ctx context.Context, tb testing.TB, args ...string) {
+	tb.Helper()
+	assertSuccess(tb, r.run(ctx, args...))
+}
+
+func (r ipTablesRunner) run(ctx context.Context, args ...string) error {
+	c := exec.CommandContext(ctx, r.binary, args...)
+	return commands.RunAndReadError(c)
+}
+
+var iptablesModeRegex = regexp.MustCompile(`^ip6?tables.*\((.+)\).*`)
+
+func readIPTablesMode(ctx context.Context, tb testing.TB, iptables string) string {
+	tb.Helper()
+	var out bytes.Buffer
+	c := exec.CommandContext(ctx, iptables, "-V")
+	c.Stdout = &out
+	assertSuccess(tb, commands.RunAndReadError(c))
+
+	outIPTablesVersion := out.String()
+	matches := iptablesModeRegex.FindStringSubmatch(outIPTablesVersion)
+	if len(matches) != 2 {
+		tb.Fatalf("Can't read `%s -V` output format: %s", iptables, outIPTablesVersion)
+	}
+
+	tb.Logf("Output of `%s -V`: %s", iptables, outIPTablesVersion)
+
+	mode := matches[1]
+	return mode
+}
+
+func assertSuccess(tb testing.TB, err error) {
+	tb.Helper()
+	if err != nil {
+		tb.Fatal(err.Error())
+	}
+}

From d2556dcaf87b59c0f4322ad5d20fd12eeb3c677c Mon Sep 17 00:00:00 2001
From: Guillermo Gaston <gaslor@amazon.com>
Date: Wed, 9 Oct 2024 20:53:49 +0000
Subject: [PATCH 2/2] Simplify distroless image by moving installation to
 binary

---
 Makefile                                  |  2 +-
 internal/iptables/alternatives.go         | 26 ++++++++-----
 iptables-wrapper-installer.sh             |  6 +--
 main.go                                   | 46 ++++++++++++++++++++++-
 test/Dockerfile.test-alpine               |  2 +-
 test/Dockerfile.test-debian               |  2 +-
 test/Dockerfile.test-distroless           | 34 +++--------------
 test/build/clean-distroless.sh            | 43 ---------------------
 test/build/stage-binaries-from-package.sh |  4 +-
 test/run-test.sh                          |  2 +
 test/wrapper_test.go                      |  3 ++
 11 files changed, 79 insertions(+), 91 deletions(-)
 delete mode 100755 test/build/clean-distroless.sh

diff --git a/Makefile b/Makefile
index 95a11da..7051f8a 100644
--- a/Makefile
+++ b/Makefile
@@ -21,7 +21,7 @@ fmt: ## Check formatting
 build-tests: $(BIN_DIR)
 	$(GO) test ./test -c -o $(BIN_DIR)/tests
 
-check: check-debian check-debian-nosanity check-debian-backports check-fedora check-alpine
+check: check-debian check-debian-nosanity check-debian-backports check-fedora check-alpine check-distroless
 
 check-debian: build build-tests
 	./test/run-test.sh --build-fail debian
diff --git a/internal/iptables/alternatives.go b/internal/iptables/alternatives.go
index 655ba09..e0b89cf 100644
--- a/internal/iptables/alternatives.go
+++ b/internal/iptables/alternatives.go
@@ -42,7 +42,7 @@ func BuildAlternativeSelector(sbinPath string) AlternativeSelector {
 		return updateAlternativesSelector{sbinPath: sbinPath}
 	} else {
 		// if we don't find any tool to managed the alternatives, handle it manually with symlinks
-		return symlinkSelector{sbinPath: sbinPath}
+		return NewSymlinker(sbinPath)
 	}
 }
 
@@ -79,17 +79,25 @@ func (a alternativesSelector) UseMode(ctx context.Context, mode Mode) error {
 	return nil
 }
 
-// symlinkSelector  manages an iptables setup by manually creating symlinks
-// that point to the proper "mode" binaries.
+func NewSymlinker(sbinPath string) Symlinker {
+	return Symlinker{sbinPath: sbinPath}
+}
+
+// Symlinker manages an iptables setup by manually creating symlinks
+// for all iptables commands.
 // It configures: `iptables`, `iptables-save`, `iptables-restore`,
 // `ip6tables`, `ip6tables-save` and `ip6tables-restore`.
-type symlinkSelector struct {
+type Symlinker struct {
 	sbinPath string
 }
 
-func (s symlinkSelector) UseMode(ctx context.Context, mode Mode) error {
-	modeStr := string(mode)
-	xtablesForModePath := XtablesPath(s.sbinPath, mode)
+// UseMode configures the system to use the selected iptables mode.
+func (s Symlinker) UseMode(ctx context.Context, mode Mode) error {
+	return s.LinkAll(ctx, XtablesPath(s.sbinPath, mode))
+}
+
+// LinkAll creates symlinks for all iptables commands to the targetPath.
+func (s Symlinker) LinkAll(ctx context.Context, targetPath string) error {
 	cmds := []string{"iptables", "iptables-save", "iptables-restore", "ip6tables", "ip6tables-save", "ip6tables-restore"}
 
 	for _, cmd := range cmds {
@@ -97,8 +105,8 @@ func (s symlinkSelector) UseMode(ctx context.Context, mode Mode) error {
 		// If deleting fails, ignore it and try to create symlink regardless
 		_ = os.RemoveAll(cmdPath)
 
-		if err := os.Symlink(xtablesForModePath, cmdPath); err != nil {
-			return fmt.Errorf("creating %s symlink for mode %s: %v", cmd, modeStr, err)
+		if err := os.Symlink(targetPath, cmdPath); err != nil {
+			return fmt.Errorf("creating %s symlink to %s: %v", cmd, targetPath, err)
 		}
 	}
 
diff --git a/iptables-wrapper-installer.sh b/iptables-wrapper-installer.sh
index 36b249f..4ab46f6 100755
--- a/iptables-wrapper-installer.sh
+++ b/iptables-wrapper-installer.sh
@@ -96,9 +96,9 @@ case "${altstyle}" in
             --install /usr/sbin/iptables iptables /usr/sbin/iptables-wrapper 100 \
             --slave /usr/sbin/iptables-restore iptables-restore /usr/sbin/iptables-wrapper \
             --slave /usr/sbin/iptables-save iptables-save /usr/sbin/iptables-wrapper \
-            --slave /usr/sbin/ip6tables iptables /usr/sbin/iptables-wrapper \
-            --slave /usr/sbin/ip6tables-restore iptables-restore /usr/sbin/iptables-wrapper \
-            --slave /usr/sbin/ip6tables-save iptables-save /usr/sbin/iptables-wrapper
+            --slave /usr/sbin/ip6tables ip6tables /usr/sbin/iptables-wrapper \
+            --slave /usr/sbin/ip6tables-restore ip6tables-restore /usr/sbin/iptables-wrapper \
+            --slave /usr/sbin/ip6tables-save ip6tables-save /usr/sbin/iptables-wrapper
 	;;
 
     debian)
diff --git a/main.go b/main.go
index caebb05..1407125 100644
--- a/main.go
+++ b/main.go
@@ -32,6 +32,11 @@ it will enter an infinite loop, calling itself recursively.
 
 It's important to note that this proxy behavior will only happen on the first iptables-*
 execution. Following invocations will use directly the binaries for the selected mode.
+
+If the command is executed with the `install` argument, it will create symlinks for all iptables binaries
+pointing to itself. The second argument must be the path where the symlinks should be installed, preferably
+the same path where iptables is installed.
+This is useful for the installation process of the iptables-wrapper itself before first execution.
 */
 package main
 
@@ -41,6 +46,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"path/filepath"
 
 	"github.com/kubernetes-sigs/iptables-wrappers/internal/iptables"
 )
@@ -48,10 +54,41 @@ import (
 func main() {
 	ctx := context.Background()
 
+	if installMode() {
+		install(ctx)
+		return
+	}
+
+	forward(ctx)
+}
+
+func installMode() bool {
+	return len(os.Args) == 3 && os.Args[1] == "install"
+}
+
+func installFolder() string {
+	return os.Args[2]
+}
+
+// install creates symlinks for all iptables binaries in the given folder
+// pointing to the current binary being executed.
+func install(ctx context.Context) {
+	wrapperPath, err := os.Executable()
+	if err != nil {
+		fatal(err)
+	}
+	wrapperPath = filepath.Clean(wrapperPath)
+
+	if err := iptables.NewSymlinker(installFolder()).LinkAll(ctx, wrapperPath); err != nil {
+		fatal(err)
+	}
+}
+
+// forward detects the iptables mode to use and re-executes the exact same command passed to this program.
+func forward(ctx context.Context) {
 	sbinPath, err := iptables.DetectBinaryDir()
 	if err != nil {
-		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
-		os.Exit(1)
+		fatal(err)
 	}
 
 	// We use `xtables-<mode>-multi` binaries by default to inspect the installed rules,
@@ -91,3 +128,8 @@ func main() {
 		os.Exit(code)
 	}
 }
+
+func fatal(err error) {
+	fmt.Fprintf(os.Stderr, "Error: %s\n", err)
+	os.Exit(1)
+}
diff --git a/test/Dockerfile.test-alpine b/test/Dockerfile.test-alpine
index b7cc734..45363bf 100644
--- a/test/Dockerfile.test-alpine
+++ b/test/Dockerfile.test-alpine
@@ -16,7 +16,7 @@
 
 FROM alpine:3.15
 
-RUN apk add --no-cache iptables
+RUN apk add --no-cache iptables ip6tables
 COPY iptables-wrapper-installer.sh /
 COPY bin/iptables-wrapper /
 RUN /iptables-wrapper-installer.sh
diff --git a/test/Dockerfile.test-debian b/test/Dockerfile.test-debian
index ad86f8e..14b69c4 100644
--- a/test/Dockerfile.test-debian
+++ b/test/Dockerfile.test-debian
@@ -19,7 +19,7 @@ FROM debian:buster
 ARG INSTALL_ARGS=
 ARG REPO=buster
 
-RUN echo deb http://deb.debian.org/debian buster-backports main >> /etc/apt/sources.list; \
+RUN echo deb http://archive.debian.org/debian buster-backports main >> /etc/apt/sources.list; \
     apt-get update; \
     apt-get -t ${REPO} -y --no-install-recommends install iptables
 
diff --git a/test/Dockerfile.test-distroless b/test/Dockerfile.test-distroless
index bd3ffec..ce4cc50 100644
--- a/test/Dockerfile.test-distroless
+++ b/test/Dockerfile.test-distroless
@@ -17,7 +17,7 @@
 
 ARG STAGE_DIR="/opt/stage"
 
-FROM debian:bullseye-slim as build
+FROM debian:bullseye-slim AS build
 ARG STAGE_DIR
 
 COPY test/build/stage-binaries-from-package.sh /
@@ -25,36 +25,12 @@ COPY test/build/package-utils.sh /
 COPY test/build/stage-binary-and-deps.sh /
 
 RUN mkdir -p "${STAGE_DIR}" && \
-    /stage-binaries-from-package.sh "${STAGE_DIR}" conntrack \
-    ebtables    \
-    ipset       \
-    iptables    \
-    kmod		&& \
-	`# below binaries and dash are used by iptables-wrapper-installer.sh` \
-    /stage-binary-and-deps.sh "${STAGE_DIR}" /bin/dash \
-    /bin/mv \
-    /bin/chmod \
-    /bin/grep \
-    /bin/ln  \
-    /bin/rm \
-    /bin/sleep \
-    /usr/bin/wc
+    /stage-binaries-from-package.sh "${STAGE_DIR}" iptables
 
-RUN ln -sf /bin/dash "${STAGE_DIR}"/bin/sh
-
-FROM gcr.io/distroless/static as intermediate
+FROM gcr.io/distroless/static
 ARG STAGE_DIR
 
-COPY test/build/clean-distroless.sh /clean-distroless.sh
 COPY --from=build "${STAGE_DIR}" /
-COPY iptables-wrapper-installer.sh /
-COPY bin/iptables-wrapper /
-# iptables-wrapper-installer needs to know that iptables exists before doing all its magic
-RUN echo "" >  /usr/sbin/iptables && \
-	/iptables-wrapper-installer.sh && \
-    /clean-distroless.sh
-
-FROM scratch
-
-COPY --from=intermediate / /
+COPY bin/iptables-wrapper /usr/sbin/iptables-wrapper
+RUN ["/usr/sbin/iptables-wrapper", "install", "/usr/sbin"]
 COPY bin/tests /
diff --git a/test/build/clean-distroless.sh b/test/build/clean-distroless.sh
deleted file mode 100755
index b46844c..0000000
--- a/test/build/clean-distroless.sh
+++ /dev/null
@@ -1,43 +0,0 @@
-#!/bin/sh
-
-# Copyright 2022 The Kubernetes Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# USAGE: clean-distroless.sh
-
-# Modified version of https://github.com/kubernetes/release/blob/master/images/build/distroless-iptables/distroless/clean-distroless.sh
-
-REMOVE="/usr/share/base-files
-/usr/share/man
-/usr/lib/*-linux-gnu/gconv/
-/usr/bin/c_rehash
-/usr/bin/openssl
-/bin/mv
-/bin/chmod
-/bin/grep
-/bin/ln
-/bin/sleep
-/usr/bin/wc
-/iptables-wrapper-installer.sh
-/bin/sh
-/bin/dash
-/clean-distroless.sh
-/bin/rm"
-
-IFS="
-"
-
-for item in ${REMOVE}; do
-    rm -rf "${item}"
-done
diff --git a/test/build/stage-binaries-from-package.sh b/test/build/stage-binaries-from-package.sh
index b08bbaa..cb664dd 100755
--- a/test/build/stage-binaries-from-package.sh
+++ b/test/build/stage-binaries-from-package.sh
@@ -49,11 +49,11 @@ get_dependent_packages() {
 main() {
     STAGE_DIR="${1}/"
     mkdir -p "${STAGE_DIR}"/var/lib/dpkg/status.d/
-    apt -y update
+    apt-get -y update
     shift
     while (( "$#" )); do        # While there are arguments still to be shifted
         PACKAGE="${1}"
-        apt -y install "${PACKAGE}"
+        apt-get -y install "${PACKAGE}"
         stage_file_list "${PACKAGE}" "$STAGE_DIR"
         while IFS= read -r c_dep; do
             stage_file_list "${c_dep}" "${STAGE_DIR}"
diff --git a/test/run-test.sh b/test/run-test.sh
index b08a6df..b813357 100755
--- a/test/run-test.sh
+++ b/test/run-test.sh
@@ -112,6 +112,8 @@ if ! build "${tag}" ${build_arg}; then
     FAIL "build failed unexpectedly"
 fi
 
+# We execute each test in a separate container to avoid having to cleanup after each test and just start with a clean slate.
+
 if ! docker run --privileged "iptables-wrapper-test-${tag}" /tests -test.v -test.run "^TestIPTablesWrapperLegacy$" ; then
     FAIL "failed legacy iptables"
 fi
diff --git a/test/wrapper_test.go b/test/wrapper_test.go
index 2af20f1..d052503 100644
--- a/test/wrapper_test.go
+++ b/test/wrapper_test.go
@@ -166,14 +166,17 @@ func (tt iptablesWrapperTest) isIPTablesWrapper(binaryRealPath string) bool {
 }
 
 func (tt iptablesWrapperTest) iptablesRealPath(tb testing.TB) string {
+	tb.Helper()
 	return binaryRealPath(tb, tt.iptablesPath)
 }
 
 func (tt iptablesWrapperTest) ip6tablesRealPath(tb testing.TB) string {
+	tb.Helper()
 	return binaryRealPath(tb, tt.ip6tablesPath)
 }
 
 func binaryRealPath(tb testing.TB, binary string) string {
+	tb.Helper()
 	realPath, err := filepath.EvalSymlinks(binary)
 	assertSuccess(tb, err)