#!/usr/bin/perl
use strict;
use warnings;
use feature 'state';

=head1 NAME

portblocker-enforce

=head1 SYNOPSIS

    ## Важно: обычно этот скрипт вручную запускать не надо, он сам запускается из cron

    ## сделать правила iptables как надо
    $ portblocker-enforce --set

    ## получить данные про нужные кондукторные группы и сохранить в кэш
    $ portblocker-enforce --fetch

    ## и то и другое: получить группы из кондуктора, сделать правила как надо
    $ portblocker-enforce --fetch --set

    ## вывести команды iptables/ip6tables, которые сделают правила как надо
    $ portblocker-enforce --set --simulate

    ## вывести, какие кондукторные группы скрипту интересны
    $ portblocker-enforce --fetch --simulate

    ## так делать нельзя: пока скрипт не получил кондукторные группы, он не
    ## может сказать, как приводить правила в нужный вид:
    $ portblocker-enforce --fetch --set --simulate
    Running portblocker with --fetch --set --simulate all at the same time is unsupported;
    <...>

=head1 DESCRIPTION

Внутренний скрипт, обычно вручную запускать не надо.

Часть системы portblocker. Что такое portblocker, написано в portblocker --help

Делает две вещи:
1. получает содержимое кондукторных групп (складывает в /var/spool/portblocker/hosts.cache)
2. делает правила iptables/ip6tables как надо

При обычных запусках (без --simulate) записывает действия в syslog. Как примерно посмотреть:

    $ grep portblocker /var/log/syslog
    Dec 28 15:50:01 $hostname CRON[11423]: (root) CMD (/usr/local/bin/portblocker.pl --fetch)
    Dec 28 15:50:01 $hostname portblocker[11425]: start action_fetch
    Dec 28 15:50:01 $hostname portblocker[11425]: finish action_fetch
    Dec 28 15:51:01 $hostname CRON[11452]: (root) CMD (/usr/local/bin/portblocker.pl --set)
    Dec 28 15:51:01 $hostname portblocker[11454]: start action_set
    Dec 28 15:51:01 $hostname portblocker[11454]: running: /sbin/iptables -N portblocker
    Dec 28 15:51:01 $hostname portblocker[11454]: running: /sbin/iptables -A portblocker -p tcp -m tcp --dport 3314 -j REJECT --reject-with tcp-reset
    <...>
    Dec 28 15:51:02 $hostname portblocker[11454]: running: /sbin/ip6tables -I INPUT -p tcp -m tcp --dport 17601 -j portblocker
    Dec 28 15:51:02 $hostname portblocker[11454]: finish action_set

Скрипт редактирует правила iptables и ip6tables. Пакеты к "интересным" портам заворачиваются в цепочку portblocker.
Там проверяется хост отправителя: если он есть в белом списке -- ACCEPT, если нет -- REJECT.

Скрипт обрабатывает только инстансы в состоянии "ограничен". Инстансов в состоянии "открыт" для него как будто бы нет.

Что скрипт делает, более подробно:
1. создаётся цепочка portblocker, если её нет
2. пакеты в сторону портов инстансов направляются на обработку в этой цепочке
3. для каждого порта инстанса и каждого хоста, с которого должен быть доступ к этому порту, в цепочке portblocker
   делается правило ACCEPT
4. для каждого порта инстанса в конце цепочки portblocker делается правило REJECT --reject-with tcp-reset
5. если есть правила, которые направляют пакеты в цепочку portblocker, но не соответствуют ни одному инстансу,
   эти правила удаляются
6. если в цепочке portblocker находятся правила, которых там не должно быть (в пп. 3 и 4 их не делали),
   эти правила удаляются

Когда скрипт portblocker "открывает" инстанс, скрипт portblocker-enforce удалит правила для этого инстанса
(потому что не узнает их в пп. 5 и 6).

Почему именно REJECT и именно --reject-with tcp-reset:
1. DROP -- отправители пакетов долго ждали бы таймаута
2. REJECT --reject-with icmp-port-unreachable (--reject-with по умолчанию) -- не закрывал бы старые соединения

