When using -av option, now VT_Count returned for each sample consider…

…s only the AVs in the give list

avclass2/avclass2_labeler.py

@@ -169,8 +169,8 @@ def main(args):
                 # sys.stderr.flush()
                 continue
 
-            # Compute VT_Count
-            vt_count = len(sample_info.labels)
+            # Compute VT_Count (using list of AV engines if provided)
+            vt_count = av_labels.get_sample_vt_count(sample_info)
 
             # Get the distinct tokens from all the av labels in the report
             # And print them. 

avclass2/lib/avclass2_common.py

@@ -400,6 +400,8 @@ def read_avs(avs_file):
         '''Read AV engine set from given file'''
         with open(avs_file) as fd:
             avs = set(map(str.strip, fd.readlines()))
+        sys.stderr.write("[-] Using %d AV engines in %s\n" % (len(avs),
+                                                              avs_file))
         return avs
 
     @staticmethod
@@ -669,3 +671,15 @@ def rank_tags(self, av_dict, threshold=1):
                     if len(avs) > threshold)
         return sorted(pairs, key=itemgetter(1,0), reverse=True)
 
+    def get_sample_vt_count(self, sample_info):
+        ''' Return number of detections for sample
+            in the provided AV whitelist (if any) '''
+        if self.avs is None:
+            return len(sample_info.labels)
+        else:
+            cnt = 0
+            for (av_name, label) in sample_info.labels:
+                if av_name in self.avs:
+                    cnt += 1
+            return cnt
+