packet_io: Refactor the Trough a bit

diff --git a/src/main/snort_module.cc b/src/main/snort_module.cc
index dc8462a6f4aacf1441755e7f1f84726e750dfdbe..95a04d16b6b066186f1e693ce40445aef11904f2 100644 (file)
--- a/src/main/snort_module.cc
+++ b/src/main/snort_module.cc
@@ -706,8 +706,8 @@ bool SnortModule::set(const char*, Value& v, SnortConfig* sc)
     }
     else if ( v.is("-r") || v.is("--pcap-list") )
     {
-        Trough::add_source(Trough::SOURCE_LIST, v.get_string());
         sc->run_flags |= RUN_FLAG__READ;
+        Trough::add_source(Trough::SOURCE_LIST, v.get_string());
     }
     else if ( v.is("-S") )
         config_set_var(sc, v.get_string());
@@ -889,13 +889,13 @@ bool SnortModule::set(const char*, Value& v, SnortConfig* sc)
 
     else if ( v.is("--pcap-file") )
     {
-        Trough::add_source(Trough::SOURCE_FILE_LIST, v.get_string());
         sc->run_flags |= RUN_FLAG__READ;
+        Trough::add_source(Trough::SOURCE_FILE_LIST, v.get_string());
     }
     else if ( v.is("--pcap-dir") )
     {
-        Trough::add_source(Trough::SOURCE_DIR, v.get_string());
         sc->run_flags |= RUN_FLAG__READ;
+        Trough::add_source(Trough::SOURCE_DIR, v.get_string());
     }
     else if ( v.is("--pcap-filter") )
         Trough::set_filter(v.get_string());

diff --git a/src/packet_io/trough.cc b/src/packet_io/trough.cc
index 81c17272f43516c584a7d0cb481ef7191dbfe210..9bfcdd2ca06b60f5959d6d9a31fdcc385db0627f 100644 (file)
--- a/src/packet_io/trough.cc
+++ b/src/packet_io/trough.cc
@@ -23,7 +23,6 @@
 
 #include "trough.h"
 
-#include <dirent.h>
 #include <fnmatch.h>
 #include <sys/stat.h>
 
@@ -45,158 +44,154 @@ std::vector<std::string>::const_iterator Trough::pcap_queue_iter;
 unsigned Trough::pcap_loop_count = 0;
 unsigned Trough::file_count = 0;
 
