From f12cc8daf4073fb746ee19771f2374e5101d07b5 Mon Sep 17 00:00:00 2001
From: greninja <sshadab84@gmail.com>
Date: Wed, 28 Dec 2016 21:31:10 +0530
Subject: [PATCH] Enhancement features #15

---
 lib/cuckoo/common/config.py    |  1 +
 lib/cuckoo/common/constants.py |  2 ++
 modules/processing/cuckooml.py | 14 ++++++++++++++
 3 files changed, 17 insertions(+)

diff --git a/lib/cuckoo/common/config.py b/lib/cuckoo/common/config.py
index f3c5c7f8..76228a70 100644
--- a/lib/cuckoo/common/config.py
+++ b/lib/cuckoo/common/config.py
@@ -72,3 +72,4 @@ def parse_options(options):
 def emit_options(options):
     """Emit the analysis options from a dictionary to a string."""
     return ",".join("%s=%s" % (k, v) for k, v in options.items())
+
diff --git a/lib/cuckoo/common/constants.py b/lib/cuckoo/common/constants.py
index 3418fd18..c23bf6d5 100644
--- a/lib/cuckoo/common/constants.py
+++ b/lib/cuckoo/common/constants.py
@@ -8,6 +8,8 @@
 _current_dir = os.path.abspath(os.path.dirname(__file__))
 CUCKOO_ROOT = os.path.normpath(os.path.join(_current_dir, "..", "..", ".."))
 
+print os.listdir(CUCKOO_ROOT + "/" + "sample_data/dict")
+
 CUCKOO_VERSION = "2.0-dev"
 CUCKOO_GUEST_PORT = 8000
 CUCKOO_GUEST_INIT = 0x001
diff --git a/modules/processing/cuckooml.py b/modules/processing/cuckooml.py
index 14912c24..9d8e174c 100644
--- a/modules/processing/cuckooml.py
+++ b/modules/processing/cuckooml.py
@@ -797,6 +797,13 @@ def filter_dataset(self, dataset=None, feature_coverage=0.1,
 
     def detect_abnormal_behaviour(self, count_dataset=None, figures=True):
         """Detect samples that behave significantly different than others."""
+        
+        # Safety check for plotting
+        if not Config("cuckooml").cuckooml.plotting and figures:
+            print >> sys.stderr, "Warning: 'plotting' and 'figures' do not match. Plotting modules might not be imported."
+            figures = False
+
+
         if count_dataset is None:
             # Pull all count features
             count_features = self.feature_category(":count:")
@@ -1133,6 +1140,13 @@ def performance_metric(clustering, labels, data, noise):
 
     def clustering_label_distribution(self, clustering, labels, plot=False):
         """Get statistics about number of ground truth labels per cluster."""
+        
+        # Safety check for plotting
+        if not Config("cuckooml").cuckooml.plotting and plot:
+            print >> sys.stderr, "Warning: 'plotting' and 'plot' do not match. Plotting modules might not be imported."
+            plot = False
+
+
         cluster_ids = set(clustering["label"].tolist())
         labels_ids = set(labels["label"].tolist())
         cluster_distribution = {}