From 2caa6e42ea58efa98cc6414ef23e9c8d4168c4e3 Mon Sep 17 00:00:00 2001
From: Yogesh Ojha <yogesh.ojha11@gmail.com>
Date: Thu, 11 Nov 2021 13:59:20 +0530
Subject: [PATCH] Added color code for trees

---
 .../templates/startScan/detail_scan.html      |  22 ++-
 web/static/assets/css/app-dark.min.css        |   4 +
 web/static/custom/custom.css                  |  32 ++++-
 web/static/custom/mitch.js                    | 127 +++++++++++++++++-
 4 files changed, 175 insertions(+), 10 deletions(-)

diff --git a/web/startScan/templates/startScan/detail_scan.html b/web/startScan/templates/startScan/detail_scan.html
index 3a9538c8..8762bb0d 100644
--- a/web/startScan/templates/startScan/detail_scan.html
+++ b/web/startScan/templates/startScan/detail_scan.html
@@ -18,6 +18,8 @@ Detailed Scan
 <link href="{% static 'plugins/lightbox/css/lightbox.css' %}" rel="stylesheet" />
 <script src="{% static 'plugins/gridzy/gridzy.js' %}"></script>
 <script src="https://d3js.org/d3.v4.min.js"></script>
+<script src="https://cdn.rawgit.com/eligrey/canvas-toBlob.js/f1a01896135ab378aa5c0118eadd81da55e698d8/canvas-toBlob.js"></script>
+<script src="https://cdn.rawgit.com/eligrey/FileSaver.js/e9d941381475b5df8b7d7691013401e171014e89/FileSaver.min.js"></script>
 {% endblock custom_js_css_link %}
 
 {% block breadcrumb_title %}
@@ -990,9 +992,23 @@ Scan Findings
 				<div class="card">
 					<div class="card-body">
 						<div id="visualisation-filter" style="display:none">
-							<div class="form-check form-switch">
-								<input type="checkbox" class="form-check-input" name="expand-nodes-checkbox">
-								<label class="form-check-label" for="customSwitch1">&nbsp; &nbsp; Expand all nodes</label>
+							<div class="row">
+								<div class="col-6">
+									<div class="form-check form-switch">
+										<input type="checkbox" class="form-check-input" name="expand-nodes-checkbox">
+										<label class="form-check-label" for="customSwitch1">&nbsp; &nbsp; Expand all nodes</label>
+									</div>
+								</div>
+								<div class="col-6">
+									<div class="dropdown float-end">
+										<a href="#" class="dropdown-toggle arrow-none card-drop" data-bs-toggle="dropdown" aria-expanded="false">
+											<i class="mdi mdi-dots-vertical"></i>
+										</a>
+										<div class="dropdown-menu dropdown-menu-end" style="">
+											<a href="#" class="dropdown-item" id="saveButton">Download Visualization as png</a>
+										</div>
+									</div>
+								</div>
 							</div>
 						</div>
 						<hr>
diff --git a/web/static/assets/css/app-dark.min.css b/web/static/assets/css/app-dark.min.css
index fef56451..2edf1222 100644
--- a/web/static/assets/css/app-dark.min.css
+++ b/web/static/assets/css/app-dark.min.css
@@ -10,3 +10,7 @@
 .page-item.disabled .page-link{
   background: #263238 !important;
 }
+
+#visualisation-svg{
+  background-color: #36404A !important;
+}
diff --git a/web/static/custom/custom.css b/web/static/custom/custom.css
index e6a68139..6a2235b1 100644
--- a/web/static/custom/custom.css
+++ b/web/static/custom/custom.css
@@ -289,7 +289,7 @@ div.dataTables_wrapper div.dataTables_processing {
 }
 
 .node text {
-  font: 12px sans-serif;
+  font: 16px sans-serif;
 }
 
 .link {
@@ -297,3 +297,33 @@ div.dataTables_wrapper div.dataTables_processing {
   stroke: #ccc;
   stroke-width: 1px;
 }
+
+
+#visualisation-svg{
+  background-color: #fff;
+}
+
+
+.svg-text-primary {
+  fill: #0d6efd;
+}
+
+.svg-text-success {
+  fill: #198754;
+}
+
+.svg-text-danger {
+  fill: #dc3545;
+}
+
+.svg-text-critical {
+  fill: #dc3545;
+}
+
+.svg-text-warning {
+  fill: #ffc107;
+}
+
+.svg-text-low {
+  fill: #ffc107;
+}
diff --git a/web/static/custom/mitch.js b/web/static/custom/mitch.js
index 1a8fb158..67d28044 100644
--- a/web/static/custom/mitch.js
+++ b/web/static/custom/mitch.js
@@ -30,11 +30,11 @@ function visualise_scan_results(scan_id)
     // appends a 'group' element to 'svg'
     // moves the 'group' element to the top left margin
     var svg = d3.select("#visualisation").append("svg")
-    .attr("width", width + margin.right + margin.left)
-    .attr("height", height + margin.top + margin.bottom)
+    .attr("id", "visualisation-svg")
+    .attr("width", width)
+    .attr("height", height)
     .call(d3.zoom().on("zoom", function () {svg.attr("transform", d3.event.transform)}).scaleExtent([0.2, 3]))
-    .append("g")
-    .attr("transform", "translate("+ margin.left + "," + -1 * (height/2-250) + ")");
+    .append("g").attr("transform", "translate("+ margin.left + "," + -1 * (height/2-250) + ")");
 
     var i = 0,
     duration = 500,
@@ -121,8 +121,34 @@ function visualise_scan_results(scan_id)
       .attr("text-anchor", function(d) {
         return d.children || d._children ? "end" : "start";
       })
