ip2net

#!/usr/bin/env python3
# ip2net (part of ossobv/vcutil) // wdoekes/2023 // Public Domain
#
# Reduce/merge IPs to lists of larger subnets. This is useful when
# you have a list of IPs, from e.g. CloudFront that you want to place in
# an ACL. If you don't want huge ACLs, you might be able to make them a
# bit smaller using ip2net.
#
# Right now some values are hardcoded based on a CloudFront sample:
# - if two IPs in a /24 subnet match, they are placed into such a subnet
# - subnets will not get larger than /15
# - between /24 and /15 subnets will be merged if they are consecutive
#   (e.g. 1.1.2.0/24 and 1.1.3.0/24, but not 1.1.3.0/24 ansd 1.1.4.0/24)
#
# You'll likely want to use the --unique (-u) option, unless you want to
# see how many IPs are in use. In that case you might do:
#
#   $ ip2net < ips | sort -V | uniq -c
#
# Example usage:
#
#   $ for ip in 192.168.1.1 1.2.2.4 1.2.2.5 1.2.3.10 1.2.3.128; do
#       echo $ip
#     done | ip2net -u
#   1.2.2.0/23
#   192.168.1.1
#
# A more practical example, where the github action IPs/networks are
# reduced from 3500 ranges to 700 ranges:
#
#   $ curl -sSfL https://api.github.com/meta | jq -r .actions[]
#   ...
#   40.123.140.0/22
#   40.123.144.0/26
#   40.123.144.64/29
#   ...
#
#   $ curl -sSfL https://api.github.com/meta | jq -r .actions[] |
#     wc -l
#   3583
#
#   $ curl -sSfL https://api.github.com/meta | jq -r .actions[] |
#     ip2net | sort -V | uniq -c | sort -nr | head -n3
#       351 2603:1030::/32
#       174 40.123.0.0/16
#       173 40.86.0.0/15
#
#   $ curl -sSfL https://api.github.com/meta | jq -r .actions[] |
#     ip2net -u | wc -l
#   743
#
import sys

from argparse import ArgumentParser, RawDescriptionHelpFormatter
from ipaddress import IPv4Network, IPv6Network, ip_network
from os import environ
from unittest import TestCase, main as main_unittest


def ipreduce(nets, smallest_prefix, largest_prefix):
    nets.sort()

    new = []
    idx = 0
    while idx < len(nets):
        net = nets[idx]

        # Make at most /15 (largest_prefix), make at least /24
        # (smallest_prefix).
        supernet = net
        if supernet.prefixlen > largest_prefix:
            supernet = supernet.supernet()
        while supernet.prefixlen > smallest_prefix:
            supernet = supernet.supernet()

        idx2 = idx + 1
        while idx2 < len(nets) and nets[idx2].subnet_of(supernet):
            idx2 += 1
        if idx2 == idx + 1:
            new.append(net)
            idx += 1
        else:
            new.append(supernet)
            idx = idx2
    return new


def process(nets, reduce_func):
    nets = list(set(nets))
    old_nets = []
    new_nets = nets

    while len(old_nets) != len(new_nets):
        old_nets = new_nets
        new_nets = reduce_func(old_nets)

    # Check and translate.
    translation = {}
    for net in nets:
        for n in new_nets:
            if net.subnet_of(n):
                translation[net] = n
                break
        else:
            assert False, (net, 'not in', new_nets)

    return translation


def print_nets(nets, verbose=False):
    for net in nets:
        if net.num_addresses == 1 and not verbose:
            print(net.network_address.compressed)
        else:
            print(net.with_prefixlen)


def compress_ips(ips, args):
    ip4s = [i for i in ips if isinstance(i, IPv4Network)]
    ip6s = [i for i in ips if isinstance(i, IPv6Network)]

    translation4 = process(ip4s, (lambda x: ipreduce(x, 24, 15)))
    translation6 = process(ip6s, (lambda x: ipreduce(x, 32, 48)))

    nets4 = [translation4[ip] for ip in ip4s]
    nets6 = [translation6[ip] for ip in ip6s]

    if args.unique:
        nets4 = list(sorted(set(nets4)))
        nets6 = list(sorted(set(nets6)))

    print_nets(nets4, verbose=args.verbose)
    print_nets(nets6, verbose=args.verbose)