Как смотреть правила iptables и ip6tables:

    $ sudo iptables -vnL
    Chain INPUT (policy ACCEPT 119 packets, 8930 bytes)
     pkts bytes target     prot opt in     out     source               destination
        0     0 portblocker  tcp  --  *      *       0.0.0.0/0            0.0.0.0/0            tcp dpt:17601

    <здесь цепочки OUTPUT и FORWARD, их portblocker не трогает>

    Chain portblocker (4 references)
     pkts bytes target     prot opt in     out     source               destination
        0     0 ACCEPT     tcp  --  *      *       37.140.140.240       0.0.0.0/0            tcp dpt:17601
        <...>
        0     0 REJECT     tcp  --  *      *       0.0.0.0/0            0.0.0.0/0            tcp dpt:17601 reject-with tcp-reset

Если кажется, что portblocker "делает что-то не то" и хочется его быстро отключить:
1. убрать запуск portblocker-enforce из crontab-файла /etc/cron.d/yandex-du-portblocker
2. удалить правила iptables и ip6tables
    echo '#!/bin/bash -ex' > iptables-cleanup.sh
    iptables -S INPUT | grep portblocker | sed 's/^-A/-D/' | xargs -d "\n" -n1 echo iptables >> iptables-cleanup.sh
    ip6tables -S INPUT | grep portblocker | sed 's/^-A/-D/' | xargs -d "\n" -n1 echo ip6tables >> iptables-cleanup.sh
    echo 'iptables -F portblocker' >> iptables-cleanup.sh
    echo 'iptables -X portblocker' >> iptables-cleanup.sh
    echo 'ip6tables -F portblocker' >> iptables-cleanup.sh
    echo 'ip6tables -X portblocker' >> iptables-cleanup.sh
    chmod a+x iptables-cleanup.sh
    ./iptables-cleanup.sh

=cut

use File::Basename;
use File::Slurp;
use Getopt::Long;
use List::MoreUtils qw( any );
use Net::DNS;
use Net::IP;
use Pid::File::Flock ':auto';
use Sys::Syslog qw( :standard :macros );
use YAML 'LoadFile';

use Yandex::Conductor;
use Yandex::HashUtils;
use Yandex::Shell;

use PortBlocker;

# $Yandex::Shell::PRINT_COMMANDS = 1;

=head1 CONFIGURATION

=head2 $HOST_CACHE_PATH

Где хранится кэш с информацией о хостах в группах кондуктора.

=cut

my $HOST_CACHE_PATH = '/var/spool/portblocker/hosts.cache';

=head2 $IPTABLES_CHAIN_NAME

Цепочка iptables, правилами в которой управляет этот скрипт. В конфигурации INPUT нужные порты
перенаправляются на эту цепочку.

=cut

my $IPTABLES_CHAIN_NAME = 'PORTBLOCKER';

=head2 $IPTABLES_BLOCK_ACTION

Каким образом блокировать пакеты, которые нужно заблокировать.

=cut

my $IPTABLES_BLOCK_ACTION = '-j REJECT --reject-with tcp-reset';

my ( $opt_fetch, $opt_set, $opt_simulate );

Getopt::Long::GetOptions(
    'fetch'      => \$opt_fetch,
    'set'        => \$opt_set,
    'simulate'   => \$opt_simulate,
    'help|usage' => \&usage,
)
    or die "Invalid options";

openlog 'portblocker', 'ndelay,pid', LOG_USER;

action_fetch() if $opt_fetch;
action_set() if $opt_set;

exit 0;

=head1 SUBROUTINES/METHODS

=head2 get_host_addresses

=cut

sub get_host_addresses {
    my ($hostname) = @_;

    state $resolver = Net::DNS::Resolver->new;
    state $cache = {};

    return $cache->{$hostname} if $cache->{$hostname};

    $cache->{$hostname} = { ipv4 => [], ipv6 => [] };
    for my $proto ( qw( ipv4 ipv6 ) ) {
        my $queryres = $resolver->query( $hostname, $proto eq 'ipv4' ? 'A' : 'AAAA' );
        next unless $queryres;

        my @answers = $queryres->answer;
        $cache->{$hostname}->{$proto} = [ map { $_->address } @answers ];
    }

    return $cache->{$hostname};
}

=head2 run_iptables_change_cmd

=cut