-int Trough::get_pcaps(std::vector<struct PcapReadObject> &pol)
+bool Trough::add_pcaps_dir(const std::string& dirname, const std::string& filter)
 {
-    for (const PcapReadObject &pro : pol)
+    Directory pcap_dir(dirname.c_str(), filter.c_str());
+    if (pcap_dir.error_on_open())
+    {
+        ErrorMessage("Error getting pcaps under dir: %s: %s\n",
+                dirname.c_str(), get_error(pcap_dir.error_on_open()));
+        return false;
+    }
+
+    std::vector<std::string> tmp_queue;
+    const char* pcap_filename;
+    while ((pcap_filename = pcap_dir.next()))
+        tmp_queue.emplace_back(pcap_filename);
+    std::sort(tmp_queue.begin(), tmp_queue.end());
+
+    pcap_queue.reserve(pcap_queue.size() + tmp_queue.size());
+    pcap_queue.insert(pcap_queue.end(), tmp_queue.begin(), tmp_queue.end());
+
+    return true;
+}
+
+bool Trough::add_pcaps_list_file(const std::string& list_filename, const std::string& filter)
+{
+    std::ifstream pcap_list_file(list_filename);
+    if (!pcap_list_file.is_open())
+    {
+        ErrorMessage("Could not open pcap list file: %s: %s\n", list_filename.c_str(), get_error(errno));
+        return false;
+    }
+
+    std::string pcap_name;
+    while (getline(pcap_list_file, pcap_name))
     {
-        const std::string& arg = pro.arg;
-        const std::string& filter = pro.filter;
+        /* Trim leading and trailing whitespace. */
+        constexpr const char* whitespace = " \f\n\r\t\v";
+        pcap_name.erase(0, pcap_name.find_first_not_of(whitespace));
+        pcap_name.erase(pcap_name.find_last_not_of(whitespace) + 1);
+
+        if (pcap_name.empty())
+            continue;
+
+        /* do a quick check to make sure file exists */
+        struct stat sb;
+        if (stat(pcap_name.c_str(), &sb) == -1)
+        {
+            ErrorMessage("Error getting stat on pcap file: %s: %s\n", pcap_name.c_str(), get_error(errno));
+            pcap_list_file.close();
+            return false;
+        }
+        if (S_ISDIR(sb.st_mode))
+        {
+            if (!add_pcaps_dir(pcap_name, filter))
+            {
+                pcap_list_file.close();
+                return false;
+            }
+        }
+        else if (S_ISREG(sb.st_mode))
+        {
+            if (filter.empty() || (fnmatch(filter.c_str(), pcap_name.c_str(), 0) == 0))
+                pcap_queue.emplace_back(pcap_name);
+        }
+        else
+        {
+            ErrorMessage("Specified entry in \'%s\' is not a regular file or directory: %s\n",
+                    list_filename.c_str(), pcap_name.c_str());
+            pcap_list_file.close();
+            return false;
+        }
+    }
+    pcap_list_file.close();
 
+    return true;
+}
+
+bool Trough::add_pcaps_list(const std::string& list)
+{
+    if (list.empty())
+    {
+        ErrorMessage("No pcaps specified in pcap list\n");
+        return false;
+    }
+
+    std::string pcap_name;
+    size_t i = 0;
+    size_t pos = 0;
+
+    do
+    {
+        pos = list.find(' ', i);
+        if (pos == std::string::npos)
+            pcap_name = list.substr(i);
+        else
+        {
+            pcap_name = list.substr(i, pos - i);
+            i = ++pos;
+        }
+        /* do a quick check to make sure file exists */
+        if (pcap_name != "-")
+        {
+            struct stat sb;
+            if (stat(pcap_name.c_str(), &sb) == -1)
+            {
+                ErrorMessage("Error getting stat on file: %s: %s (%d)\n",
+                        pcap_name.c_str(), get_error(errno), errno);
+                return false;
+            }
+            if (!(sb.st_mode & (S_IFREG|S_IFIFO)))
+            {
+                ErrorMessage("Specified pcap is not a regular file: %s\n", pcap_name.c_str());
+                return false;
+            }
+        }
+
+        pcap_queue.emplace_back(pcap_name);
+    } while (pos != std::string::npos);
+
+    return true;
+}
+
+bool Trough::get_pcaps(std::vector<struct PcapReadObject> &pol)
+{
+    for (const PcapReadObject &pro : pol)
+    {
         switch (pro.type)
         {
             case SOURCE_FILE_LIST:
                 /* arg should be a file with a list of pcaps in it */
-                {
-                    const char* whitespace = " \f\n\r\t\v";
-                    std::ifstream pcap_list_file(arg);
-                    std::string pcap_name;
-                    struct stat sb;
-
-                    if (!pcap_list_file.is_open())
-                    {
-                        ErrorMessage("Could not open pcap list file: %s: %s\n",
-                                arg.c_str(), get_error(errno));
-                        return -1;
-                    }
-
-                    while (getline(pcap_list_file, pcap_name))
-                    {
-                        /* Trim leading and trailing whitespace. */
-                        pcap_name.erase(0, pcap_name.find_first_not_of(whitespace));
-                        pcap_name.erase(pcap_name.find_last_not_of(whitespace) + 1);
-
-                        if (pcap_name.empty())
-                            continue;
-
-                        /* do a quick check to make sure file exists */
-                        if (snort::SnortConfig::read_mode() && stat(pcap_name.c_str(), &sb) == -1)
-                        {
-                            ErrorMessage("Error getting stat on pcap file: %s: %s\n",
-                                    pcap_name.c_str(), get_error(errno));
-                            pcap_list_file.close();
-                            return -1;
-                        }
-                        else if (snort::SnortConfig::read_mode() && S_ISDIR(sb.st_mode))
-                        {
-                            Directory pcap_dir(pcap_name.c_str(), filter.c_str());
-                            std::vector<std::string> tmp_queue;
-                            const char* pcap_filename;
-
-                            if (pcap_dir.error_on_open())
-                            {
-                                ErrorMessage("Error getting pcaps under dir: %s: %s\n",
-                                        pcap_name.c_str(), get_error(pcap_dir.error_on_open()));
-                                pcap_list_file.close();
-                                return -1;
-                            }
-                            while ((pcap_filename = pcap_dir.next()))
-                                tmp_queue.emplace_back(pcap_filename);
-                            std::sort(tmp_queue.begin(), tmp_queue.end());
-                            pcap_queue.reserve(pcap_queue.size() + tmp_queue.size());
-                            pcap_queue.insert(pcap_queue.end(), tmp_queue.begin(), tmp_queue.end());
-                        }
-                        else if (!snort::SnortConfig::read_mode() || S_ISREG(sb.st_mode))
-                        {
-                            if (filter.empty() ||
-                                (fnmatch(filter.c_str(), pcap_name.c_str(), 0) == 0))
-                                pcap_queue.emplace_back(pcap_name);
-                        }
-                        else
-                        {
-                            ErrorMessage("Specified entry in \'%s\' is not a regular file or "
-                                    "directory: %s\n", arg.c_str(), pcap_name.c_str());
-                            pcap_list_file.close();
-                            return -1;
-                        }
-                    }
-
-                    pcap_list_file.close();
-                }
-
+                if (!add_pcaps_list_file(pro.arg, pro.filter))
+                    return false;
                 break;
 
             case SOURCE_LIST:
                 /* arg should be a space separated list of pcaps */
-                {
-                    struct stat sb;
-                    std::string pcap_name;
-                    size_t i = 0;
-                    size_t pos = 0;
-
-                    if (arg.empty())
-                    {
-                        ErrorMessage("No pcaps specified in pcap list\n");
-                        return -1;
-                    }
-
-                    do
-                    {
-                        pos = arg.find(' ', i);
-                        if (pos == std::string::npos)
-                            pcap_name = arg.substr(i);
-                        else
-                        {
-                            pcap_name = arg.substr(i, pos - i);
-                            i = ++pos;
-                        }
-                        /* do a quick check to make sure file exists */
-                        if (snort::SnortConfig::read_mode() and pcap_name != "-")
-                        {
-                            if (stat(pcap_name.c_str(), &sb) == -1)
-                            {
-                                ErrorMessage("Error getting stat on file: %s: %s (%d)\n",
-                                        pcap_name.c_str(), get_error(errno), errno);
-                                return -1;
-                            }
-                            if (!(sb.st_mode & (S_IFREG|S_IFIFO)))
-                            {
-                                ErrorMessage("Specified pcap is not a regular file: %s\n",
-                                        pcap_name.c_str());
-                                return -1;
-                            }
-                        }
-
-                        pcap_queue.emplace_back(pcap_name);
-                    } while (pos != std::string::npos);
-                }
-
+                if (!add_pcaps_list(pro.arg))
+                    return false;
                 break;
 
             case SOURCE_DIR:
                 /* arg should be a directory name */
-                {
-                    Directory pcap_dir(arg.c_str(), filter.c_str());
-                    std::vector<std::string> tmp_queue;
-                    const char* pcap_filename;
-
-                    if (pcap_dir.error_on_open())
-                    {
-                        ErrorMessage("Error getting pcaps under dir: %s: %s\n",
-                                arg.c_str(), get_error(pcap_dir.error_on_open()));
-                        return -1;
-                    }
-                    while ((pcap_filename = pcap_dir.next()))
-                        tmp_queue.emplace_back(pcap_filename);
-                    std::sort(tmp_queue.begin(), tmp_queue.end());
-                    pcap_queue.reserve(pcap_queue.size() + tmp_queue.size());
-                    pcap_queue.insert(pcap_queue.end(), tmp_queue.begin(), tmp_queue.end());
-                }
-
+                if (!add_pcaps_dir(pro.arg, pro.filter))
+                    return false;
                 break;
         }
     }
 
-    return 0;
+    return true;
 }
 
 void Trough::add_source(SourceType type, const char* list)
