Skip to content

Commit

Permalink
Merge pull request #10 from rusakovichma/data-flow-diagram-generation
Browse files Browse the repository at this point in the history
Graphviz DFD rendering
  • Loading branch information
rusakovichma authored May 22, 2023
2 parents 1a02881 + 91972a8 commit 903ec1c
Show file tree
Hide file tree
Showing 13 changed files with 346 additions and 12 deletions.
8 changes: 7 additions & 1 deletion pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

<groupId>com.github.rusakovichma</groupId>
<artifactId>tictaac</artifactId>
<version>1.2.3</version>
<version>1.3.0</version>

<name>TicTaaC</name>
<url>https://github.com/rusakovichma/TicTaaC.git</url>
Expand Down Expand Up @@ -487,6 +487,12 @@
<version>${xchart.version}</version>
</dependency>

<dependency>
<groupId>guru.nidi</groupId>
<artifactId>graphviz-java</artifactId>
<version>0.18.1</version>
</dependency>

<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter-engine</artifactId>
Expand Down
6 changes: 5 additions & 1 deletion src/main/java/com/github/rusakovichma/tictaac/Launcher.java
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import com.github.rusakovichma.tictaac.engine.StandardThreatEngine;
import com.github.rusakovichma.tictaac.engine.ThreatEngine;
import com.github.rusakovichma.tictaac.model.Threat;
import com.github.rusakovichma.tictaac.model.ThreatModel;
import com.github.rusakovichma.tictaac.model.ThreatRisk;
import com.github.rusakovichma.tictaac.model.ThreatsCollection;
import com.github.rusakovichma.tictaac.model.exception.QualityGateFailed;
Expand Down Expand Up @@ -193,7 +194,9 @@ public static void main(String[] args) {
ThreatEngine threatEngine = getThreatEngine(rulesProvider, mitigator);

ThreatModelProvider threatModelProvider = getThreatModel(singleValueParams);
ThreatsCollection threats = threatEngine.generateThreats(threatModelProvider.getModel());

ThreatModel threatModel = threatModelProvider.getModel();
ThreatsCollection threats = threatEngine.generateThreats(threatModel);

modelThreats.put(threatModelsParam, threats.getThreats());

Expand All @@ -202,6 +205,7 @@ public static void main(String[] args) {
new ReportHeader(
threats.getName(), threats.getVersion(), new Date()
),
threatModel,
threats.getThreats()
);
} catch (IOException ex) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,14 @@
import com.github.rusakovichma.tictaac.model.OwaspCategory;
import com.github.rusakovichma.tictaac.model.Threat;
import com.github.rusakovichma.tictaac.model.ThreatCategory;
import com.github.rusakovichma.tictaac.model.ThreatModel;
import com.github.rusakovichma.tictaac.model.mitigation.MitigationStatus;
import com.github.rusakovichma.tictaac.model.threatmodel.boundary.BoundaryCategory;
import com.github.rusakovichma.tictaac.reporter.analytics.ThreatAnalytics;
import com.github.rusakovichma.tictaac.reporter.chart.ChartPlotter;
import com.github.rusakovichma.tictaac.reporter.chart.XChartPlotter;
import com.github.rusakovichma.tictaac.reporter.dfd.DataFlowRender;
import com.github.rusakovichma.tictaac.reporter.dfd.GraphvizDataFlowRender;
import com.github.rusakovichma.tictaac.util.ResourceUtil;

import java.io.*;
Expand All @@ -39,6 +42,8 @@ public class StreamThreatsReporter implements ThreatsReporter {
private final OutputStream outputStream;
private final ReportFormat reportFormat;

private DataFlowRender dataFlowRender = new GraphvizDataFlowRender();

private String headerTemplate;
private String entryTemplate;
private String entrySeparator;
Expand Down Expand Up @@ -151,19 +156,22 @@ private String[] getCharts(Collection<Threat> threats) {
return charts.toArray(new String[charts.size()]);
}

private void writeModel(ReportHeader header, Collection<Threat> threats)
private void writeModel(ReportHeader header, ThreatModel threatModel, Collection<Threat> threats)
throws IOException {
if (reportFormat == ReportFormat.html) {
String[] charts = getCharts(threats);
headerTemplate = headerTemplate.replace("%owasp-chart%", charts[0])
.replace("%stride-chart%", charts[1])
.replace("%vectors-chart%", charts[2])
.replace("%mitigations-chart%", charts[3]);
.replace("%mitigations-chart%", charts[3])
.replace("%data-flow-diagram%",
Base64.getEncoder().encodeToString(
dataFlowRender.createDataFlow(threatModel)));
}

outputStream.write(
String.format(headerTemplate, header.getName(), header.getVersion(),
new SimpleDateFormat("yyyy-MM-dd HH:mm").format(new Date()))
new SimpleDateFormat("yyyy-MM-dd HH:mm").format(new Date()))
.getBytes()
);

Expand Down Expand Up @@ -193,8 +201,8 @@ private void writeModel(ReportHeader header, Collection<Threat> threats)
}

