Merge pull request #10 from nprint/index

Index

Author: Jordan Holland

Makefile.am

@@ -2,14 +2,16 @@ AUTOMAKE_OPTIONS = subdir-objects
 
 bin_PROGRAMS = pcapml
 
-pcapml_SOURCES = include/dir/dir.hpp include/label/label.hpp include/label/labeler.hpp \
+pcapml_SOURCES = include/label/dir.hpp include/label/label.hpp include/label/labeler.hpp \
+				 include/label/pcap_labeler.hpp  \
 				 include/pcap/reader_pcap.hpp include/pcap/writer_pcap.hpp \
 				 include/pcapng/block_pcapng.hpp include/pcapng/reader_pcapng.hpp \
 				 include/pcapng/sorter_pcapng.hpp include/pcapng/splitter_pcapng.hpp \
 				 include/pcapng/writer_pcapng.hpp include/pcapng/stripper_pcapng.hpp \
 				 include/sample/sample.hpp include/sample/sampler.hpp \
 				 include/util.hpp \
-				 src/dir/dir.cpp src/label/label.cpp src/label/labeler.cpp \
+				 src/label/dir.cpp src/label/label.cpp src/label/labeler.cpp \
+				 src/label/pcap_labeler.cpp \
 				 src/pcap/reader_pcap.cpp src/pcap/writer_pcap.cpp \
 				 src/pcapng/block_pcapng.cpp src/pcapng/reader_pcapng.cpp \
 				 src/pcapng/sorter_pcapng.cpp src/pcapng/splitter_pcapng.cpp \

include/dir/dir.hpp

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2020 nPrint
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy
+ * of the License at https://www.apache.org/licenses/LICENSE-2.0
+ */
+
+#ifndef INCLUDE_LABEL_DIR_HPP_
+#define INCLUDE_LABEL_DIR_HPP_
+
+#define MAX_FILE_LEN 2056
+
+#include <stdio.h>
+#include <dirent.h>
+#include <string.h>
+#include <sys/stat.h>
+
+#include <map>
+#include <string>
+
+#include "labeler.hpp"
+
+class DirLabeler : public Labeler {
+ public:
+    void print_stats();
+    uint32_t process_packet(PcapPacketInfo *pi);
+    uint32_t label_dir(std::string dir, std::string label_file,
+                       std::string outfile, bool stats_out);
+ private:
+    uint64_t files_processed = 0;
+    uint64_t files_skipped = 0;
+    Label *active_file_label;
+    int16_t cur_linktype = 999;
+    std::map<std::string, Label *> file_labels;
+
+    void build_file_label_map();
+    uint32_t label_file(Label *l);
+    uint32_t process_directory(std::string dir);
+};
+
+#endif  // INCLUDE_LABEL_DIR_HPP_

include/label/dir.hpp

@@ -18,20 +18,22 @@
 class Label {
  public:
     void print();
-    bool set_info(std::string label, std::string bpf_filter = "",
+    uint32_t set_info(std::string label, std::string bpf_filter = "",
+                  std::string file = "", std::string hash_key = "",
                   uint64_t ts_start = 0, uint64_t ts_end = UINT64_MAX,
                   pcap_t *handle = NULL);
+    std::string get_file();
     std::string get_label();
     std::string get_sample_id();
     std::string get_comment_string();
     std::string get_unhashed_sample_id();
-    bool match_packet(pcap_packet_info *pi);
+    bool match_packet(PcapPacketInfo *pi);
  private:
     bool info_set = false;
+    uint64_t ts_start, ts_end;
     std::string bpf_string_filter;
     struct bpf_program *bpf_pcap_filter = NULL;
-    uint64_t ts_start, ts_end;
-    std::string sample_id, unhashed_sample_id, label, comment_str;
+    std::string sample_id, unhashed_sample_id, label, comment_str, file, hash_key;
 };
 
 #endif  // INCLUDE_LABEL_LABEL_HPP_

include/label/label.hpp

@@ -10,32 +10,39 @@
 
 #include <signal.h>
 
+static volatile int stop = 0;
+
+#include <string>
 #include <vector>
 #include <fstream>
 
-static volatile int stop = 0;
-
 #include "util.hpp"
 #include "label.hpp"
 #include "reader_pcap.hpp"
 #include "writer_pcapng.hpp"
 
-#define LABEL_FILE_LOC  0
-#define FILTER_LOC 1
-#define TS_START   2
-#define TS_END     3
+#define TRAFFIC_LOC  0
+#define METADATA_LOC 1
+#define HASHKEY_LOC  2
 
-class PcapMLLabeler {
+class Labeler {
  public:
-    void print_stats(FILE *stream);
-    bool label_pcap(char *label_file, char *pcap, char *outfile,
-                    bool infile_is_device, bool print_stats);
- private:
+    virtual void print_stats() = 0;
+    virtual uint32_t process_packet(PcapPacketInfo *pi) = 0;
+
+    int process_traffic(PcapReader r);
+    int load_labels(std::string label_file, pcap_t *handle = NULL);
+ protected:
+    PcapNGWriter w;
     std::vector<Label *> labels;
-    uint64_t packets_matched = 0;
+    uint64_t packets_matched =  0;
     uint64_t packets_received = 0;
-
-    bool load_labels(char *label_file, pcap_t *handle = NULL);
+ private:
+    Label *process_label_line(std::string line, pcap_t *handle = NULL);
+    Label *process_traffic_filter(std::string traffic_filter,
+                                  std::string hash_key,
+                                  std::string metadata,
+                                  pcap_t *handle = NULL);
 };
 
 #endif  // INCLUDE_LABEL_LABELER_HPP_

include/label/labeler.hpp

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2020 nPrint
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy
+ * of the License at https://www.apache.org/licenses/LICENSE-2.0
+ */
+
+#ifndef INCLUDE_LABEL_PCAP_LABELER_HPP_
+#define INCLUDE_LABEL_PCAP_LABELER_HPP_
+
+#include "labeler.hpp"
+
+class PcapLabeler : public Labeler {
+ public:
+    void print_stats();
+    uint32_t process_packet(PcapPacketInfo *pi);
+    uint32_t label_pcap(char *label_file, char *pcap, char *outfile,
+                        bool infile_is_device, bool print_stats);
+};
+
+#endif  // INCLUDE_LABEL_PCAP_LABELER_HPP_

include/label/pcap_labeler.hpp

@@ -26,12 +26,12 @@
 class PcapReader {
  public:
     void close_file();
+    void print_stats();
     pcap_t *get_pcap_t();
     uint16_t get_linktype();
     int open_live(char *devce);
     int open_file(char *infile);
-    void print_stats(FILE *stream);
-    pcap_packet_info *get_next_packet();
+    PcapPacketInfo *get_next_packet();
  private:
     pcap_t *f = NULL;
 };

include/pcap/reader_pcap.hpp

@@ -42,7 +42,7 @@ class PcapNGWriter {
     void close_file();
     int write_section_block();
     int write_interface_block(uint16_t link_type, uint32_t snap_len);
-    int write_epb_from_pcap_pkt(pcap_packet_info *pi, std::string comment);
+    int write_epb_from_pcap_pkt(PcapPacketInfo *pi, std::string comment);
  private:
     FILE *pcapng;
 

include/pcapng/writer_pcapng.hpp

@@ -19,7 +19,7 @@
 
 #include "block_pcapng.hpp"
 
-struct pcap_packet_info {
+struct PcapPacketInfo {
     struct pcap_pkthdr *hdr;
     const uint8_t *buf;
     int32_t pcap_next_rv;