Merge pull request #152 from jaredj/geoip-headers

Add message headers from GeoIP data

config.sample/add_geoip_headers

@@ -0,0 +1 @@
+1

plugins/ident/geoip

@@ -73,6 +73,14 @@ The path to the GeoIP database directory.
 
 Default: /usr/local/share/GeoIP
 
+=head2 add_headers <true|false>
+
+Add message headers with GeoIP data
+
+  ident/geoip [ add_headers (true|false) ]
+
+Default: true
+
 =head1 LIMITATIONS
 
 The distance calculations are more concerned with being fast than accurate.
@@ -129,8 +137,18 @@ sub register {
     $self->{_args} = {@args};
     $self->{_args}{db_dir} ||= '/usr/local/share/GeoIP';
 
-    $self->load_geoip2() and return;
+    $self->load_geoip() or return;
-    $self->load_geoip1();
+    my $enabled = $self->{_args}{add_headers};
+    $enabled = 'true' if ! defined $enabled;
+    return if $enabled =~ /false/i;
+    $self->register_hook( data_post => 'add_headers' );
+}
+
+sub load_geoip {
+    my ( $self ) = @_;
+    $self->load_geoip2() and return 1;
+    $self->load_geoip1() and return 1;
+    return 0;
 }
 
 sub load_geoip1 {
@@ -152,6 +170,7 @@ sub load_geoip1 {
     $self->init_my_country_code();
 
     $self->register_hook('connect', 'geoip_lookup');
+    return 1;
 }
 
 sub load_geoip2 {
@@ -191,6 +210,16 @@ sub load_geoip2 {
     return;
 }
 
+sub add_headers {
+    my ( $self, $txn ) = @_;
+    for my $h (qw( Country Continent City ASN )) {
+        my $note = lc "geoip_$h";
+        next if ! $self->connection->notes($note);
+        $txn->header->delete("X-GeoIP-$h");
+        $txn->header->add( "X-GeoIP-$h", $self->connection->notes($note), 0 );
+    }
+}
+
 sub geoip2_lookup {
     my $self = shift;
 

t/config/add_geoip_headers

@@ -0,0 +1 @@
+1

t/plugin_tests/ident/geoip

@@ -27,6 +27,7 @@ sub register_tests {
         $self->register_test('test_set_continent');
         $self->register_test('test_set_distance');
         $self->register_test('test_set_asn');
+        $self->register_test('test_add_headers');
     }
 }
 
@@ -46,6 +47,30 @@ sub test_geoip2_lookup {
     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;