-      .attr('class', 'nodeText')
-      .text(function(d) { return d.data.description; });
+      .attr('class', function(data){
+        text_class = "";
+        if (data.data.http_status >= 400 || data.data.is_uncommon || data.data.description == 'High' || data.data.title == 'Exposed Creds') {
+          text_class = " svg-text-danger "
+        }
+        else if (data.data.http_status > 200 && data.data.http_status < 400 || data.data.description == 'Medium') {
+          text_class = " svg-text-warning "
+        }
+        else if (data.data.http_status == 200) {
+          text_class = " svg-text-success "
+        }
+        else if (data.data.description == 'Critical'){
+          text_class = " svg-text-critical "
+        }
+        else if (data.data.description == 'Low'){
+          text_class = " svg-text-low "
+        }
+        else if (data.data.description == 'Informational'){
+          text_class = " svg-text-primary "
+        }
+        return 'textNode' + text_class;
+      })
+      .text(function(d) {
+        if (d.data.title == 'Interesting') {
+          return '(Interesting) ' + d.data.description;
+        }
+        return d.data.description;
+      });
 
       // UPDATE
       var nodeUpdate = nodeEnter.merge(node);
@@ -232,6 +258,95 @@ function visualise_scan_results(scan_id)
       update(root);
     });
 
+    // source: view-source:https://www.demo2s.com/javascript/javascript-d3.js-save-svg-to-png-image-demo-b1a10.htm
+
+    d3.select('#saveButton').on('click', function(){
+      // get domain name
+      var domain = $("[aria-current]").text();
+      var svgString = getSVGString(d3.select('#visualisation-svg').node());
+      svgString2Image( svgString, 2*width, 2*height, 'png', save ); // passes Blob and filesize String to the callback
+      function save( dataBlob, filesize ){
+        saveAs( dataBlob, domain + '_visualisation.png' ); // FileSaver.js function
+      }
+    });
+
+    // Below are the functions that handle actual exporting:
+    // getSVGString ( svgNode ) and svgString2Image( svgString, width, height, format, callback )
+    function getSVGString( svgNode ) {
+      svgNode.setAttribute('xlink', 'http://www.w3.org/1999/xlink');
+      var cssStyleText = getCSSStyles( svgNode );
+      appendCSS( cssStyleText, svgNode );
+      var serializer = new XMLSerializer();
+      var svgString = serializer.serializeToString(svgNode);
+      svgString = svgString.replace(/(\w+)?:?xlink=/g, 'xmlns:xlink='); // Fix root xlink without namespace
+      svgString = svgString.replace(/NS\d+:href/g, 'xlink:href'); // Safari NS namespace fix
+      return svgString;
+      function getCSSStyles( parentElement ) {
+        var selectorTextArr = [];
+        // Add Parent element Id and Classes to the list
+        selectorTextArr.push( '#'+parentElement.id );
+        for (var c = 0; c < parentElement.classList.length; c++)
+        if ( !contains('.'+parentElement.classList[c], selectorTextArr) )
+        selectorTextArr.push( '.'+parentElement.classList[c] );
+        // Add Children element Ids and Classes to the list
+        var nodes = parentElement.getElementsByTagName("*");
+        for (var i = 0; i < nodes.length; i++) {
+          var id = nodes[i].id;
+          if ( !contains('#'+id, selectorTextArr) )
+          selectorTextArr.push( '#'+id );
+          var classes = nodes[i].classList;
+          for (var c = 0; c < classes.length; c++)
+          if ( !contains('.'+classes[c], selectorTextArr) )
+          selectorTextArr.push( '.'+classes[c] );
+        }
+        // Extract CSS Rules
+        var extractedCSSText = "";
+        for (var i = 0; i < document.styleSheets.length; i++) {
+          var s = document.styleSheets[i];
+          try {
+            if(!s.cssRules) continue;
+          } catch( e ) {
+            if(e.name !== 'SecurityError') throw e; // for Firefox
+            continue;
+          }
+          var cssRules = s.cssRules;
+          for (var r = 0; r < cssRules.length; r++) {
+            if ( contains( cssRules[r].selectorText, selectorTextArr ) )
+            extractedCSSText += cssRules[r].cssText;
+          }
+        }
+        return extractedCSSText;
+        function contains(str,arr) {
+          return arr.indexOf( str ) === -1 ? false : true;
+        }
+      }
+      function appendCSS( cssText, element ) {
+        var styleElement = document.createElement("style");
+        styleElement.setAttribute("type","text/css");
+        styleElement.innerHTML = cssText;
+        var refNode = element.hasChildNodes() ? element.children[0] : null;
+        element.insertBefore( styleElement, refNode );
+      }
+    }
+    function svgString2Image( svgString, width, height, format, callback ) {
+      var format = format ? format : 'png';
+      var imgsrc = 'data:image/svg+xml;base64,'+ btoa( unescape( encodeURIComponent( svgString ) ) ); // Convert SVG string to data URL
+      var canvas = document.createElement("canvas");
+      var context = canvas.getContext("2d");
+      canvas.width = width;
+      canvas.height = height;
+      var image = new Image();
+      image.onload = function() {
+        context.clearRect ( 0, 0, width, height );
+        context.drawImage(image, 0, 0, width, height);
+        canvas.toBlob( function(blob) {
+          var filesize = Math.round( blob.length/1024 ) + ' KB';
+          if ( callback ) callback( blob, filesize );
+        });
+      };
+      image.src = imgsrc;
+    }
+
   }).fail(function(){
     $('#visualisation-loader').empty();
     $("#visualisation-loader").append(`<h5 class="text-danger">Sorry, could not visualize.</h5>`);