class Ip2NetTest(TestCase):
    def _collect(self, *args, **kwargs):
        self._collected.append((args, kwargs))

    def _expect_simple_prints(self):
        ret = []
        for args, kwargs in self._collected:
            assert len(args) == 1, (args, kwargs)
            assert len(kwargs) == 0, (args, kwargs)
            ret.append(args[0])
        return ret

    def setUp(self):
        global print
        self._old_print = print
        self._collected = []
        print = self._collect

    def tearDown(self):
        global print
        print = self._old_print

    def _2net4(self, ips):
        return [IPv4Network(ip) for ip in ips]

    def _2net6(self, ips):
        return [IPv6Network(ip) for ip in ips]

    def _run_ip4_process_test(self, input_, expected, unique=True):
        ip4s = self._2net4(input_)
        expected = self._2net4(expected)
        translation4 = process(ip4s, (lambda x: ipreduce(x, 24, 15)))
        nets4 = [translation4[ip] for ip in ip4s]
        if unique:
            nets4 = list(set(nets4))
            nets4.sort()
        self.assertEqual(expected, nets4)

    def test_process_1(self):
        self._run_ip4_process_test(
            ('10.10.10.10 8.8.8.8 10.10.10.20 192.168.0.255 '
             '10.10.10.30').split(),
            ('10.10.10.0/24 8.8.8.8 10.10.10.0/24 192.168.0.255 '
             '10.10.10.0/24').split(),
            unique=False)

    def test_process_2(self):
        # This auto-sorts, but two records are not improved by merging them.
        self._run_ip4_process_test(
            ['192.168.2.1', '20.20.20.20'], ['20.20.20.20', '192.168.2.1'])

    def test_process_3(self):
        # 1.1.1.0/24 + 1.1.2.0/24 does NOT become 1.1.0.0/22
        # 1.1.2.0/24 + 1.1.3.0/24 does become     1.1.2.0/23
        self._run_ip4_process_test(
            '1.1.1.5 1.1.1.8 1.1.2.5 1.1.2.6 1.1.2.7 1.1.3.8'.split(),
            '1.1.1.0/24 1.1.2.0/23'.split())
        # Here we do have 1.1.0.0/24 + 1.1.1.0/24 + 1.1.2.0/24 + 1.1.3.0/24.
        self._run_ip4_process_test(
            '1.1.1.5 1.1.0.8 1.1.0.9 1.1.2.5 1.1.2.6 1.1.2.7 1.1.3.8'.split(),
            '1.1.0.0/22'.split())
        # BUG: This is kind of flaky:
        # - 1.1.0.8 will not merge into a /24 on its own,
        # - 10.1.0.0/24 + 1.1.1.1 will become a /23
        self._run_ip4_process_test(
            '1.1.0.8 1.1.1.1'.split(),
            '1.1.0.8 1.1.1.1'.split())
        self._run_ip4_process_test(
            '1.1.0.8 1.1.1.1 1.1.1.2'.split(),
            '1.1.0.8 1.1.1.0/24'.split())
        self._run_ip4_process_test(
            '1.1.0.8 1.1.0.9 1.1.1.2'.split(),
            '1.1.0.0/23'.split())  # not 10.1.0.0/24 + 1.1.1.2 (?)

    def test_print_nets(self):
        print_nets(self._2net4([
            '1.2.3.0/24', '1.2.3.0/24', '1.2.3.0/24', '5.5.5.5']))
        result = self._expect_simple_prints()
        self.assertEqual(
            ['1.2.3.0/24', '1.2.3.0/24', '1.2.3.0/24', '5.5.5.5'],
            result)

    def test_compress_ips_badval(self):
        class args:
            unique = False
            verbose = False

        compress_ips(['aap', 'noot', 'mies'], args=args)
        result = self._expect_simple_prints()
        self.assertEqual([], result)

    def test_compress_ips(self):
        class args:
            unique = False
            verbose = False

        compress_ips(self._2net4([
            '1.2.3.0/24', '5.5.5.5/32', '1.2.3.0/24', '5.5.5.5']),
            args=args)
        result = self._expect_simple_prints()
        self.assertEqual(
            ['1.2.3.0/24', '5.5.5.5', '1.2.3.0/24', '5.5.5.5'], result)

    def test_compress_ips_u(self):
        class args:
            unique = True
            verbose = False

        compress_ips(self._2net4([
            '1.2.3.0/24', '1.2.3.0/24', '1.2.3.0/24', '5.5.5.5']),
            args=args)
        result = self._expect_simple_prints()
        self.assertEqual(['1.2.3.0/24', '5.5.5.5'], result)

    def test_print_nets_u_v(self):
        class args:
            unique = True
            verbose = True

        compress_ips(self._2net4([
            '1.2.3.0/24', '1.2.3.0/24', '1.2.3.0/24', '5.5.5.5']),
            args=args)
        result = self._expect_simple_prints()
        self.assertEqual(['1.2.3.0/24', '5.5.5.5/32'], result)


def main():
    parser = ArgumentParser()
    parser = ArgumentParser(
        formatter_class=RawDescriptionHelpFormatter,
        description='''\
Compress multiple IPs to fewer networks.

Reads IPs from stdin, outputs networks to stdout.

For instance:

    $ for ip in 192.168.1.1 1.2.2.4 1.2.2.5 1.2.3.10 1.2.3.128; do
        echo $ip
      done | ip2net -u

    1.2.2.0/23
    192.168.1.1
''')
    parser.add_argument('-u', '--unique', action='store_true', help=(
        'output only the first of equal networks'))
    parser.add_argument('-v', '--verbose', action='store_true', help=(
        'output network prefix even for single IPs'))
    args = parser.parse_args()
    ips = [ip_network(line.rstrip('\r\n')) for line in sys.stdin]
    compress_ips(ips, args)


if __name__ == '__main__':
    if environ.get('RUNTESTS', '') not in ('', '0'):
        main_unittest()
        raise RuntimeError('(unreachable code)')

    main()