restore

plugins/ident/geoip

@@ -175,23 +175,38 @@ sub load_geoip1 {
 sub load_geoip2 {
     my $self = shift;
 
-    eval 'use Geo::IP';
+    eval 'use GeoIP2::Database::Reader';
     if ($@) {
-        warn "could not load Geo::IP";
-        $self->log(LOGERROR, "could not load Geo::IP");
+        $self->log(LOGERROR, "could not load GeoIP2");
         return;
     }
 
-    $self->open_geoip_db();
+    warn "Using GeoIP2."
+       . " ASN data is not currently available using the GeoIP2 module!\n";
 
-    # Note that opening the GeoIP DB only in register has caused problems before:
-    # https://github.com/smtpd/qpsmtpd/commit/29ea9516806e9a8ca6519fcf987dbd684793ebdd#plugins/ident/geoip
-    # Opening the DB anew for every connection is horribly inefficient.
-    # Instead, attempt to reopen upon connect if the DB connection fails.
-    $self->init_my_country_code();
+   eval {
+        $self->{_geoip2_city} = GeoIP2::Database::Reader->new(
+            file => $self->{_args}{db_dir} . '/GeoLite2-City.mmdb',
+        );
+    };
+    if ($@) {
+        $self->log(LOGERROR, "unable to load GeoLite2-City.mmdb");
+    }
 
-    $self->register_hook('connect', 'geoip_lookup');
-    return 1;
+    eval {
+        $self->{_geoip2_country} = GeoIP2::Database::Reader->new(
+            file => $self->{_args}{db_dir} . '/GeoLite2-Country.mmdb',
+        );
+    };
+    if ($@) {
+        $self->log(LOGERROR, "unable to load GeoLite2-Country.mmdb");
+    }
+
+    if ($self->{_geoip2_city} || $self->{_geoip2_country}) {
+        $self->register_hook('connect', 'geoip2_lookup');
+        return 1;
+    }
+    return;
 }
 
 sub add_headers {
@@ -205,6 +220,36 @@ sub add_headers {
     return DECLINED;
 }
 
+sub geoip2_lookup {
+    my $self = shift;
+
+    my $ip = $self->qp->connection->remote_ip;
+    return DECLINED if $self->is_localhost($ip);
+
+    if ($self->{_geoip2_city}) {
+        my $city_rec = $self->{_geoip2_city}->city(ip => $ip);
+        if ($city_rec) {
+            $self->qp->connection->notes('geoip_country', $city_rec->country->iso_code());
+            $self->qp->connection->notes('geoip_country_name', $city_rec->country->name());
+            $self->qp->connection->notes('geoip_continent', $city_rec->continent->code());
+            $self->qp->connection->notes('geoip_city', $city_rec->city->name());
+            $self->qp->connection->notes('geoip_asn', $city_rec->traits->autonomous_system_number());
+            return DECLINED;
+        }
+    }
+
+    if ($self->{_geoip2_country}) {
+        my $country_rec = $self->{_geoip2_country}->country(ip => $ip);
+        if ($country_rec) {
+            $self->qp->connection->notes('geoip_country', $country_rec->country->iso_code());
+            $self->qp->connection->notes('geoip_country_name', $country_rec->country->name());
+            $self->qp->connection->notes('geoip_continent', $country_rec->continent->code());
+        };
+    }
+
+    return DECLINED;
+}
+
 sub geoip_lookup {
     my $self = shift;