sub run_iptables_change_cmd {
    my ($cmd) = @_;

    if ($opt_simulate) {
        print "Run: $cmd\n";
    } else {
        syslog LOG_NOTICE, "running: $cmd";
        yash_system($cmd);
    }
}

=head2 normalize_ip_mask

=cut

sub normalize_ip_mask {
    my ($mask) = @_;
    my $ip = Net::IP->new($mask);
    return $ip->ip . '/' . $ip->prefixlen;
}

=head2 chain_exists

=cut

sub chain_exists {
    my ($iptables_cmd) = @_;

    my @iptables_dump = split /\n/, yash_qx("$iptables_cmd -S");
    return any { $_ eq "-N $IPTABLES_CHAIN_NAME" } @iptables_dump;
}

=head2 ensure_chain_created

Если "нашей" цепочки нет, пытается её создать с помощью run_iptables_change_cmd,
которая на самом деле запускает iptables -N или выводит команду на экран в режиме
--simulate.

=cut

sub ensure_chain_exists {
    my ($iptables_cmd) = @_;

    # эта переменная нужна для того, чтобы в режиме simulate не делать вид,
    # что мы хотим создать цепочку дважды
    state $chain_created = {};

    return if chain_exists($iptables_cmd) || $chain_created->{$iptables_cmd};
    run_iptables_change_cmd("$iptables_cmd -N $IPTABLES_CHAIN_NAME");
    $chain_created->{$iptables_cmd}++;
}

=head2 get_chain_dump

Возвращает содержимое "нашей" цепочки ($IPTABLES_CHAIN_NAME), если она есть.
Если её нет, делает вид, что она есть пустая, и возвращает одну команду создания.

=cut

sub get_chain_dump {
    my ($iptables_cmd) = @_;

    unless ( chain_exists($iptables_cmd) ) {
        return "-N $IPTABLES_CHAIN_NAME";
    }

    return yash_qx("$iptables_cmd -S $IPTABLES_CHAIN_NAME");
}

=head2 fix_rules

=cut

sub fix_rules {
    my ( $iptables_cmd, $port, $masks ) = @_;

    # в @desired_rules есть только правила, относящиеся к конкретным хостам;
    # -j REJECT/DROP без -s, которое должно быть последним, обрабатывается отдельно
    my @desired_rules = map { "-s $_ -p tcp -m tcp --dport $port -j ACCEPT" } map { normalize_ip_mask($_) } @$masks;

    my $chain_dump = get_chain_dump($iptables_cmd);

    my @existing_rules;
    for my $iptables_out_line ( split /\n/, $chain_dump ) {
        next unless $iptables_out_line =~ /--dport $port/;

        $iptables_out_line =~ s/-s (\S+)/ '-s ' . normalize_ip_mask($1) /e;

        if ( $iptables_out_line =~ /^-A $IPTABLES_CHAIN_NAME (.*)$/ ) {
            push @existing_rules, $1;
        }
    }
    my %is_rule_desired = map { $_ => 1 } @desired_rules;

    ensure_chain_exists($iptables_cmd);

    # есть ли в конце списка правил -j REJECT?
    my $last_rule = $existing_rules[-1];
    if ( $last_rule && $last_rule eq "-p tcp -m tcp --dport $port $IPTABLES_BLOCK_ACTION" ) {
        # если есть, нужно убрать его из списка @existing_rules, чтобы при сравнении дальше оно не удалилось
        pop @existing_rules;
    } else {
        # если нет, его нужно добавить
        run_iptables_change_cmd(
            "$iptables_cmd -A $IPTABLES_CHAIN_NAME -p tcp -m tcp --dport $port $IPTABLES_BLOCK_ACTION" );
    }

    my %is_rule_existing = map { $_ => 1 } @existing_rules;

    # а теперь нужно сравнить @existing_rules и @desired_rules и попросить iptables привести всё
    # к виду в @desired_rules
    for my $rule_to_remove ( grep { !$is_rule_desired{$_} } @existing_rules ) {
        run_iptables_change_cmd("$iptables_cmd -D $IPTABLES_CHAIN_NAME $rule_to_remove");
    }

    for my $rule_to_add ( grep { !$is_rule_existing{$_} } @desired_rules ) {
        run_iptables_change_cmd("$iptables_cmd -I $IPTABLES_CHAIN_NAME $rule_to_add");
    }

    return;
}

