Implemented custom category loading.

Modified ndpiReader with -c <file> (use mining_hosts.txt as example) for loading name-based categories

1 files changed, 139 insertions, 100 deletions

diff --git a/example/ndpiReader.c b/example/ndpiReader.c
index 6a4c31ade..3069adf8e 100644
--- a/example/ndpiReader.c
+++ b/example/ndpiReader.c
@@ -1,7 +1,7 @@
 /*
  * ndpiReader.c
  *
- * Copyright (C) 2011-17 - ntop.org
+ * Copyright (C) 2011-18 - ntop.org
  *
  * nDPI is free software: you can redistribute it and/or modify
  * it under the terms of the GNU Lesser General Public License as published by
@@ -67,6 +67,7 @@ static FILE *results_file           = NULL;
 static char *results_path           = NULL;
 static char * bpfFilter             = NULL; /**< bpf filter  */
 static char *_protoFilePath         = NULL; /**< Protocol file path  */
+static char *_customCategoryFilePath= NULL; /**< Custom categories file path  */
 #ifdef HAVE_JSON_C
 static char *_statsFilePath         = NULL; /**< Top stats file path */
 static char *_diagnoseFilePath      = NULL; /**< Top stats file path */
@@ -104,8 +105,8 @@ static u_int32_t num_flows;
 static struct ndpi_detection_module_struct *ndpi_info_mod = NULL;
 
 struct flow_info {
-	struct ndpi_flow_info *flow;
-	u_int16_t thread_id;
+  struct ndpi_flow_info *flow;
+  u_int16_t thread_id;
 };
 
 static struct flow_info *all_flows;
