From 7963634f831561bd5595fecd1c9de7d2547cfc54 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 |  57 +++++++++++
 lib/puppet_x/extlib/cidr.rb         | 151 ++++++++++++++++++++++++++++
 2 files changed, 208 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..edeb28c
--- /dev/null
+++ b/lib/puppet/datatypes/extlib/cidr.rb
@@ -0,0 +1,57 @@
+# 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],
+      address => { type => String[1], kind => derived },
+      canonical => { type => String[1], kind => derived },
+      netmask => { type => String[1], kind => derived },
+      family => { type => String[1], kind => derived },
+      prefix_length => { type => Integer[0], kind => derived },
+      is_single_address => { type => Boolean, kind => derived },
+      is_link_local => { type => Boolean, kind => derived },
+      is_loopback => { type => Boolean, kind => derived },
+      is_private => { type => Boolean, kind => derived },
+      first => { type => String[1], kind => derived },
+      last => { type => String[1], kind => derived },
+      count => { type => Integer[0], kind => derived },
+      host_count => { type => Integer[0], kind => derived },
+      host_min => { type => String[1], kind => derived },
+      host_max => { type => String[1], kind => derived },
+      host_next => { type => String[1], kind => derived },
+      host_prev => { type => String[1], kind => derived },
+    },
+    functions => {
+      host_rewind => Callable[[String[1]], 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..70d0f65
--- /dev/null
+++ b/lib/puppet_x/extlib/cidr.rb
@@ -0,0 +1,151 @@
+# 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
+        @single_address = (@ip.ipv4? && @ip.prefix == 32) || (@ip.ipv6? && @ip.prefix == 128)
+
+        if (@ip.ipv4? && @ip.prefix.between?(31, 32)) || (@ip.ipv6? && @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 = @first
+      end
+
+      def address
+        @ip.to_s
+      end
+
+      def canonical
+        "#{@ip.to_string}/#{@ip.prefix}"
+      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 is_single_address
+        @single_address
+      end
+
+      def is_link_local
+        @ip.link_local?
+      end
+
+      def is_loopback
+        @ip.loopback?
+      end
+
+      def is_private
+        @ip.private?
+      end
+
+      def first
+        @first.to_s
+      end
+
+      def last
+        @last.to_s
+      end
+
+      # Total amount of addresses in the range
+      def count
+        @last.to_i - @first.to_i + 1
+      end
+
+      # Amount of usable addresses in the range (except first and last addresses)
+      def host_count
+        @host_max.to_i - @host_min.to_i + 1
+      end
+
+      def host_rewind(new_ip)
+        raise ArgumentError, "IP should be between #{@first} and #{@last}" unless new_ip.between?(@first, @last)
+
+        @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
+
+      # If our CIDR is just an IP address then just return next one.
+      # Don't care about network range.
+      def host_next_by_step(step = 1)
+        next_ip = IPAddr.new(@host_current.to_i + step, @ip.family)
+
+        if !@single_address
+          if !next_ip.between?(@host_min, @host_max)
+            return nil
+          end
+        end
+
+        @host_current = next_ip
+        return @host_current.to_s
+      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
+        if @single_address
+          @ip.to_s
+        else
+          "#{@ip.to_s}/#{@ip.prefix}"
+        end
+      end
+
+      def to_i
+        @ip.to_i
+      end
+
+      # Compare IP addresses
+      def <=>(other)
+        @ip.to_i <=> if other.is_a? Extlib::CIDR
+                       other.to_i
+                     else
+                       IPAddr.new(other).to_i
+                     end
+      end
+      alias compare <=>
+    end
+  end
+end