=head2 fix_input_redirects

=cut

sub fix_input_redirects {
    my ( $iptables_cmd, $ports ) = @_;

    my %want_port = map { $_ => 1 } @$ports;

    my @input_dump = split /\n/, yash_qx("$iptables_cmd -S INPUT");

    my %port_redirected;

    for my $line (@input_dump) {
        my $rule;
        if ( $line =~ /^-A INPUT (.*)/ ) {
            $rule = $1;
        }

        next unless $rule;

        my $rule_has_redirect = ( $rule =~ /-j $IPTABLES_CHAIN_NAME/ ? 1 : 0 );

        my $port = '';
        if ( $rule =~ /--dport (\d+)/ ) {
            $port = $1;
        }

        my $correct_redirect_rule = $rule eq "-p tcp -m tcp --dport $port -j $IPTABLES_CHAIN_NAME" ? 1 : 0;

        my $delete_cmd = "$iptables_cmd -D INPUT $rule";

        if ( $rule_has_redirect && $want_port{$port} && $correct_redirect_rule ) {
            # это правило, которое мы хотим, пометим, что мы его видели
            $port_redirected{$port} = 1;
            next;
        }

        if ( $want_port{$port} && !$correct_redirect_rule ) {
            # это неправильное правило для интересного нам порта, его нужно удалить
            run_iptables_change_cmd($delete_cmd);
            next;
        }

        if ( $rule_has_redirect && !$want_port{$port} ) {
            # это правило перенаправляет в нашу цепочку какие-то пакеты с порта, который нам неинтересен,
            # его нужно удалить
            run_iptables_change_cmd($delete_cmd);
            next;
        }

        if ( !$rule_has_redirect && $want_port{$port} ) {
            # это правило что-то делает с портом, который нам интересен, но при этом это не перенаправление
            # в нашу цепочку, так что его нужно удалить
            run_iptables_change_cmd($delete_cmd);
            next;
        }

        # оставшийся случай никак обрабатывать не нужно:
        # !$rule_has_redirect && !$want_port{$port}
    }

    # а ещё нужно перенаправить все порты, для которых мы не нашли правильных перенаправлений
    for my $port ( grep { !$port_redirected{$_} } @$ports ) {
        run_iptables_change_cmd("$iptables_cmd -I INPUT -p tcp -m tcp --dport $port -j $IPTABLES_CHAIN_NAME");
    }
}

=head2 cleanup_chain

=cut

sub cleanup_chain {
    my ( $iptables_cmd, $ports ) = @_;

    my %want_port = map { $_ => 1 } @$ports;

    my @chain_dump = split /\n/, get_chain_dump($iptables_cmd);

    my %port_redirected;

    ensure_chain_exists($iptables_cmd);

    for my $line (@chain_dump) {
        my $rule;
        if ( $line =~ /^-A $IPTABLES_CHAIN_NAME (.*)/ ) {
            $rule = $1;
        }

        next unless $rule;

        my $port = '';
        if ( $rule =~ /--dport (\d+)/ ) {
            $port = $1;
        }

        if ( !$port || !$want_port{$port} ) {
            run_iptables_change_cmd("$iptables_cmd -D $IPTABLES_CHAIN_NAME $rule");
        }
    }
}

=head2 action_fetch

=cut

sub action_fetch {
    my $config = get_configuration();

    syslog LOG_NOTICE, 'start action_fetch';

    my $cache = {};

    for my $rule (@$config) {
        for my $allow_item ( @{ $rule->{allow} } ) {
            next unless $allow_item->{type} eq 'conductor-group';

            my $group = $allow_item->{group};
            next if $cache->{$group};

            $cache->{$group} = { ipv4 => [], ipv6 => [] };

            if ($opt_simulate) {
                print "Would fetch '$group' from conductor\n";
                next;
            }

            my $hosts = conductor_groups2hosts($group);
            for my $host (@$hosts) {
                my $addresses = get_host_addresses($host);
                push @{ $cache->{$group}->{ipv4} }, @{ $addresses->{ipv4} };
                push @{ $cache->{$group}->{ipv6} }, @{ $addresses->{ipv6} };
            }
        }
    }

    if ( $opt_simulate && ! ( keys %$cache ) ) {
        print "(no conductor groups to fetch)\n";
    }

    unless ($opt_simulate) {
        write_file( $HOST_CACHE_PATH, { atomic => 1 }, YAML::Dump( $cache ) );
    }

    syslog LOG_NOTICE, 'finish action_fetch';

    return;
}