@Override
public void publish(ReportHeader header, Collection<Threat> threats)
public void publish(ReportHeader header, ThreatModel threatModel, Collection<Threat> threats)
throws IOException {
writeModel(header, threats);
writeModel(header, threatModel, threats);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,14 @@
package com.github.rusakovichma.tictaac.reporter;

import com.github.rusakovichma.tictaac.model.Threat;
import com.github.rusakovichma.tictaac.model.ThreatModel;

import java.io.IOException;
import java.util.Collection;

public interface ThreatsReporter {

public void publish(ReportHeader header, Collection<Threat> threats)
public void publish(ReportHeader header, ThreatModel threatModel, Collection<Threat> threats)
throws IOException;

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
/*
* This file is part of TicTaaC.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
* Copyright (c) 2022 Mikhail Rusakovich. All Rights Reserved.
*/
package com.github.rusakovichma.tictaac.reporter.dfd;

import com.github.rusakovichma.tictaac.model.ThreatModel;

import java.io.IOException;

public interface DataFlowRender {

public byte[] createDataFlow(ThreatModel threatModel) throws IOException;

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
/*
* This file is part of TicTaaC.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
* Copyright (c) 2022 Mikhail Rusakovich. All Rights Reserved.
*/
package com.github.rusakovichma.tictaac.reporter.dfd;

import com.github.rusakovichma.tictaac.model.ThreatModel;
import com.github.rusakovichma.tictaac.model.threatmodel.Boundary;
import com.github.rusakovichma.tictaac.model.threatmodel.DataFlow;
import com.github.rusakovichma.tictaac.model.threatmodel.Element;
import com.github.rusakovichma.tictaac.model.threatmodel.boundary.BoundaryCategory;
import com.github.rusakovichma.tictaac.util.ImageUtils;
import guru.nidi.graphviz.attribute.*;
import guru.nidi.graphviz.engine.Format;
import guru.nidi.graphviz.engine.Graphviz;
import guru.nidi.graphviz.model.Graph;
import guru.nidi.graphviz.model.LinkSource;
import guru.nidi.graphviz.model.Node;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import static guru.nidi.graphviz.attribute.Rank.RankDir.LEFT_TO_RIGHT;
import static guru.nidi.graphviz.model.Factory.graph;
import static guru.nidi.graphviz.model.Factory.node;

public class GraphvizDataFlowRender implements DataFlowRender {

private LinkSource getNodeFromElement(Element element) {
switch (element.getType()) {
case interactor:
case externalService:
case internalService:
return node(element.getId()).with(Label.of(element.getName()))
.with(Shape.BOX);
case proxyServer:
case webServer:
case process:
return node(element.getId()).with(Shape.M_RECORD,
Label.of(String.format("{<f0> |<f1> %s\n\n\n}", element.getName())));
case database:
return node(element.getId()).with(Shape.M_RECORD,
Label.of(String.format("<f0> |<f1> %s\n\n\n", element.getName()))
);
default:
throw new RuntimeException("Unsupported entity");
}
}

private Graph getGraphFromBoundary(Boundary boundary) {
Graph graph = graph(boundary.getId()).cluster();

if (boundary.getCategory() == BoundaryCategory.globalNetwork) {
graph = graph.graphAttr().with(Color.WHITE, Rank.newRank().noCluster());
} else {
graph = graph.graphAttr().with(Color.BLUE, Label.of(boundary.getId()));
}

LinkSource[] boundaryLinkSources = new LinkSource[boundary.getElements().size()];
for (int i = 0; i < boundary.getElements().size(); i++) {
boundaryLinkSources[i] = getNodeFromElement(boundary.getElements().get(i));
}

graph = graph.with(boundaryLinkSources);

return graph;
}

@Override
public byte[] createDataFlow(ThreatModel threatModel) throws IOException {
if (threatModel == null) {
return new byte[0];
}

Graph dataFlowGraph = graph(threatModel.getName()).directed()
.graphAttr().with(Rank.dir(LEFT_TO_RIGHT));

List<LinkSource> dataFlowObjects = new ArrayList<>();
for (Boundary boundary : threatModel.getBoundaries()) {
dataFlowObjects.add(getGraphFromBoundary(boundary));
}

//if some elements are not within any boundary
for (Element element : threatModel.getElements()) {
if (!threatModel.getBoundaries().stream()
.anyMatch(boundary -> boundary.getElements().contains(element))) {
dataFlowObjects.add(getNodeFromElement(element));
}
}

for (DataFlow dataFlow : threatModel.getDataFlows()) {
Node node = node(dataFlow.getSource().getId())
.link(dataFlow.getTarget().getId());
if (dataFlow.getTitle() != null) {
node = node.with(Label.of(dataFlow.getTitle()));
}

dataFlowObjects.add(node);
}

dataFlowGraph = dataFlowGraph.with(dataFlowObjects);

return ImageUtils.toByteArray(
Graphviz.fromGraph(dataFlowGraph).render(Format.PNG).toImage(),
"png");
}
}
45 changes: 45 additions & 0 deletions src/main/java/com/github/rusakovichma/tictaac/util/ImageUtils.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
package com.github.rusakovichma.tictaac.util;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;

/*
* This file is part of TicTaaC.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
* Copyright (c) 2022 Mikhail Rusakovich. All Rights Reserved.
*/
public class ImageUtils {

public static byte[] toByteArray(BufferedImage bi, String format) throws IOException {
ByteArrayOutputStream baos = new ByteArrayOutputStream();
ImageIO.write(bi, format, baos);
byte[] bytes = baos.toByteArray();
return bytes;

}

public static BufferedImage toBufferedImage(byte[] bytes) throws IOException {
InputStream is = new ByteArrayInputStream(bytes);
BufferedImage bi = ImageIO.read(is);
return bi;
}

public static void saveToFile(byte[] image, String format, String path) throws IOException {
ByteArrayInputStream bais = new ByteArrayInputStream(image);
final BufferedImage bufferedImage = ImageIO.read(bais);
ImageIO.write(bufferedImage, format, new File(path));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@
<p>The threat modeling process should be driven by application designers, however, developers and testers should be involved.</p>
<p><strong>Results of this threat modeling exercise are prioritized by risk level list of identified threats and list of proposed mitigations for those threats.</strong></p>
<p>&nbsp;</p>
<h2>Data Flow Diagram</h2>
<img src="data:image/png;base64, %data-flow-diagram%" alt="Data Flow Diagram" />
<p>&nbsp;</p>
<h2>Threat Analysis</h2>
<table cellspacing="0" cellpadding="0">
<thead>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,6 @@ void accept() {
File dir = new File("src/test/resources");
File[] modelFiles = dir.listFiles(new ThreatModelFilter());

assertTrue(modelFiles.length == 5);
assertTrue(modelFiles.length == 6);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ void publishJson() throws Exception {
ByteArrayOutputStream stream = new ByteArrayOutputStream();
StreamThreatsReporter reporter = new StreamThreatsReporter(stream, ReportFormat.json);

reporter.publish(getReportHeader(), getThreats());
reporter.publish(getReportHeader(), null, getThreats());

assertTrue(!new String(stream.toByteArray()).isEmpty());
}
Expand All @@ -63,7 +63,7 @@ void publishHtml() throws Exception {
ByteArrayOutputStream stream = new ByteArrayOutputStream();
StreamThreatsReporter reporter = new StreamThreatsReporter(stream, ReportFormat.html);

reporter.publish(getReportHeader(), getThreats());
reporter.publish(getReportHeader(), null, getThreats());

assertTrue(!new String(stream.toByteArray()).isEmpty());
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
package com.github.rusakovichma.tictaac.reporter.dfd;

import com.github.rusakovichma.tictaac.mapper.ThreatModelMapper;
import com.github.rusakovichma.tictaac.model.ThreatModel;
import com.github.rusakovichma.tictaac.parser.impl.NodeTreeParser;
import com.github.rusakovichma.tictaac.parser.model.NodeTree;
import com.github.rusakovichma.tictaac.util.FileUtil;
import com.github.rusakovichma.tictaac.util.ImageUtils;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.io.InputStream;

import static org.junit.jupiter.api.Assertions.*;

class GraphvizDataFlowRenderTest {

@Test
void createDataFlow() throws IOException {
InputStream fileInputStream = FileUtil.fileToInputStream("src/test/resources/data-flow-render-test.yml");

NodeTree tree = new NodeTreeParser().getNodeTree(fileInputStream);
ThreatModelMapper mapper = new ThreatModelMapper(tree);

ThreatModel threatModel = mapper.getModel();

GraphvizDataFlowRender render = new GraphvizDataFlowRender();

assertTrue(render.createDataFlow(threatModel).length > 0);
}
}
Loading

0 comments on commit 903ec1c

Please sign in to comment.