From 240ae9790bbb7e1ff5f07e5f428fbe8e8ffc67c8 Mon Sep 17 00:00:00 2001
From: Seth Hall <seth@icir.org>
Date: Thu, 11 Aug 2011 14:59:01 -0400
Subject: [PATCH] Small updates for notice framework.

- New ACTION_ADD_GEODATA to add geodata to notices in an extension
  field named remote_location.

- Loading extend-email/hostnames by default now that it only
  does anything when the ACTION_EMAIL action is applied (finally).
---
 scripts/base/frameworks/notice/__load__.bro   |  7 +--
 .../frameworks/notice/actions/add-geodata.bro | 47 +++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 scripts/base/frameworks/notice/actions/add-geodata.bro

diff --git a/scripts/base/frameworks/notice/__load__.bro b/scripts/base/frameworks/notice/__load__.bro
index bbc1fcae0d..2cc93ee933 100644
--- a/scripts/base/frameworks/notice/__load__.bro
+++ b/scripts/base/frameworks/notice/__load__.bro
@@ -6,7 +6,8 @@
 @load ./actions/drop
 @load ./actions/email_admin
 @load ./actions/page
+@load ./actions/add-geodata
 
-# Load the script to add hostnames to emails by default.
-# NOTE: this exposes a memleak in async DNS lookups.
-#@load ./extend-email/hostnames
+# There shouldn't be any defaul toverhead from loading these since they 
+# *should* only do anything when notices have the ACTION_EMAIL action applied.
+@load ./extend-email/hostnames
diff --git a/scripts/base/frameworks/notice/actions/add-geodata.bro b/scripts/base/frameworks/notice/actions/add-geodata.bro
new file mode 100644
index 0000000000..71e9c6b490
--- /dev/null
+++ b/scripts/base/frameworks/notice/actions/add-geodata.bro
@@ -0,0 +1,47 @@
+##! This script adds geographic location data to notices for the "remote"
+##! host in a connection.  It does make the assumption that one of the 
+##! addresses in a connection is "local" and one is "remote" which is 
+##! probably a safe assumption to make in most cases.  If both addresses
+##! are remote, it will use the $src address.
+
+module Notice;
+
+export {
+	redef enum Action += {
+		## Indicates that the notice should have geodata added for the
+		## "remote" host.  :bro:id:`Site::local_nets` must be defined
+		## in order for this to work.
+		ACTION_ADD_GEODATA
+	};
+	
+	redef record Info += {
+		## If libGeoIP support is built in, notices can have geographic
+		## information attached to them.
+		remote_location: geo_location  &log &optional;
+	};
+	
+	## Notice types which should have the "remote" location looked up.
+	## If GeoIP support is not built in, this does nothing.
+	const lookup_location_types: set[Notice::Type] = {} &redef;
+	
+	## Add a helper to the notice policy for looking up GeoIP data.
+	redef Notice::policy += {
+		[$pred(n: Notice::Info) = { return (n$note in Notice::lookup_location_types); },
+		 $priority = 10],
+	};
+}
+
+# This is handled at a high priority in case other notice handlers 
+# want to use the data.
+event notice(n: Notice::Info) &priority=10
+	{
+	if ( ACTION_ADD_GEODATA in n$actions &&
+	     |Site::local_nets| > 0 &&
+	     ! n?$remote_location )
+		{
+		if ( n?$src && ! Site::is_local_addr(n$src) )
+			n$remote_location = lookup_location(n$src);
+		else if ( n?$dst && ! Site::is_local_addr(n$dst) )
+			n$remote_location = lookup_location(n$dst);
+		}
+	}
\ No newline at end of file