qpsmtpd/t/plugin_tests/ident/geoip

#!perl -w

use strict;
use warnings;

use lib 'lib';
use Qpsmtpd::Constants;

sub register_tests {
    my $self = shift;

    eval 'use GeoIP2::Database::Reader';
    if ( !$@ ) {
        $self->register_test('test_geoip2_lookup');
    }

    eval 'use Geo::IP';
    if ( !$@ ) {
        $self->register_test('test_geoip_lookup');
        $self->register_test('test_geoip_load_db');
        $self->register_test('test_geoip_init_cc');
        $self->register_test('test_set_country_code');
        $self->register_test('test_set_country_name');
        $self->register_test('test_set_continent');
        $self->register_test('test_set_distance');
        $self->register_test('test_set_asn');
        $self->register_test('test_add_headers');
    }
}

sub test_geoip2_lookup {
    my $self = shift;

    $self->qp->connection->remote_ip('24.24.24.24');
    cmp_ok( $self->geoip2_lookup(), '==', DECLINED, "exit code DECLINED");

    if (!$self->load_geoip2()) {
        warn "failed to load GeoIP2\n";
    }

    cmp_ok( $self->connection->notes('geoip_country'), 'eq', 'US', "24.24.24.24 is in country US");
    cmp_ok( $self->connection->notes('geoip_country_name'), 'eq', 'United States', "24.24.24.24 is in country United States");
    cmp_ok( $self->connection->notes('geoip_continent'), 'eq', 'NA', "24.24.24.24 is in continent NA");
    cmp_ok( $self->connection->notes('geoip_city'), 'eq', 'Deer Park', "24.24.24.24 is in city of Deer Park");
}

sub test_add_headers {
    my ( $self ) = @_;
    my @notes = qw( geoip_country geoip_continent geoip_city geoip_asn );
    $self->connection->notes( $_ => "test $_" ) for @notes;
    my $header = $self->transaction->header( Mail::Header->new );
    my @tags = (qw( X-GeoIP-Country X-GeoIP-Continent X-GeoIP-City X-GeoIP-ASN ));
    $header->add( $_ => 'DELETETHIS' ) for @tags;
    $self->add_headers($self->transaction);
    is( $self->all_headers('X-GeoIP-Country'), 'test geoip_country',
        'X-GeoIP-Country header added' );
    is( $self->all_headers('X-GeoIP-Continent'), 'test geoip_continent',
        'X-GeoIP-Continent header added' );
    is( $self->all_headers('X-GeoIP-City'), 'test geoip_city',
        'X-GeoIP-City header added' );
    is( $self->all_headers('X-GeoIP-ASN'), 'test geoip_asn',
        'X-GeoIP-ASN header added' );
}

sub all_headers {
    # Return all instances of a given message header
    my ( $self, $tag ) = @_;
    return join " | ", map { chomp $_; $_ } $self->transaction->header->get($tag);
}

sub test_geoip_lookup {
    my $self = shift;

    $self->qp->connection->remote_ip('24.24.24.24');
    cmp_ok( $self->geoip_lookup(), '==', DECLINED, "exit code");

    cmp_ok( $self->connection->notes('geoip_country'), 'eq', 'US', "24.24.24.24 is in the US");
}

sub test_geoip_load_db {
    my $self = shift;

    $self->open_geoip_db();

    if ( $self->{_geoip_city} ) {
        ok( ref $self->{_geoip_city}, "loaded GeoIP city db" );
    }
    else {
        ok( "no GeoIP city db" );
    }

    if ( $self->{_geoip} ) {
        ok( ref $self->{_geoip}, "loaded GeoIP db" );
    }
    else {
        ok( "no GeoIP db" );
    }
}

sub test_geoip_init_cc {
    my $self = shift;

    $self->{_my_country_code} = undef;
    ok( ! $self->{_my_country_code}, "undefined");

    my $test_ip = '208.175.177.10';
    $self->{_args}{distance} = $test_ip;
    $self->init_my_country_code( $test_ip );
    cmp_ok( $self->{_my_country_code}, 'eq', 'US', "country set and matches");
}

