From f97e1a6e9ac29ff8a30f6b4c57dc46e3f826ed98 Mon Sep 17 00:00:00 2001
From: Yury Bushmelev <jay4mail@gmail.com>
Date: Mon, 1 Apr 2024 20:51:02 +0300
Subject: [PATCH] Add Extlib::CIDR DataType

Example:
```
$ip = Extlib::CIDR('192.168.1.0/24')
notice $ip.host_min, $ip.host_max
```
---
 lib/puppet/datatypes/extlib/cidr.rb |  52 ++++++++++++
 lib/puppet_x/extlib/cidr.rb         | 120 ++++++++++++++++++++++++++++
 2 files changed, 172 insertions(+)
 create mode 100644 lib/puppet/datatypes/extlib/cidr.rb
 create mode 100644 lib/puppet_x/extlib/cidr.rb

diff --git a/lib/puppet/datatypes/extlib/cidr.rb b/lib/puppet/datatypes/extlib/cidr.rb
new file mode 100644
index 0000000..fdcb16a
--- /dev/null
+++ b/lib/puppet/datatypes/extlib/cidr.rb
@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+# The `Extlib::CIDR` object represents a CIDR.
+#
+# @param cidr
+# @!method address
+# @!method canonical
+# @!method netmask
+# @!method family
+# @!method prefix_length
+# @!method first
+# @!method last
+# @!method host_rewind
+# @!method host_next
+# @!method host_prev
+# @!method host_count
+# @!method host_min
+# @!method host_max
+# @!method to_i
+# @!method compare
+# @!method include
+#
+Puppet::DataTypes.create_type('Extlib::CIDR') do
+  interface <<-PUPPET
+    attributes => {
+      cidr => String[1],
+    },
+    functions => {
+      address => Callable[[], String[1]],
+      canonical => Callable[[], String[1]],
+      netmask => Callable[[], String[1]],
+      family => Callable[[], String[1]],
+      prefix_length => Callable[[], Integer[0]],
+      first => Callable[[], String[1]],
+      last => Callable[[], String[1]],
+      host_rewind => Callable[[String[1]], String[1]],
+      host_next => Callable[[], Optional[String[1]]],
+      host_prev => Callable[[], Optional[String[1]]],
+      host_count => Callable[[], Integer[0]],
+      host_min => Callable[[], String[1]],
+      host_max => Callable[[], String[1]],
+      to_i => Callable[[], Integer[0]],
+      compare => Callable[[Variant[Extlib::CIDR, String[1]]], Integer[-1,1]],
+      include => Callable[[Variant[Extlib::CIDR, String[1]]], Boolean],
+    }
+  PUPPET
+
+  load_file('puppet_x/extlib/cidr')
+
+  PuppetX::Extlib::CIDR.include(Puppet::Pops::Types::PuppetObject)
+  implementation_class PuppetX::Extlib::CIDR
+end
diff --git a/lib/puppet_x/extlib/cidr.rb b/lib/puppet_x/extlib/cidr.rb
new file mode 100644
index 0000000..cfe0037
--- /dev/null
+++ b/lib/puppet_x/extlib/cidr.rb
@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+require 'ipaddr'
+
+module PuppetX
+  class Extlib
+    class CIDR
+      attr_reader :ip
+
+      def initialize(cidr)
+        @ip = IPAddr.new(cidr)
+        @first = @ip.to_range.first
+        @last = @ip.to_range.last
+
+        if (@ip.ipv4? and @ip.prefix.between?(31,32)) or (@ip.ipv6? and @ip.prefix.between?(127,128))
+          @host_min = @first
+          @host_max = @last
+        else
+          @host_min = IPAddr.new(@first.to_i + 1, @ip.family)
+          @host_max = IPAddr.new(@last.to_i - 1, @ip.family)
+        end
+        @host_current = @host_min
+      end
+
+      def address
+        @ip.to_s
+      end
+
+      def canonical
+        @ip.to_string
+      end
+
+      def netmask
+        @ip.inspect.gsub(%r{.*/(.*)>}, '\1')
+      end
+
+      # ipv4/ipv6
+      def family
+        @ip.inspect.split(':')[1].strip.downcase
+      end
+
+      def prefix_length
+        @ip.prefix
+      end
+
+      def first
+        @first.to_s
+      end
+
+      def last
+        @last.to_s
+      end
+
+      def host_rewind(new_ip)
+        raise ArgumentError, "IP should be between #{@host_min} and #{@host_max}" unless new_ip.between?(@host_min, @host_max)
+        @host_current = IPAddr.new(new_ip)
+        new_ip.to_s
+      end
+
+      def host_next
+        host_next_by_step(1)
+      end
+
+      def host_prev
+        host_next_by_step(-1)
+      end
+
+      def host_next_by_step(step = 1)
+        ret_ip = @host_current
+        next_ip = IPAddr.new(@host_current.to_i + step, @ip.family)
+        # If our CIDR is just an IP address then just return next one.
+        # Don't care about network range.
+        if (@ip.ipv4? and @ip.prefix == 32) or (@ip.ipv6? and @ip.prefix == 128) or @host_current.between?(@host_min, @host_max)
+          @host_current = next_ip
+          return ret_ip.to_s
+        end
+        nil
+      end
+
+      def host_count
+        @host_max.to_i - @host_min.to_i + 1
+      end
+
+      def host_min
+        @host_min.to_s
+      end
+
+      def host_max
+        @host_max.to_s
+      end
+
+      def include?(other)
+        if other.is_a? Extlib::CIDR
+          @ip.include? other.ip
+        else
+          @ip.include? IPAddr.new(other)
+        end
+      end
+      alias include include?
+
+      def to_s
+        @ip.to_s
+      end
+
+      def to_i
+        @ip.to_i
+      end
+
+      # Compare IP addresses
+      def <=>(other)
+        if other.is_a? Extlib::CIDR
+          @ip.to_i <=> other.to_i
+        else
+          @ip.to_i <=> IPAddr.new(other).to_i
+        end
+      end
+      alias compare <=>
+    end
+  end
+end