#!/usr/bin/perl -p
# save-combine - filter and combine output from ipset(1) save
# Copyright (C)2026 Corwin Brust <corwin@bru.st>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
#################
#
# Use a flat-file "DB" of ipset per row in the form NAME MAX
# Take DB file-name from first program arg, otherwise "ipset-counts.txt"
# Drop create/add lines for any set not mentioned in the DB
# Input is expected on stdin in the same format as output by ipset save
#
#################

use strict;
use warnings;

# if no DB file
our ($default_db_file) =  q(max-counts.txt);;

# specify as -# or --check=# (e.g. -0 or --check=0)
# 0, no; 1, -exist; 2, ipset test, 3, both exit and test
our ($check);

our @check_bits = map unpack("B*", pack("N", $_)), 0,1;

# whether to print to STDERR for each add skipped
# this only affects $check = 2 (--check=2/-2)
our ($checkwarn);

our (%m,%h);

sub usage {
    my $message = shift;
    $message .= "\n" if $message;
    die <<END_OF_USAGE
${message}USAGE: ipset save >fil && \\
       $0 [max-count-file] <fil
END_OF_USAGE
}

# return the first power of two larger than argument
sub max2 {
  my $val = $_[0] || 1;
  my $rv = 2;
  $rv *= 2 until $rv > $val;
  return $rv;
}

sub db {
  my %m;
  my( $file ) = grep defined&&length,(
    @_,
    $default_db_file,
    q(max-counts.txt)
  ); #$default_db_file;
  open my$FH, '<', $file
      or usage( q(ERROR: cannot open DB)
		. qq( "$file": $!)
		. ' ('.( 0+$! ).')'
	    );
  while (my $line = <$FH>) {
    chomp $line;
    if ($line) {
      my($k,$v) = split /\s+/, $line;
      if ($k and $v) {
	$m{$k} = $v;
      }
    }
  }
  return %m;
}

BEGIN # program is a filter so we must wrap start-up processing
{
    # avoid extra "BEGIN failed--" messages
    $SIG{__DIE__} = sub {warn @_; exit 1};

    # display usage if requested
    usage() if grep /^-+[?h]/, @ARGV;

    # process options (must come first, must start with -)
    while (@ARGV and $ARGV[0] =~ /^-+(.*)/) {
	local $_ = $1;
	if (/^(?:c(?:check)?)?=?([0123]|t(?:est)?)$/) {
	    $check = $1;
	    $check = 2 # --check=test
		if lc($1) =~ /^t/;
	    shift @ARGV;
	}
	elsif (/^w?(?:arn)?$/) {
	    $checkwarn = 1;
	    shift @ARGV;
	}
	else {
	    usage( qq(ERROR: unknown option "$ARGV[0]") );
	}
    }

    # default: use -exist option to add (via ipset restore)
    $checkwarn = 0 unless defined $checkwarn;
    $check = 1 unless defined $check;
    $check = unpack("B*", pack("N", $check)) if $check;


    # and unless we have input
    die usage(
	qq(ERROR: STDIN is non a pipe or redirection)
    ) if -t STDIN;

    # read db, check remaining command-line DB file
    %m = db( @ARGV ); # use Data::Dumper; die Dumper( \%m );
}

if( /^create (\S+)/ )
{
    if( exists $h{$1} ) { # check if create issued
	$_ = ''; # don't print again
    } elsif( exists $m{ $1 } ) {
	$h{$1} = 1; # ensure this is the only printing
	my $n = $1; # grab the name
	my $v = $m{$n}; # lookup max
	# mangle the create to inject maxelem from DB
	s/^create $n (.*?maxelem) \d+ (.*)$/create $n $1 $v $2/
	    #and warn qq[set $n=$v]
    } else {
	$_ = ''; # skip create when no max defined
    }
}
elsif( /^add (\S+)/ )
{
    my $set = $1;
    $_ = '', next unless exists $m{$1};
    if ($check) {
      if ($check & ( 1<< $check_bits[1])) {
	my( $ip ) = /$set (\S+)/s;
	unless (system( qq(ipset test "$set" "$ip" >/dev/null 2>&1) )) {
	  warn qq[skip $set $ip (ipset test -eq 0)\n]
	      if $checkwarn;
	  $_ = '';
	  next;
	}
      }
      if ($check & ( 1<< $check_bits[0])) {
	s/$/ -exist/
      }
    }
}