adding additional features to the bar charts

Author: summitt

pom.xml

@@ -4,7 +4,7 @@
 	<modelVersion>4.0.0</modelVersion>
 	<groupId>com.factionsecurity</groupId>
 	<artifactId>faction-vulnerability-bar-chart</artifactId>
-	<version>1.0</version>
+	<version>1.1</version>
 	<name>Faction Vulnerability Bar Chart</name>
 	<dependencies>
 		<dependency>
@@ -17,6 +17,18 @@
 			<artifactId>jfreechart</artifactId>
 			<version>1.5.4</version>
 		</dependency>
+		 <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-api</artifactId>
+            <version>5.9.1</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-engine</artifactId>
+            <version>5.9.1</version>
+            <scope>test</scope>
+        </dependency>
 	</dependencies>
 	<build>
 		<plugins>

src/main/java/com/faction/VulnerabilityBarChart.java

@@ -1,9 +1,12 @@
 package com.faction;
 
 import java.util.ArrayList;
+
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -16,27 +19,54 @@
 import org.jfree.chart.ChartFactory;
 import org.jfree.chart.ChartUtils;
 import org.jfree.chart.JFreeChart;
+import org.jfree.chart.axis.CategoryAxis;
 import org.jfree.chart.axis.NumberAxis;
 import org.jfree.chart.axis.NumberTickUnit;
+import org.jfree.chart.axis.ValueAxis;
+import org.jfree.chart.labels.ItemLabelAnchor;
+import org.jfree.chart.labels.ItemLabelPosition;
+import org.jfree.chart.labels.StandardCategoryItemLabelGenerator;
 import org.jfree.chart.plot.CategoryPlot;
 import org.jfree.chart.plot.PlotOrientation;
 import org.jfree.chart.renderer.category.BarRenderer;
 import org.jfree.chart.renderer.category.StandardBarPainter;
+import org.jfree.chart.ui.TextAnchor;
 import org.jfree.chart.renderer.category.CategoryItemRenderer;
+import org.jfree.chart.renderer.category.CategoryItemRendererState;
 import org.jfree.data.category.CategoryDataset;
 import org.jfree.data.category.DefaultCategoryDataset;
 import org.sonatype.plexus.components.cipher.Base64;
 
+import java.awt.BasicStroke;
 import java.awt.Color;
 import java.awt.Font;
+import java.awt.FontMetrics;
+import java.awt.Graphics2D;
 import java.awt.Paint;
