diff --git a/plugins/ident/geoip b/plugins/ident/geoip index 53e81fe..9f09671 100644 --- a/plugins/ident/geoip +++ b/plugins/ident/geoip @@ -341,13 +341,8 @@ sub init_my_country_code { sub set_country_code { my $self = shift; - my $remote_ip = $self->qp->connection->remote_ip; - - my $code = $self->{_geoip_city} - ? $self->get_country_code_gc($remote_ip) - : $self->get_country_code($remote_ip); - - return if ! $code; + my $ip = $self->qp->connection->remote_ip; + my $code = $self->get_country_code($ip) or return; $self->qp->connection->notes('geoip_country', $code); return $code; } @@ -357,8 +352,11 @@ sub get_country_code { my $ip = shift || $self->qp->connection->remote_ip; if ($self->{_geoip_city}) { return $self->get_country_code_gc($ip); - }; - return $self->{_geoip}->country_code_by_addr($ip); + } + if ($self->{_geoip}) { + return $self->{_geoip}->country_code_by_addr($ip); + } + return undef; } sub get_country_code_gc { @@ -371,17 +369,24 @@ sub get_country_code_gc { sub set_country_name { my $self = shift; - my $remote_ip = $self->qp->connection->remote_ip; - - my $name = $self->{_geoip_city} - ? $self->get_country_name_gc($remote_ip) - : $self->{_geoip}->country_name_by_addr($remote_ip); - - return if ! $name; + my $ip = $self->qp->connection->remote_ip; + my $name = $self->get_country_name($ip) or return; $self->qp->connection->notes('geoip_country_name', $name); return $name; } +sub get_country_name { + my $self = shift; + my $ip = shift || $self->qp->connection->remote_ip; + if ($self->{_geoip_city}) { + return $self->get_country_name_gc($ip); + } + if ($self->{_geoip}) { + return $self->{_geoip}->country_name_by_addr($ip); + } + return undef; +} + sub get_country_name_gc { my $self = shift; return if !$self->{_geoip_record}; @@ -390,17 +395,24 @@ sub get_country_name_gc { sub set_continent { my ($self, $country_code) = @_; - $country_code or return; - - my $continent = $self->{_geoip_city} - ? $self->get_continent_gc() - : $self->{_geoip}->continent_code_by_country_code($country_code); - - $continent or return; + return if !$country_code; + my $continent = $self->get_continent($country_code) or return; $self->qp->connection->notes('geoip_continent', $continent); return $continent; } +sub get_continent { + my ($self, $country_code) = @_; + return if !$country_code; + if ($self->{_geoip_city}) { + return $self->get_continent_gc(); + } + if ($self->{_geoip}) { + return $self->{_geoip}->continent_code_by_country_code($country_code); + } + return undef; +} + sub get_continent_gc { my $self = shift; return if !$self->{_geoip_record}; diff --git a/t/plugin_tests/ident/geoip b/t/plugin_tests/ident/geoip index d28ed10..5436866 100644 --- a/t/plugin_tests/ident/geoip +++ b/t/plugin_tests/ident/geoip @@ -117,6 +117,10 @@ sub test_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"); @@ -134,6 +138,11 @@ sub test_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"); @@ -152,6 +161,11 @@ sub test_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'); @@ -200,6 +214,10 @@ sub test_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"); @@ -208,3 +226,15 @@ sub test_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; +}