=head2 action_set

=cut

sub action_set {
    if ( $opt_fetch && $opt_simulate ) {
        die "Running portblocker with --fetch --set --simulate all at the same time is unsupported;\n" .
            "--fetch --simulate doesn't actually fetch the groups from conductor, so it isn't possible\n" .
            "to know what --set would do after.\n\n" .
            "Instead, you may back up the hosts cache file, then run portblocker-enforce --fetch and\n" .
            "portblocker-enforce --set --simulate, then restore the hosts cache file.\n\n" .
            "Host cache file is: $HOST_CACHE_PATH\n";
    }

    unless ( -f $HOST_CACHE_PATH ) {
        die "$HOST_CACHE_PATH not found, run $0 --fetch first before running $0 --set\n";
    }

    syslog LOG_NOTICE, 'start action_set' unless $opt_simulate;

    my $config = get_configuration();
    my $host_cache = YAML::LoadFile($HOST_CACHE_PATH);

    my $conffile_enabled = get_enabled_conffiles_hash();

    my @all_ports;

    for my $rule (@$config) {
        next unless $conffile_enabled->{ $rule->{conffile} };

        my $ports = $rule->{ports};
        my $allow = $rule->{allow};

        my ( @ipv4_hosts, @ipv6_hosts, @ipv4_masks, @ipv6_masks );
        for my $allow_item ( @{ $rule->{allow} || [] } ) {
            my $type = $allow_item->{type};

            if ( $type eq 'host' ) {
                push @ipv4_hosts, $allow_item->{'ip-address'};
                push @ipv6_hosts, $allow_item->{'ip6-address'};
                next;
            }

            if ( $type eq 'host-mask' ) {
                push @ipv4_masks, $allow_item->{'ip-mask'}  if $allow_item->{'ip-mask'};
                push @ipv6_masks, $allow_item->{'ip6-mask'} if $allow_item->{'ip6-mask'};
                next;
            }

            if ( $type eq 'conductor-group' ) {
                my $addresses = $host_cache->{ $allow_item->{group} };

                unless ($addresses) {
                    die "group $allow_item->{group} isn't present in the cache of conductor groups,\n" .
                        "try running --fetch one more time?\n";
                }

                die "group $allow_item->{group} seems to be empty"
                    unless @{ $addresses->{ipv4} } || @{ $addresses->{ipv6} };

                push @ipv4_hosts, @{ $addresses->{ipv4} };
                push @ipv6_hosts, @{ $addresses->{ipv6} };
                next;
            }
        }

        while ( my $host = shift @ipv4_hosts ) {
            next unless $host;
            push @ipv4_masks, "$host/32";
        }

        while ( my $host = shift @ipv6_hosts ) {
            next unless $host;
            push @ipv6_masks, "$host/128";
        }

        for my $port (@$ports) {
            fix_rules( '/sbin/iptables',  $port, \@ipv4_masks );
            fix_rules( '/sbin/ip6tables', $port, \@ipv6_masks );
        }

        push @all_ports, @$ports;
    }

    fix_input_redirects( '/sbin/iptables',  \@all_ports );
    fix_input_redirects( '/sbin/ip6tables', \@all_ports );

    cleanup_chain( '/sbin/iptables',  \@all_ports );
    cleanup_chain( '/sbin/ip6tables', \@all_ports );

    syslog LOG_NOTICE, 'finish action_set' unless $opt_simulate;

    return;
}

=head2 usage

=cut

sub usage {
    my $base_cmd = "podselect -section NAME -section SYNOPSIS -section DESCRIPTION $0 | pod2text-utf8";

    if ( my $pager = $ENV{PAGER} ) {
        system("$base_cmd | $pager");
    } else {
        system($base_cmd);
    }

    exit(0);
}