@@ -230,7 +231,7 @@ static void help(u_int long_help) {
 
   printf("ndpiReader -i <file|device> [-f <filter>][-s <duration>][-m <duration>]\n"
 	 "          [-p <protos>][-l <loops> [-q][-d][-h][-t][-v <level>]\n"
-	 "          [-n <threads>] [-w <file>] [-j <file>] [-x <file>] \n\n"
+	 "          [-n <threads>][-w <file>][-c <file>][-j <file>][-x <file>]\n\n"
 	 "Usage:\n"
 	 "  -i <file.pcap|device>     | Specify a pcap file/playlist to read packets from or a\n"
 	 "                            | device for live capture (comma-separated list)\n"
@@ -249,6 +250,7 @@ static void help(u_int long_help) {
 	 "  -q                        | Quiet mode\n"
 	 "  -t                        | Dissect GTP/TZSP tunnels\n"
 	 "  -r                        | Print nDPI version and git revision\n"
+	 "  -c <path>                 | Load custom categories from the specified file\n"
 	 "  -w <path>                 | Write test output on the specified file. This is useful for\n"
 	 "                            | testing purposes in order to compare results across runs\n"
 	 "  -h                        | This help\n"
@@ -276,7 +278,7 @@ static void help(u_int long_help) {
 	 "  --fifo <path to file or pipe>\n"
 	 "  --debug\n"
 	 "  --dbg-proto proto|num[,...]\n"
-	 );
+    );
 #endif
 
   if(long_help) {
@@ -360,7 +362,7 @@ int cmpFlows(const void *_a, const void *_b) {
   uint64_t a_size = fa->src2dst_bytes + fa->dst2src_bytes;
   uint64_t b_size = fb->src2dst_bytes + fb->dst2src_bytes;
   if(a_size != b_size)
-	  return a_size < b_size ? 1 : -1;
+    return a_size < b_size ? 1 : -1;
 
 // copy from ndpi_workflow_node_cmp();
 
@@ -449,7 +451,7 @@ static void parseOptions(int argc, char **argv) {
   if(trace) fprintf(trace, " #### %s #### \n", __FUNCTION__);
 #endif
 
-  while ((opt = getopt_long(argc, argv, "df:g:i:hp:l:s:tv:V:n:j:rp:w:q0123:456:7:89:m:b:x:", longopts, &option_idx)) != EOF) {
+  while ((opt = getopt_long(argc, argv, "c:df:g:i:hp:l:s:tv:V:n:j:rp:w:q0123:456:7:89:m:b:x:", longopts, &option_idx)) != EOF) {
 #ifdef DEBUG_TRACE
     if(trace) fprintf(trace, " #### -%c [%s] #### \n", opt, optarg ? optarg : "");
 #endif
@@ -507,6 +509,10 @@ static void parseOptions(int argc, char **argv) {
       _protoFilePath = optarg;
       break;
 
+    case 'c':
+      _customCategoryFilePath = optarg;
+      break;
+
     case 's':
       capture_for = atoi(optarg);
       capture_until = capture_for + time(NULL);
@@ -528,8 +534,8 @@ static void parseOptions(int argc, char **argv) {
       nDPI_LogLevel  = atoi(optarg);
       if(nDPI_LogLevel < 0) nDPI_LogLevel = 0;
       if(nDPI_LogLevel > 3) {
-	 nDPI_LogLevel = 3;
-      	 _debug_protocols = strdup("all");
+	nDPI_LogLevel = 3;
+	_debug_protocols = strdup("all");
       }
       break;
 
@@ -738,7 +744,7 @@ static void printFlow(u_int16_t id, struct ndpi_flow_info *flow, u_int16_t threa
 	    flow->bidirectional ? "<->" : "->",
 	    (flow->ip_version == 6) ? "[" : "",
 	    flow->dst_name, (flow->ip_version == 6) ? "]" : "", ntohs(flow->dst_port)
-	    );
+      );
 
     if(flow->vlan_id > 0) fprintf(out, "[VLAN: %u]", flow->vlan_id);
 
@@ -834,9 +840,9 @@ static void node_print_unknown_proto_walker(const void *node, ndpi_VISIT which,
   if(flow->detected_protocol.app_protocol != NDPI_PROTOCOL_UNKNOWN) return;
 
   if((which == ndpi_preorder) || (which == ndpi_leaf)) {
-     /* Avoid walking the same node multiple times */
-     all_flows[num_flows].thread_id = thread_id, all_flows[num_flows].flow = flow;
-     num_flows++;
+    /* Avoid walking the same node multiple times */
+    all_flows[num_flows].thread_id = thread_id, all_flows[num_flows].flow = flow;
+    num_flows++;
   }
 }
 
@@ -851,9 +857,9 @@ static void node_print_known_proto_walker(const void *node, ndpi_VISIT which, in
   if(flow->detected_protocol.app_protocol == NDPI_PROTOCOL_UNKNOWN) return;
 
   if((which == ndpi_preorder) || (which == ndpi_leaf)) {
-     /* Avoid walking the same node multiple times */
-     all_flows[num_flows].thread_id = thread_id, all_flows[num_flows].flow = flow;
-     num_flows++;
+    /* Avoid walking the same node multiple times */
+    all_flows[num_flows].thread_id = thread_id, all_flows[num_flows].flow = flow;
+    num_flows++;
   }
 }
 
@@ -948,7 +954,7 @@ void updateScanners(struct single_flow_info **scanners, u_int32_t saddr,
 /* *********************************************** */
 
 int updateIpTree(u_int32_t key, u_int8_t version,
-                  addr_node **vrootp, const char *proto) {
+		 addr_node **vrootp, const char *proto) {
   addr_node *q;
   addr_node **rootp = vrootp;
 
@@ -1116,7 +1122,7 @@ static int receivers_sort_asc(void *_a, void *_b) {
 /* ***************************************************** */
 /*@brief removes first (size - max) elements from hash table.
  * hash table is ordered in ascending order.
-*/
+ */
 static struct receiver *cutBackTo(struct receiver **receivers, u_int32_t size, u_int32_t max) {
   struct receiver *r, *tmp;
   int i=0;
@@ -1143,7 +1149,7 @@ static struct receiver *cutBackTo(struct receiver **receivers, u_int32_t size, u
  * if element already in the second table
  *  then updates its value
  * else adds it to the second table
-*/
+ */
 static void mergeTables(struct receiver **primary, struct receiver **secondary) {
   struct receiver *r, *s, *tmp;
 
@@ -1191,7 +1197,7 @@ static void deleteReceivers(struct receiver *receivers) {
  * }
  * else
  *   update table1
-*/
+ */
 static void updateReceivers(struct receiver **receivers, u_int32_t dst_addr,
                             u_int8_t version, u_int32_t num_pkts,
                             struct receiver **topReceivers) {
@@ -1202,7 +1208,7 @@ static void updateReceivers(struct receiver **receivers, u_int32_t dst_addr,
   HASH_FIND_INT(*receivers, (int *)&dst_addr, r);
   if(r == NULL) {
     if(((size = HASH_COUNT(*receivers)) < MAX_TABLE_SIZE_1)
-    || ((a = acceptable(num_pkts)) != 0)){
+       || ((a = acceptable(num_pkts)) != 0)){
       r = (struct receiver *)malloc(sizeof(struct receiver));
       if(!r) return;
 
@@ -1320,7 +1326,7 @@ static void port_stats_walker(const void *node, ndpi_VISIT which, int depth, voi
 
     if(((r = strcmp(ipProto2Name(flow->protocol), "TCP")) == 0)
        && (flow->src2dst_packets == 1) && (flow->dst2src_packets == 0)) {
-       updateScanners(&scannerHosts, flow->src_ip, flow->ip_version, dport);
+      updateScanners(&scannerHosts, flow->src_ip, flow->ip_version, dport);
     }
 
     updateReceivers(&receivers, flow->dst_ip, flow->ip_version,
@@ -1463,6 +1469,39 @@ static void setupDetection(u_int16_t thread_id, pcap_t * pcap_handle) {
 
   if(_protoFilePath != NULL)
     ndpi_load_protocols_file(ndpi_thread_info[thread_id].workflow->ndpi_struct, _protoFilePath);
+
+  if(_customCategoryFilePath) {
+    FILE *fd = fopen(_customCategoryFilePath, "r");
+
+    if(fd) {
+      while(fd) {
+	char buffer[512], *line, *name, *category;
+	int i;
+	
+	if(!(line = fgets(buffer, sizeof(buffer), fd)))
+	  break;
+	
+	if(((i = strlen(line)) <= 1) || (line[0] == '#'))
+	  continue;
+	else
+	  line[i-1] = '\0';
+
+	name = strtok(line, "\t");
+	if(name) {
+	  category = strtok(NULL, "\t");
+
+	  if(category) {
+	    // printf("Loading %s\t%s\n", name, category);
+	    ndpi_load_hostname_category(ndpi_thread_info[thread_id].workflow->ndpi_struct,
+					name, (ndpi_protocol_category_t)atoi(category));
+	  }
+	}
+      }
+
+      ndpi_enable_loaded_categories(ndpi_thread_info[thread_id].workflow->ndpi_struct);
+    } else
+      printf("ERROR: Unable to read file %s\n", _customCategoryFilePath);
+  }
 }
 
 
@@ -2045,8 +2084,8 @@ static void printResults(u_int64_t tot_usec) {
       total_flows += ndpi_thread_info[thread_id].workflow->num_allocated_flows;
 
     if((all_flows = (struct flow_info*)malloc(sizeof(struct flow_info)*total_flows)) == NULL) {
-       printf("Fatal error: not enough memory\n");
-       exit(-1);
+      printf("Fatal error: not enough memory\n");
+      exit(-1);
     }
 
     if(!json_flag) fprintf(out, "\n");
@@ -2130,12 +2169,12 @@ static void printResults(u_int64_t tot_usec) {
     saveScannerStats(&jObj_stats, &scannerHosts);
 
     if((count = HASH_COUNT(topReceivers)) == 0){
-        HASH_SORT(receivers, receivers_sort);
-        saveReceiverStats(&jObj_stats, &receivers, cumulative_stats.ip_packet_count);
+      HASH_SORT(receivers, receivers_sort);
+      saveReceiverStats(&jObj_stats, &receivers, cumulative_stats.ip_packet_count);
     }
     else{
-        HASH_SORT(topReceivers, receivers_sort);
-        saveReceiverStats(&jObj_stats, &topReceivers, cumulative_stats.ip_packet_count);
+      HASH_SORT(topReceivers, receivers_sort);
+      saveReceiverStats(&jObj_stats, &topReceivers, cumulative_stats.ip_packet_count);
     }
 
     u_int64_t total_src_addr = getTopStats(srcStats);
@@ -2151,7 +2190,7 @@ static void printResults(u_int64_t tot_usec) {
 #endif
   }
 
- free_stats:
+free_stats:
   if(scannerHosts) {
     deleteScanners(scannerHosts);
     scannerHosts = NULL;
@@ -2213,7 +2252,7 @@ static int getNextPcapFileFromPlaylist(u_int16_t thread_id, char filename[], u_i
       return -1;
   }
 
- next_line:
+next_line:
   if(fgets(filename, filename_len, playlist_fp[thread_id])) {
     int l = strlen(filename);
     if(filename[0] == '\0' || filename[0] == '#') goto next_line;
@@ -2272,7 +2311,7 @@ static pcap_t * openPcapFileOrDevice(u_int16_t thread_id, const u_char * pcap_fi
       if(strstr((char*)pcap_file, (char*)".pcap"))
 	printf("ERROR: could not open pcap file %s: %s\n", pcap_file, pcap_error_buffer);
       else if((getNextPcapFileFromPlaylist(thread_id, filename, sizeof(filename)) != 0)
-	 || ((pcap_handle = pcap_open_offline(filename, pcap_error_buffer)) == NULL)) {
+	      || ((pcap_handle = pcap_open_offline(filename, pcap_error_buffer)) == NULL)) {
         printf("ERROR: could not open playlist %s: %s\n", filename, pcap_error_buffer);
         exit(-1);
       } else {
@@ -2361,8 +2400,8 @@ static void pcap_process_packet(u_char *args,
      && ((extcap_packet_filter == (u_int16_t)-1)
 	 || (p.app_protocol == extcap_packet_filter)
 	 || (p.master_protocol == extcap_packet_filter)
-	 )
-     ) {
+       )
+    ) {
     struct pcap_pkthdr h;
     uint32_t *crc, delta = sizeof(struct ndpi_packet_trailer) + 4 /* ethernet trailer */;
     struct ndpi_packet_trailer *trailer;
@@ -2454,7 +2493,7 @@ void * processing_thread(void *_thread_id) {
 #endif
     if((!json_flag) && (!quiet_mode)) printf("Running thread %ld...\n", thread_id);
 
- pcap_loop:
+pcap_loop:
   runPcapLoop(thread_id);
 
   if(playlist_fp[thread_id] != NULL) { /* playlist: read next file */
@@ -2881,8 +2920,8 @@ void getSourcePorts(struct json_object *jObj_stat, int srcPortArray[], int size,
 
 
     if((flows_packets > FLOWS_PACKETS_THRESHOLD)
-        && (flows_percent >= FLOWS_PERCENT_THRESHOLD)
-        && packets_number >= threshold) {
+       && (flows_percent >= FLOWS_PERCENT_THRESHOLD)
+       && packets_number >= threshold) {
       if((res = json_object_object_get_ex(src_pkts_stat, "port", &jObj_port)) == 0) {
 	fprintf(stderr, "ERROR: can't get \"port\", use -x flag only with .json files generated by ndpiReader -b flag.\n");
 	exit(-1);
@@ -2967,7 +3006,7 @@ void getScannerHosts(struct json_object *jObj_stat, int duration,
 
 #ifdef HAVE_JSON_C
 void getDestinationHosts(struct json_object *jObj_stat, int duration,
-                 const char *dstHostArray[16], int size) {
+			 const char *dstHostArray[16], int size) {
   int j;
 
   for(j=0; j<json_object_array_length(jObj_stat); j++) {
@@ -3119,7 +3158,7 @@ static void produceBpfFilter(char *filePath) {
   jObj_bpfFilter = json_object_new_object();
 
   bpf_filter_pkt_peak_filter(&jObj_bpfFilter, filterSrcPorts, PORT_ARRAY_SIZE,
-                  filterSrcHosts, HOST_ARRAY_SIZE, filterPktDstHosts, HOST_ARRAY_SIZE/2);
+			     filterSrcHosts, HOST_ARRAY_SIZE, filterPktDstHosts, HOST_ARRAY_SIZE/2);
 
   bpf_filter_host_peak_filter(&jObj_bpfFilter, filterDstHosts, HOST_ARRAY_SIZE);
 
@@ -3142,54 +3181,54 @@ static void produceBpfFilter(char *filePath) {
 #ifdef APP_HAS_OWN_MAIN
 int orginal_main(int argc, char **argv) {
 #else
-int main(int argc, char **argv) {
+  int main(int argc, char **argv) {
 #endif
-  int i;
+    int i;
 
-  if(ndpi_get_api_version() != NDPI_API_VERSION) {
-    printf("nDPI Library version mismatch: please make sure this code and the nDPI library are in sync\n");
-    return(-1);
-  }
+    if(ndpi_get_api_version() != NDPI_API_VERSION) {
+      printf("nDPI Library version mismatch: please make sure this code and the nDPI library are in sync\n");
+      return(-1);
+    }
   
-  automataUnitTest();
+    automataUnitTest();
 
-  ndpi_info_mod = ndpi_init_detection_module();
-  if (ndpi_info_mod == NULL) return -1;
+    ndpi_info_mod = ndpi_init_detection_module();
+    if (ndpi_info_mod == NULL) return -1;
 
-  memset(ndpi_thread_info, 0, sizeof(ndpi_thread_info));
+    memset(ndpi_thread_info, 0, sizeof(ndpi_thread_info));
 
-  parseOptions(argc, argv);
+    parseOptions(argc, argv);
 
-  if(bpf_filter_flag) {
+    if(bpf_filter_flag) {
 #ifdef HAVE_JSON_C
-    produceBpfFilter(_diagnoseFilePath);
-    return 0;
+      produceBpfFilter(_diagnoseFilePath);
+      return 0;
 #endif
-  }
+    }
 
-  if((!json_flag) && (!quiet_mode)) {
-    printf("\n-----------------------------------------------------------\n"
-	   "* NOTE: This is demo app to show *some* nDPI features.\n"
-	   "* In this demo we have implemented only some basic features\n"
-	   "* just to show you what you can do with the library. Feel \n"
-	   "* free to extend it and send us the patches for inclusion\n"
-	   "------------------------------------------------------------\n\n");
+    if((!json_flag) && (!quiet_mode)) {
+      printf("\n-----------------------------------------------------------\n"
+	     "* NOTE: This is demo app to show *some* nDPI features.\n"
+	     "* In this demo we have implemented only some basic features\n"
+	     "* just to show you what you can do with the library. Feel \n"
+	     "* free to extend it and send us the patches for inclusion\n"
+	     "------------------------------------------------------------\n\n");
 
-    printf("Using nDPI (%s) [%d thread(s)]\n", ndpi_revision(), num_threads);
-  }
+      printf("Using nDPI (%s) [%d thread(s)]\n", ndpi_revision(), num_threads);
+    }
 
-  signal(SIGINT, sigproc);
+    signal(SIGINT, sigproc);
 
-  for(i=0; i<num_loops; i++)
-    test_lib();
+    for(i=0; i<num_loops; i++)
+      test_lib();
 
-  if(results_path)  free(results_path);
-  if(results_file)  fclose(results_file);
-  if(extcap_dumper) pcap_dump_close(extcap_dumper);
-  if(ndpi_info_mod) ndpi_exit_detection_module(ndpi_info_mod);
+    if(results_path)  free(results_path);
+    if(results_file)  fclose(results_file);
+    if(extcap_dumper) pcap_dump_close(extcap_dumper);
+    if(ndpi_info_mod) ndpi_exit_detection_module(ndpi_info_mod);
 
-  return 0;
-}
+    return 0;
+  }
 
 #ifdef WIN32
 #ifndef __GNUC__
@@ -3201,42 +3240,42 @@ int main(int argc, char **argv) {
 /**
    @brief Timezone
 **/
-struct timezone {
-  int tz_minuteswest; /* minutes W of Greenwich */
-  int tz_dsttime;     /* type of dst correction */
-};
+  struct timezone {
+    int tz_minuteswest; /* minutes W of Greenwich */
+    int tz_dsttime;     /* type of dst correction */
+  };
 
 
 /**
    @brief Set time
 **/
-int gettimeofday(struct timeval *tv, struct timezone *tz) {
-  FILETIME        ft;
-  LARGE_INTEGER   li;
-  __int64         t;
-  static int      tzflag;
-
-  if(tv) {
-    GetSystemTimeAsFileTime(&ft);
-    li.LowPart  = ft.dwLowDateTime;
-    li.HighPart = ft.dwHighDateTime;
-    t  = li.QuadPart;       /* In 100-nanosecond intervals */
-    t -= EPOCHFILETIME;     /* Offset to the Epoch time */
-    t /= 10;                /* In microseconds */
-    tv->tv_sec  = (long)(t / 1000000);
-    tv->tv_usec = (long)(t % 1000000);
-  }
+  int gettimeofday(struct timeval *tv, struct timezone *tz) {
+    FILETIME        ft;
+    LARGE_INTEGER   li;
+    __int64         t;
+    static int      tzflag;
+
+    if(tv) {
+      GetSystemTimeAsFileTime(&ft);
+      li.LowPart  = ft.dwLowDateTime;
+      li.HighPart = ft.dwHighDateTime;
+      t  = li.QuadPart;       /* In 100-nanosecond intervals */
+      t -= EPOCHFILETIME;     /* Offset to the Epoch time */
+      t /= 10;                /* In microseconds */
+      tv->tv_sec  = (long)(t / 1000000);
+      tv->tv_usec = (long)(t % 1000000);
+    }
+
+    if(tz) {
+      if(!tzflag) {
+	_tzset();
+	tzflag++;
+      }
 
-  if(tz) {
-    if(!tzflag) {
-      _tzset();
-      tzflag++;
+      tz->tz_minuteswest = _timezone / 60;
+      tz->tz_dsttime = _daylight;
     }
 
-    tz->tz_minuteswest = _timezone / 60;
-    tz->tz_dsttime = _daylight;
+    return 0;
   }
-
-  return 0;
-}
 #endif /* WIN32 */