From d40c0c63623ac13f705b460dea5412bc2fc0bae0 Mon Sep 17 00:00:00 2001
From: Sergey Solodyagin <solodyagin@gmail.com>
Date: Fri, 29 Dec 2023 12:13:20 +0300
Subject: [PATCH] add the IP sanitize method

---
 README.md        |  2 +-
 sanitize.go      | 22 +++++++++++++++++++---
 sanitize_test.go | 40 ++++++++++++++++++++++++++++++++++++----
 3 files changed, 56 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 4dcba81..71d90b1 100644
--- a/README.md
+++ b/README.md
@@ -8,4 +8,4 @@ go get -u github.com/andoma-go/sanitize
 
 ## Documentation
 
-View the generated [documentation](https://pkg.go.dev/github.com/andoma-go/sanitize)
+View the generated [documentation](https://pkg.go.dev/github.com/andoma-go/sanitize#readme-documentation)
diff --git a/sanitize.go b/sanitize.go
index b6dd1a9..4f88097 100644
--- a/sanitize.go
+++ b/sanitize.go
@@ -1,11 +1,18 @@
 package sanitize
 
-// Integer returns only numbers
+import (
+	"net"
+	"regexp"
+)
+
+var reIP = regexp.MustCompile(`[^a-zA-Z0-9:.]`) // IPv4 and IPv6 characters only
+
+// Integer returns numbers only
 func Integer(input string) string {
 	runes := []rune{}
 	n := 0
 	for _, r := range input {
-		if !(r == 45) && !(r >= 48 && r <= 57) { // regexp.MustCompile(`[^-\d]`).ReplaceAllString(input, "$1")
+		if !(r == 45) && !(r >= 48 && r <= 57) { // Ignore any invalid characters first
 			continue
 		}
 		switch r {
@@ -30,7 +37,7 @@ func Float(input string) string {
 	dot := false
 	n := 0
 	for i, r := range input {
-		if !(r >= 44 && r <= 46) && !(r >= 48 && r <= 57) { // regexp.MustCompile(`[^-\.,\d]`).ReplaceAllString(input, "$1")
+		if !(r >= 44 && r <= 46) && !(r >= 48 && r <= 57) { // Ignore any invalid characters first
 			continue
 		}
 		switch r {
@@ -58,3 +65,12 @@ func Float(input string) string {
 	}
 	return ""
 }
+
+// IP returns the IP address in IPv4 or IPv6 format
+func IP(input string) string {
+	sanitized := reIP.ReplaceAllString(input, "")
+	if ip := net.ParseIP(sanitized); ip != nil {
+		return ip.String()
+	}
+	return ""
+}
diff --git a/sanitize_test.go b/sanitize_test.go
index 38c6101..e9d51c0 100644
--- a/sanitize_test.go
+++ b/sanitize_test.go
@@ -6,7 +6,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-// TestInteger tests the integer sanitize method
+// TestInteger tests the Integer sanitize method
 func TestInteger(t *testing.T) {
 	t.Parallel()
 
@@ -26,14 +26,14 @@ func TestInteger(t *testing.T) {
 	}
 }
 
-// BenchmarkInteger benchmarks the integer method
+// BenchmarkInteger benchmarks the Integer method
 func BenchmarkInteger(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		_ = Integer("a-bc12.3def987asdf--")
 	}
 }
 
-// TestFloat tests the float sanitize method
+// TestFloat tests the Float sanitize method
 func TestFloat(t *testing.T) {
 	t.Parallel()
 
@@ -53,9 +53,41 @@ func TestFloat(t *testing.T) {
 	}
 }
 
-// BenchmarkFloat benchmarks the float method
+// BenchmarkFloat benchmarks the Float method
 func BenchmarkFloat(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		_ = Float("a-bc12.3def987asdf--")
 	}
 }
+
+// TestIP tests the IP sanitize method
+func TestIP(t *testing.T) {
+	t.Parallel()
+
+	var tests = []struct {
+		input    string
+		expected string
+	}{
+		{"192.168.0.1", "192.168.0.1"},
+		{"192.168.0.256", ""},
+		{"192.168", ""},
+		{"IP: 192.168.0.1", ""},
+		{"2001:0db8:85a3:0000:0000:8a2e:0370:7334", "2001:db8:85a3::8a2e:370:7334"},
+		{"2001:db8:0:0:1:0:0:1", "2001:db8::1:0:0:1"},
+		{"d8c4:e12f:eff7:df64:3995:df25:aaca:49eb", "d8c4:e12f:eff7:df64:3995:df25:aaca:49eb"},
+		{"6b71:3208:07ad:0629:a150:5734:15e8:950d", "6b71:3208:7ad:629:a150:5734:15e8:950d"},
+		{"::ffff:c000:0280", "192.0.2.128"},
+	}
+
+	for _, test := range tests {
+		output := IP(test.input)
+		assert.Equal(t, test.expected, output)
+	}
+}
+
+// BenchmarkIP benchmarks the IP method
+func BenchmarkIP(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		_ = IP("192.168.0.1")
+	}
+}