Add options to limit extracted file sizes w/ 100MB default.

src/file_analysis/analyzer/extract/CMakeLists.txt

@@ -5,4 +5,6 @@ include_directories(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}
 
 bro_plugin_begin(Bro FileExtract)
 bro_plugin_cc(Extract.cc Plugin.cc ../../Analyzer.cc)
+bro_plugin_bif(events.bif)
+bro_plugin_bif(functions.bif)
 bro_plugin_end()

src/file_analysis/analyzer/extract/Extract.cc

@@ -4,13 +4,15 @@
 
 #include "Extract.h"
 #include "util.h"
+#include "Event.h"
 #include "file_analysis/Manager.h"
 
 using namespace file_analysis;
 
-Extract::Extract(RecordVal* args, File* file, const string& arg_filename)
+Extract::Extract(RecordVal* args, File* file, const string& arg_filename,
+                 uint64 arg_limit)
     : file_analysis::Analyzer(file_mgr->GetComponentTag("EXTRACT"), args, file),
-	  filename(arg_filename)
+      filename(arg_filename), limit(arg_limit)
 	{
 	fd = open(filename.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666);
 
@@ -29,15 +31,51 @@ Extract::~Extract()
 		safe_close(fd);
 	}
 
-file_analysis::Analyzer* Extract::Instantiate(RecordVal* args, File* file)
+static Val* get_extract_field_val(RecordVal* args, const char* name)
 	{
 	using BifType::Record::Files::AnalyzerArgs;
-	Val* v = args->Lookup(AnalyzerArgs->FieldOffset("extract_filename"));
+	Val* rval = args->Lookup(AnalyzerArgs->FieldOffset(name));
 
-	if ( ! v )
+	if ( ! rval )
+		reporter->Error("File extraction analyzer missing arg field: %s", name);
+
+	return rval;
+	}
+
+file_analysis::Analyzer* Extract::Instantiate(RecordVal* args, File* file)
+	{
+	Val* fname = get_extract_field_val(args, "extract_filename");
+	Val* limit = get_extract_field_val(args, "extract_limit");
+
+	if ( ! fname || ! limit )
 		return 0;
 
-	return new Extract(args, file, v->AsString()->CheckString());
+	return new Extract(args, file, fname->AsString()->CheckString(),
+	                   limit->AsCount());
+	}
+
+static bool check_limit_exceeded(uint64 lim, uint64 off, uint64 len, uint64* n)
+	{
+	if ( lim == 0 )
+		{
+		*n = len;
+		return false;
+		}
+
+	if ( off >= lim )
+		{
+		*n = 0;
+		return true;
+		}
+
+	*n = lim - off;
+
+	if ( len > *n )
+		return true;
+	else
+		*n = len;
+
+	return false;
 	}
 
 bool Extract::DeliverChunk(const u_char* data, uint64 len, uint64 offset)
@@ -45,6 +83,26 @@ bool Extract::DeliverChunk(const u_char* data, uint64 len, uint64 offset)
 	if ( ! fd )
 		return false;
 
-	safe_pwrite(fd, data, len, offset);
-	return true;
+	uint64 towrite = 0;
+	bool limit_exceeded = check_limit_exceeded(limit, offset, len, &towrite);
+
+	if ( limit_exceeded && file_extraction_limit )
+		{
+		File* f = GetFile();
+		val_list* vl = new val_list();
+		vl->append(f->GetVal()->Ref());
+		vl->append(Args()->Ref());
+		vl->append(new Val(limit, TYPE_COUNT));
+		vl->append(new Val(offset, TYPE_COUNT));
+		vl->append(new Val(len, TYPE_COUNT));
+		f->FileEvent(file_extraction_limit, vl);
+
+		// Limit may have been modified by BIF, re-check it.
+		limit_exceeded = check_limit_exceeded(limit, offset, len, &towrite);
+		}
+
+	if ( towrite > 0 )
+		safe_pwrite(fd, data, towrite, offset);
+
+	return ( ! limit_exceeded );
 	}

src/file_analysis/analyzer/extract/Extract.h

@@ -9,6 +9,8 @@
 #include "File.h"
 #include "Analyzer.h"
 
+#include "analyzer/extract/events.bif.h"
+
 namespace file_analysis {
 
 /**
@@ -41,6 +43,13 @@ public:
 	 */
 	static file_analysis::Analyzer* Instantiate(RecordVal* args, File* file);
 
+	/**
+	 * Sets the maximum allowed extracted file size.  A value of zero means
+	 * "no limit".
+	 * @param bytes number of bytes allowed to be extracted
+	 */
+	void SetLimit(uint64 bytes) { limit = bytes; }
+
 protected:
 
 	/**
@@ -49,12 +58,15 @@ protected:
 	 * @param file the file to which the analyzer will be attached.
 	 * @param arg_filename a file system path which specifies the local file
 	 *        to which the contents of the file will be extracted/written.
+	 * @param arg_limit the maximum allowed file size.
 	 */
-	Extract(RecordVal* args, File* file, const string& arg_filename);
+	Extract(RecordVal* args, File* file, const string& arg_filename,
+	        uint64 arg_limit);
 
 private:
 	string filename;
 	int fd;
+	uint64 limit;
 };
 
 } // namespace file_analysis

src/file_analysis/analyzer/extract/Plugin.cc

@@ -18,6 +18,11 @@ protected:
 
 		AddComponent(new ::file_analysis::Component("EXTRACT",
 		        ::file_analysis::Extract::Instantiate));
+
+		extern std::list<std::pair<const char*, int> > __bif_events_init();
+		AddBifInitFunction(&__bif_events_init);
+		extern std::list<std::pair<const char*, int> > __bif_functions_init();
+		AddBifInitFunction(&__bif_functions_init);
 		}
 };
 

src/file_analysis/analyzer/extract/events.bif

@@ -0,0 +1,19 @@
+## This event is generated when a file extraction analyzer is about
+## to exceed the maximum permitted file size allowed by
+## *extract_size_limit* field of :bro:see:`Files::AnalyzerArgs`.
+## The analyzer is automatically removed from file *f*.
+##
+## f: The file.
+##
+## args: Arguments that identify a particular file extraction analyzer.
+##       This is only provided to be able to pass along to
+##       :bro:see:`FileExtract::set_limit`.
+##
+## limit: The limit, in bytes, the extracted file is about to breach.
+##
+## offset: The offset at which a file chunk is about to be written.
+##
+## len:: The length of the file chunk about to be written.
+##
+## .. bro:see:: Files::add_analyzer Files::ANALYZER_EXTRACT
+event file_extraction_limit%(f: fa_file, args: any, limit: count, offset: count, len: count%);

src/file_analysis/analyzer/extract/functions.bif

@@ -0,0 +1,19 @@
+##! Internal functions used by the extraction file analyzer.
+
+module FileExtract;
+
+%%{
+#include "file_analysis/Manager.h"
+%%}
+
+## :bro:see:`FileExtract::set_limit`.
+function FileExtract::__set_limit%(file_id: string, args: any, n: count%): bool
+    %{
+	using BifType::Record::Files::AnalyzerArgs;
+	RecordVal* rv = args->AsRecordVal()->CoerceTo(AnalyzerArgs);
+    bool result = file_mgr->SetExtractionLimit(file_id->CheckString(), rv, n);
+	Unref(rv);
+    return new Val(result, TYPE_BOOL);
+    %}
+
+module GLOBAL;