sub test_set_country_code {
    my $self = shift;

    $self->qp->connection->remote_ip('');
    my $cc = $self->set_country_code();
    ok( ! $cc, "undef");

    $self->qp->connection->remote_ip('24.24.24.24');
    $self->clear_geoip_data;
    $cc = $self->set_country_code();
    ok( ! $cc, "set_country_code() returns nothing for no geoip data");
    $self->restore_geoip_data;
    $cc = $self->set_country_code();
    cmp_ok( $cc, 'eq', 'US', "set_country_code result is $cc");

    my $note = $self->connection->notes('geoip_country');
    cmp_ok( $note, 'eq', 'US', "set_country_code set note to $cc");
}

sub test_set_country_name {
    my $self = shift;

    $self->{_geoip_record} = undef;
    $self->qp->connection->remote_ip('');
    $self->set_country_code();
    my $cn = $self->set_country_name();
    ok( ! $cn, "undef") or warn "$cn\n";

    $self->qp->connection->remote_ip('24.24.24.24');
    $self->clear_geoip_data;
    $self->set_country_code();
    $cn = $self->set_country_name();
    ok( ! $cn, "set_country_name() returns nothing for no geoip data");
    $self->restore_geoip_data;
    $self->set_country_code();
    $cn = $self->set_country_name();
    cmp_ok( $cn, 'eq', 'United States', "$cn");

    my $note = $self->connection->notes('geoip_country_name');
    cmp_ok( $note, 'eq', 'United States', "note has: $cn");
}

sub test_set_continent {
    my $self = shift;

    $self->{_geoip_record} = undef;
    $self->qp->connection->remote_ip('');
    $self->set_country_code();
    my $cn = $self->set_continent();
    ok( ! $cn, "undef") or warn "$cn\n";

    $self->qp->connection->remote_ip('24.24.24.24');
    $self->clear_geoip_data;
    $self->set_country_code();
    $cn = $self->set_continent('US');
    ok( ! $cn, 'set_continent() returns nothing for no geoip data');
    $self->restore_geoip_data;
    $self->set_country_code();
    $cn = $self->set_continent() || '';
    my $note = $self->connection->notes('geoip_continent');
    if ( $cn ) {
        cmp_ok( $cn, 'eq', 'NA', "$cn");
        cmp_ok( $note, 'eq', 'NA', "note has: $cn");
    }
    else {
        ok(1, "no continent data" );
        ok(1, "no continent data" );
    }
}

sub test_set_distance {
    my $self = shift;

    $self->{_geoip_record} = undef;
    $self->qp->connection->remote_ip('');
    $self->set_country_code();
    my $cn = $self->set_distance_gc();
    ok( ! $cn, "undef") or warn "$cn\n";

    $self->qp->connection->remote_ip('24.24.24.24');
    $self->set_country_code();
    $cn = $self->set_distance_gc();
    if ( $cn ) {
        ok( $cn, "$cn km");

        my $note = $self->connection->notes('geoip_distance');
        ok( $note, "note has: $cn");
    }
    else {
        ok( 1, "no distance data");
        ok( 1, "no distance data");
    }
}

sub test_set_asn {
    my $self = shift;

    return if !$self->{GeoIPASNum};

    $self->qp->connection->remote_ip('');
    $self->set_asn();
    my $asn = $self->set_asn();
    ok( ! $asn, "undef") or warn "$asn\n";

    $self->qp->connection->remote_ip('24.24.24.24');
    $self->clear_geoip_data;
    $asn = $self->set_asn();
    ok( ! $asn, 'set_asn() returns nothing for no ASN data' );
    $self->restore_geoip_data;
    $asn = $self->set_asn();
    ok( $self->connection->notes('geoip_asn') =~ /^11351/, "note has: $asn");

    $self->qp->connection->remote_ip('66.128.51.163');
    $asn = $self->set_asn();

    ok( $self->connection->notes('geoip_asn') =~ /^7819/, "note has: $asn");
}

my $geoip_data_bak;
my @geoip_keys = qw( _geoip _geoip_city GeoIPASNum );
sub clear_geoip_data {
    my ( $self ) = @_;
    $geoip_data_bak->{$_} = delete $self->{$_} for @geoip_keys;
}

sub restore_geoip_data {
    my ( $self ) = @_;
    $self->{$_} = delete $geoip_data_bak->{$_} for @geoip_keys;
}