+import java.awt.geom.Rectangle2D;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 
 public class VulnerabilityBarChart extends BaseExtension implements com.faction.extender.ReportManager {
 
 	private int totalLevels = 6;
 	private LinkedHashMap<String,Color> severities = new LinkedHashMap<>();
+	private HashMap<Long, SevObj> sevMap = new HashMap<>();
+	private Boolean showNumbers = false;
+	private Boolean showBorders = false;
+	private String yAxisName="";
+	
+	
+	private class SevObj {
+		public Color  color;
+		public Color  border;
+		public String  name;
+		public SevObj(String name, Color color, Color border) {
+			this.color = color;
+			this.border = border;
+			this.name = name;
+		}
+	}
 
 
 	@Override
@@ -55,31 +85,70 @@ public String reportCreate(Assessment asmt, List<Vulnerability> vulns, String re
 	 * This section represents the getters from the user supplied configs
 	 * 
 	 */
+	private Color getColorOrDefault(String value) {
+		if(value == null) {
+			return new Color(0x000000);
+		}else {
+			try {
+				return new Color(Integer.parseInt(value.trim(), 16));
+			}catch(Exception e) {
+				return new Color(0x000000);
+			}
+		}
+	}
 	private void setUpSeverities() {
 		String colorsCSV = this.getConfigs().get("Severity Colors");
+		String sevMapCSV = this.getConfigs().get("Severity Mapping");
+		String bordersCSV = this.getConfigs().get("Border Colors");
+		showNumbers = this.getConfigs().get("Show Numbers") == null? false : Boolean.parseBoolean(this.getConfigs().get("Show Numbers").trim());
+		showBorders = this.getConfigs().get("Show Borders") == null? false : Boolean.parseBoolean(this.getConfigs().get("Show Borders").trim());
+		yAxisName = this.getConfigs().get("Y Axis Label") == null? "" : this.getConfigs().get("Y Axis Label");
 		String w = this.getConfigs().get("Width");
 		if(colorsCSV == null || colorsCSV.trim().equals("")) {
-			severities.put("Critical", new Color(0xdd4b39));
-			severities.put("High", new Color(0xf39c12));
-			severities.put("Medium", new Color(0x00c0ef));
-			severities.put("Low", new Color(0x39cccc));
-			severities.put("Recommended", new Color(0x00a65a));
-			severities.put("Informational", new Color(0x95a5a6));
+			this.sevMap.put(5l, new SevObj("Critical", new Color(0xdd4b39), new Color(0xdd4b39)));
+			this.sevMap.put(4l, new SevObj("High", new Color(0xf39c12), new Color(0xf39c12)));
+			this.sevMap.put(3l, new SevObj("Medium", new Color(0x00c0ef), new Color(0x00c0ef)));
+			this.sevMap.put(2l, new SevObj("Low", new Color(0x39cccc), new Color(0x39cccc)));
+			this.sevMap.put(1l, new SevObj("Recommended", new Color(0x00a65a), new Color(0x00a65a) ));
+			this.sevMap.put(0l, new SevObj("Informational", new Color(0x95a5a6), new Color(0x95a5a6)));
 		}else {
 			String [] colors = colorsCSV.split(",");
+			String [] sevIds = sevMapCSV.split(",");
+			String [] borderColors = bordersCSV.split(",");
 			this.severities = Arrays
 					.stream(colors)
 					.map( color -> color.split(":"))
 					.collect(Collectors.toMap(
 							keyValue -> keyValue[0],
-							keyValue -> new Color(Integer.parseInt(keyValue[1].trim(), 16)),
+							keyValue -> this.getColorOrDefault(keyValue[1]),
+							(v1,v2) -> v1,
+							LinkedHashMap<String, Color>::new
+					));
+			LinkedHashMap<String,Color> borders = Arrays
+					.stream(borderColors)
+					.map( color -> color.split(":"))
+					.collect(Collectors.toMap(
+							keyValue -> keyValue[0],
+							keyValue -> this.getColorOrDefault(keyValue[1]),
 							(v1,v2) -> v1,
 							LinkedHashMap<String, Color>::new
 					));
-			this.totalLevels = colors.length;
 
-			String [] sevs = this.severities.keySet().toArray(new String[0]);
+			this.sevMap = Arrays
+					.stream(sevIds)
+					.map( sevs -> sevs.split(":"))
+					.collect(Collectors.toMap(
+							keyValue -> Long.parseLong(keyValue[1].trim()),
+							keyValue -> new SevObj(
+									keyValue[0].trim(),
+									this.severities.get(keyValue[0].trim()),
+									borders.get(keyValue[0].trim())
+									),
+							(v1,v2) -> v1,
+							HashMap<Long, SevObj>::new
+					));
 
+			this.totalLevels = colors.length;
 		}
 
 	}
@@ -106,6 +175,14 @@ private Paint[] getColors() {
 			return new Paint[] {new Color(Integer.parseInt(defaultColor.trim(), 16)) };
 		}
 	}
+	private Paint getDefaultLabelColor() {
+		String defaultColor = this.getConfigs().get("Default Label Color");
+		if(defaultColor == null || defaultColor.trim().equals("")) {
+			return new Color(0xFFFFFF) ; // Default Color
+		}else {
+			return new Color(Integer.parseInt(defaultColor.trim(), 16));
+		}
+	}
 
 	private int [] getDimensions() {
 		int [] dimensions = {900, 600};
@@ -136,7 +213,7 @@ private String vulnSeverityChartImage(List<Vulnerability> vulns, String reportTe
 			chart.getPlot().setBackgroundPaint( new Color(0,0,0,0));
 			chart.getPlot().setBackgroundImageAlpha(0);
 			// set bar colors
-			final CategoryItemRenderer renderer = new CustomRenderer(this.getColors(), this.getDefaultColor());
+			final CategoryItemRenderer renderer = new CustomRenderer(this.sevMap, this.showBorders, this.showNumbers, this.getDefaultColor()[0], this.getDefaultLabelColor());
 			final CategoryPlot plot = chart.getCategoryPlot();
 			//Match Font
 			Font font = new Font("Arial", Font.PLAIN, 25);
@@ -145,6 +222,14 @@ private String vulnSeverityChartImage(List<Vulnerability> vulns, String reportTe
 			// remove bar gradient and shaddow
 			((BarRenderer) renderer).setBarPainter(new StandardBarPainter());
 			((BarRenderer) renderer).setShadowVisible(false);
+			((BarRenderer) renderer).setDefaultItemLabelGenerator(new StandardCategoryItemLabelGenerator());
+			if(this.showNumbers) {
+			((BarRenderer) renderer).setDefaultItemLabelsVisible(true);
+				ItemLabelPosition position = new ItemLabelPosition(ItemLabelAnchor.CENTER, 
+						TextAnchor.CENTER, TextAnchor.CENTER, 0.0);
+				renderer.setDefaultPositiveItemLabelPosition(position);
+				renderer.setDefaultItemLabelPaint(this.getDefaultLabelColor());
+			}
 			plot.setRenderer(renderer);
 			// Show mininal gridlines
 			plot.setRangeGridlinesVisible(false);
@@ -164,42 +249,133 @@ private String vulnSeverityChartImage(List<Vulnerability> vulns, String reportTe
 			return image;
 	}
 
+	private List<String>getSevNames() {
+		return this.sevMap.keySet().stream()
+					.map( id -> this.sevMap.get(id).name)
+					.collect(Collectors.toList());
+	}
+	static Long [] getSevIds(HashMap<Long,SevObj> sevMap) {
+		Long [] ids = (Long[]) sevMap.keySet().toArray(new Long[0]);
+		Arrays.sort(ids, Collections.reverseOrder());
+		return ids;
+	}
+	
 
 	private CategoryDataset createDataset(List<Vulnerability> vulns) {
 
 		DefaultCategoryDataset dataset = new DefaultCategoryDataset();
-		String [] severityNames = (String[]) this.severities.keySet().toArray(new String[0]);
-		for(int level=totalLevels-1; level >=0; level-- ) {
-			final int index = level;
-			long count = vulns.stream().filter(v -> v.getOverall().equals(new Long(index))).count();
-			dataset.setValue(count, "Severity", severityNames[this.totalLevels - index -1]);
-		}
+		Arrays.asList(getSevIds(this.sevMap)).forEach(id -> {
+			long count = vulns.stream().filter(v -> v.getOverall().equals(id)).count();
+			dataset.setValue(count, this.yAxisName, this.sevMap.get(id).name);
+		});
+		
 		return dataset;
 	}
-
 	private JFreeChart createChart(CategoryDataset dataset) {
-
-		JFreeChart barChart = ChartFactory.createBarChart("", "", "Severity", dataset, PlotOrientation.VERTICAL, false,
+		String yAxisName = this.yAxisName;
+		JFreeChart barChart = ChartFactory.createBarChart("", "", yAxisName, dataset, PlotOrientation.VERTICAL, false,
 				true, false);
 		return barChart;
 	}
 
 	class CustomRenderer extends BarRenderer {
 
-		private Paint[] colors;
-		private Paint[] defaultColor;
+		private Paint defaultBarColor;
+		private Paint labelColor;
+		private HashMap<Long, SevObj> colorMap;
+		private Boolean showBorders;
+		private Boolean showNumbers;
 
-		public CustomRenderer(final Paint[] colors, final Paint[] defaultColor ) {
-			this.colors = colors;
+		public CustomRenderer(final HashMap<Long,SevObj> colorMap, Boolean showBorders, Boolean showNumbers, Paint defaultBarColor, Paint labelColor ) {
+			this.colorMap = colorMap;
+			this.showBorders = showBorders;
+			this.defaultBarColor = defaultBarColor;
+			this.labelColor = labelColor;
+			this.showNumbers = showNumbers;
 		}
-
+		@Override
 		public Paint getItemPaint(final int row, final int column) {
-			if (column >= this.colors.length) {
-				Paint[] defaultPaint = this.defaultColor;
-				return defaultPaint[0];
+			if (column >= this.colorMap.keySet().size()) {
+				Paint defaultPaint = this.defaultBarColor;
+				return defaultPaint;
 			}
-			return this.colors[column % this.colors.length];
+			Long [] sevIds = VulnerabilityBarChart.getSevIds(colorMap);			
+			Long id = sevIds[column];
+			
+			
+			return sevMap.get(id).color;
 		}
+		@Override
+		public void drawItem(Graphics2D g2, 
+		                     CategoryItemRendererState state, 
+		                     Rectangle2D dataArea, 
+		                     CategoryPlot plot, 
+		                     CategoryAxis domainAxis, 
+		                     ValueAxis rangeAxis, 
+		                     CategoryDataset dataset, 
+		                     int row, 
+		                     int column, 
+		                     int pass) {
+
+		    // Only draw bars in pass 1
+		    if (pass != 0) {
+		        return;
+		    }
+
+		    // Determine the bar's position and dimensions
+		    double value = dataset.getValue(row, column).doubleValue();
+		    double barStart = rangeAxis.valueToJava2D(
+		            rangeAxis.getLowerBound(), dataArea, plot.getRangeAxisEdge());
+		    double barEnd = rangeAxis.valueToJava2D(value, dataArea, plot.getRangeAxisEdge());
+		    double barWidth = state.getBarWidth();
+
+		    // Calculate the bar's rectangle
+		    double x = domainAxis.getCategoryStart(column, getColumnCount(), dataArea, plot.getDomainAxisEdge())
+		            + state.getBarWidth() * row;
+		    double y = Math.min(barStart, barEnd);
+		    double height = Math.abs(barEnd - barStart);
+
+		    Rectangle2D bar = new Rectangle2D.Double(x, y, barWidth, height);
+
+		    // Draw the bar (default rendering)
+		    g2.setPaint(getItemPaint(row, column));
+		    g2.fill(bar);
+
+		    // Draw the border
+			if(this.showBorders) {
+				Long [] sevIds = VulnerabilityBarChart.getSevIds(colorMap);			
+				Long id = sevIds[column];
+				Color c = sevMap.get(id).border;
+				System.out.println(c);
+				g2.setPaint(c);
+				g2.setStroke(new BasicStroke(1.0f)); // Adjust stroke as needed
+			}else {
+				g2.setPaint(Color.BLACK); // Default border color
+				g2.setStroke(new BasicStroke(0f)); // Adjust stroke as needed
+			}
+		    g2.draw(bar);
+		    
+		    if(this.showNumbers) {
+			 // Draw the item label (bar value) in the center of the bar
+				int fontSize = 12; // Set your desired font size
+				Font originalFont = g2.getFont(); // Save the original font
+				Font newFont = originalFont.deriveFont((float) fontSize);
+				g2.setFont(newFont);
+				String label = String.valueOf( (int) value); // Convert the value to a string
+				FontMetrics metrics = g2.getFontMetrics();
+				double labelWidth = metrics.stringWidth(label);
+				double labelHeight = metrics.getHeight();
+
+				// Calculate label position (center of the bar)
+				double labelX = bar.getCenterX() - (labelWidth / 2);
+				double labelY = bar.getCenterY() + (labelHeight / 4); // Adjust for vertical alignment
+
+				g2.setPaint(this.labelColor); // Set label color
+				g2.drawString(label, (float) labelX, (float) labelY);
+				g2.setFont(originalFont);
+		    }
+		}
+		
 	}