From be1b3fbf73ffd42202f45f62779f5d5c6a2c1dd1 Mon Sep 17 00:00:00 2001
From: Yair Slobodin <154875779+YairSlobodin1@users.noreply.github.com>
Date: Tue, 22 Oct 2024 17:44:21 +0300
Subject: [PATCH] added AsCidr method (#78)

* added AsCidr method

* fixed doc
---
 pkg/netset/ipblock.go      | 10 ++++++++++
 pkg/netset/ipblock_test.go |  9 +++++++++
 2 files changed, 19 insertions(+)

diff --git a/pkg/netset/ipblock.go b/pkg/netset/ipblock.go
index 1dcb173..164a336 100644
--- a/pkg/netset/ipblock.go
+++ b/pkg/netset/ipblock.go
@@ -327,6 +327,16 @@ func cidrToInterval(cidr string) (interval.Interval, error) {
 	return interval.New(int64(startNum), int64(endNum)), nil
 }
 
+// AsCidr returns the CIDR string of this IPBlock object, if it contains exactly one CIDR,
+// otherwise it returns an error
+func (b *IPBlock) AsCidr() (string, error) {
+	cidrList := b.ToCidrList()
+	if len(cidrList) != 1 {
+		return "", fmt.Errorf("ipblock contains %d cidrs", len(cidrList))
+	}
+	return cidrList[0], nil
+}
+
 // ToCidrList returns a list of CIDR strings for this IPBlock object
 func (b *IPBlock) ToCidrList() []string {
 	var cidrList []string
diff --git a/pkg/netset/ipblock_test.go b/pkg/netset/ipblock_test.go
index 83af0ce..4763fb8 100644
--- a/pkg/netset/ipblock_test.go
+++ b/pkg/netset/ipblock_test.go
@@ -100,6 +100,15 @@ func TestConversions(t *testing.T) {
 	require.Equal(t, ipRange, toPrint[0])
 
 	require.Equal(t, "", ipb1.ToIPAddressString())
+
+	_, err = ipb1.AsCidr()
+	require.NotNil(t, err)
+
+	cidr := "5.2.1.0/24"
+	ipb3, _ := netset.IPBlockFromCidr(cidr)
+	str, err := ipb3.AsCidr()
+	require.Nil(t, err)
+	require.Equal(t, str, cidr)
 }
 
 func TestDisjointIPBlocks(t *testing.T) {