@@ -210,7 +205,7 @@ void Trough::add_source(SourceType type, const char* list)
     pcap_object_list.emplace_back(pro);
 }
 
-void Trough::set_filter(const char *f)
+void Trough::set_filter(const char* f)
 {
     if (f)
         pcap_filter = f;
@@ -222,7 +217,7 @@ void Trough::setup()
 {
     if (!pcap_object_list.empty())
     {
-        if (get_pcaps(pcap_object_list) == -1)
+        if (!get_pcaps(pcap_object_list))
             FatalError("Error getting pcaps.\n");
 
         if (pcap_queue.empty())

diff --git a/src/packet_io/trough.h b/src/packet_io/trough.h
index 7161d425d4d784bc1f310b01520eea459a879010..c25abf7b78d7cd4de93c9307d605b07aaec9f0e4 100644 (file)
--- a/src/packet_io/trough.h
+++ b/src/packet_io/trough.h
@@ -65,7 +65,11 @@ private:
         std::string filter;
     };
 
-    static int get_pcaps(std::vector<struct PcapReadObject> &pol);
+    static bool add_pcaps_dir(const std::string& dirname, const std::string& filter);
+    static bool add_pcaps_list_file(const std::string& list_filename, const std::string& filter);
+    static bool add_pcaps_list(const std::string& list);
+    static bool get_pcaps(std::vector<struct PcapReadObject> &pol);
+
     static std::vector<struct PcapReadObject> pcap_object_list;
     static std::vector<std::string> pcap_queue;
     static std::vector<std::string>::const_iterator pcap_queue_iter;