VulnerabilityGroup.java

package com.denimgroup.threadfix.webapp.controller;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import com.denimgroup.threadfix.data.entities.GenericSeverity;
import com.denimgroup.threadfix.data.entities.Vulnerability;

/**
 * 
 * @author mcollins
 *
 */
public final class VulnerabilityGroup {
	
	// This class is private so it can only be generated through generateGroups.
	// It's also intentionally immutable, which is why there are no setters.
	private VulnerabilityGroup(List<Vulnerability> vulns, String name, boolean showHeader) {
		this.vulnerabilities = vulns;
		this.name = name;
		this.showHeader = showHeader;
	}
	
	private List<Vulnerability> vulnerabilities = null;
	private String name = null;
	private boolean showHeader = false;
	
	public List<Vulnerability> getVulnerabilities() {
		return vulnerabilities;
	}

	public String getName() {
		return name;
	}

	public boolean isShowHeader() {
		return showHeader;
	}
	
	public String getColorClass() {
		switch (name) {
			case GenericSeverity.CRITICAL: return "error";
			case GenericSeverity.HIGH: return "warning";
			case GenericSeverity.MEDIUM: return "success";
			default: return "info";
		}
	}

	public static List<VulnerabilityGroup> generateGroups(List<Vulnerability> vulns, boolean shouldSort) {
		if (shouldSort && vulns != null && !vulns.isEmpty()) {
			List<Vulnerability> infoVulnerabilities = new ArrayList<Vulnerability>(),
					lowVulnerabilities = new ArrayList<Vulnerability>(),
					mediumVulnerabilities = new ArrayList<Vulnerability>(),
					highVulnerabilities = new ArrayList<Vulnerability>(),
					criticalVulnerabilities = new ArrayList<Vulnerability>();

			for (Vulnerability vuln : vulns) {
				if (vuln == null || vuln.getGenericSeverity() == null ||
						vuln.getGenericSeverity().getName() == null) {
					continue;
				}
				
				switch (vuln.getGenericSeverity().getName()) {
					case GenericSeverity.INFO: infoVulnerabilities.add(vuln); break;
					case GenericSeverity.LOW: lowVulnerabilities.add(vuln); break;
					case GenericSeverity.MEDIUM: mediumVulnerabilities.add(vuln); break;
					case GenericSeverity.HIGH: highVulnerabilities.add(vuln); break;
					case GenericSeverity.CRITICAL: criticalVulnerabilities.add(vuln); break;
				}
			}
			
			return Arrays.asList(new VulnerabilityGroup[] {
				new VulnerabilityGroup(criticalVulnerabilities, "Critical", true),
				new VulnerabilityGroup(highVulnerabilities, "High", true),
				new VulnerabilityGroup(mediumVulnerabilities, "Medium", true),
				new VulnerabilityGroup(lowVulnerabilities, "Low", true),
				new VulnerabilityGroup(infoVulnerabilities, "Info", true)
			});
		} else {
			return Arrays.asList(new VulnerabilityGroup[] {
				new VulnerabilityGroup(vulns, "All", false)
			});
		}
	}
}