Skip to content

Commit

Permalink
Merge pull request #81 from target/vt-augment-and-additional-updates
Browse files Browse the repository at this point in the history
Adding VT Augment Support, Many UX Changes
  • Loading branch information
phutelmyer authored Apr 19, 2024
2 parents cb23b31 + 044236e commit 65869be
Show file tree
Hide file tree
Showing 63 changed files with 3,234 additions and 1,119 deletions.
46 changes: 46 additions & 0 deletions app/blueprints/strelka.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
get_virustotal_positives,
create_vt_zip_and_download,
download_vt_bytes,
get_virustotal_widget_url,
)
from services.insights import get_insights

Expand Down Expand Up @@ -731,6 +732,51 @@ def get_mime_type_stats(user: User) -> Tuple[Response, int]:
return jsonify(stats), 200


@strelka.route("/virustotal/widget-url", methods=["POST"])
@auth_required
def get_vt_widget_url(resource: str) -> Tuple[Response, int]:
"""
Route to get a VirusTotal widget url with customized theme colors.
Returns:
A JSON response containing the VirusTotal widget url or an error message.
"""
data = request.get_json()
if not data or "resource" not in data:
return jsonify({"error": "Resource identifier is required"}), 400

# Strelka UI Defaults
fg1 = data.get("fg1", "333333") # Dark text color
bg1 = data.get("bg1", "FFFFFF") # Light background color
bg2 = data.get("bg2", "F5F5F5") # Slightly grey background for differentiation
bd1 = data.get("bd1", "E8E8E8") # Light grey border color

api_key = os.getenv("VIRUSTOTAL_API_KEY")
if not api_key:
return jsonify({"error": "VirusTotal API key is not available."}), 500

try:
# Pass the theme colors to the function
widget_url = get_virustotal_widget_url(
api_key, data["resource"], fg1, bg1, bg2, bd1
)
return jsonify({"widget_url": widget_url}), 200
except Exception as e:
return jsonify({"error": str(e)}), 500


@strelka.route("/check_vt_api_key", methods=["GET"])
def check_vt_api_key():
"""
Endpoint to check if the VirusTotal API key is available.
Returns:
A boolean response containing the VirusTotal widget url or an error message.
"""
api_key_exists = bool(os.environ.get("VIRUSTOTAL_API_KEY"))
return jsonify({"apiKeyAvailable": api_key_exists}), 200


def submissions_to_json(submission: FileSubmission) -> Dict[str, any]:
"""
Converts the given submission to a dictionary representation that can be
Expand Down
590 changes: 513 additions & 77 deletions app/poetry.lock

Large diffs are not rendered by default.

3 changes: 2 additions & 1 deletion app/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ Pathspec = "0.9.0"
Paste = "3.5.2"
Platformdirs = "2.4.1"
Protobuf = "3.18.3"
Psycopg2-Binary = "2.9.4"
Psycopg2-Binary = "2.9.9"
Pyasn1 = "0.4.8"
PyJWT = "2.4.0"
PyOpenSSL = "20.0.1"
Expand All @@ -56,6 +56,7 @@ Typing-Extensions = "4.0.1"
Urllib3 = "1.26.18"
Waitress = "2.1.2"
Wrapt = "1.13.3"
vt-py = "0.18.0"


[tool.poetry.dev-dependencies]
Expand Down
33 changes: 33 additions & 0 deletions app/services/virustotal.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,3 +126,36 @@ def download_vt_bytes(api_key: str, file_hash: str) -> BytesIO:
error_msg = f"Error downloading file from VirusTotal: {response.text}"
logging.error(error_msg)
raise Exception(error_msg)


def get_virustotal_widget_url(
api_key: str, resource: str, fg1: str, bg1: str, bg2: str, bd1: str
) -> str:
"""
Retrieves a URL for embedding the VirusTotal widget with customized theme colors.
Args:
api_key (str): The API key for accessing VirusTotal.
resource (str): The resource identifier (file hash, URL, IP, or domain).
fg1 (str): Theme primary foreground color in hex notation.
bg1 (str): Theme primary background color in hex notation.
bg2 (str): Theme secondary background color in hex notation.
bd1 (str): Theme border color.
Returns:
str: A URL for embedding the VirusTotal widget with the specified theme.
"""
url = f"https://www.virustotal.com/api/v3/widget/url?query={resource}&fg1={fg1}&bg1={bg1}&bg2={bg2}&bd1={bd1}"
headers = {"x-apikey": api_key}

try:
response = requests.get(url, headers=headers)
response.raise_for_status()
widget_data = response.json()
return widget_data.get("data", {}).get("url", "")
except requests.HTTPError as http_err:
logging.error(f"HTTP error occurred: {http_err}")
raise
except Exception as err:
logging.error(f"An error occurred: {err}")
raise
6 changes: 3 additions & 3 deletions ui/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,13 @@
"private": true,
"homepage": "/",
"dependencies": {
"@ant-design/pro-form": "^2.23.1",
"@ant-design/pro-layout": "^7.17.16",
"@ant-design/pro-form": "^2.25.1",
"@ant-design/pro-layout": "^7.19.0",
"@testing-library/jest-dom": "^5.16.4",
"@testing-library/react": "^11.1.0",
"@testing-library/user-event": "^13.5.0",
"ansi-regex": "^6.0.1",
"antd": "^5.3.0",
"antd": "^5.16.1",
"browserslist": "^4.20.4",
"chroma-js": "^2.4.2",
"css-what": "^5.1.0",
Expand Down
10 changes: 5 additions & 5 deletions ui/src/App.css
Original file line number Diff line number Diff line change
Expand Up @@ -79,9 +79,9 @@ body {
.ant-card,
.submission-card-height,
.ant-collapse-item {
border-radius: 20px;
box-shadow: 0px 4px 6px -1px rgba(0, 0, 0, 0.1),
0px 2px 4px -1px rgba(0, 0, 0, 0.06);
border-radius: 0px;
box-shadow: -1px 4px 6px -1px rgba(0, 0, 0, 0.1),
-1px -2px 4px -1px rgba(0, 0, 0, 0.06);
}

/* Collapse Styles */
Expand All @@ -92,10 +92,10 @@ body {
.ant-collapse > .ant-collapse-item > .ant-collapse-content > .ant-collapse-content-box,
.ant-collapse-content,
.ant-collapse-content-box {
border-radius: 20px !important;
border-radius: 0px !important;
background-color: #ffffff;
border: 0;
padding-top: 5px;

padding-bottom: 5px;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,15 @@ import ReactFlow, {
MiniMap,
useReactFlow,
} from "reactflow";
import EventNode from "../FileFlow/NodeTypes/EventNode.js";
import IndexConnectEdge from "../FileFlow/EdgeTypes/IndexConnectEdge.js";
import EventNode from "./NodeTypes/EventNode.js";
import IndexConnectEdge from "./EdgeTypes/IndexConnectEdge.js";
import { initialNodes, initialEdges } from "../../data/initialData.js";
import "reactflow/dist/style.css";
import { getDagreLayout } from "../../utils/dagreLayout.js";
import DownloadImage from "../../utils/downloadImage";
import ShowFileListing from "../../utils/ShowFileListing";
import NodeSearchPanel from "../../utils/NodeSearchPanel";
import ClickGuide from "../../utils/ClickGuide";
import DownloadImage from "../../utils/downloadImage.js";
import ShowFileListing from "../../utils/ShowFileListing.js";
import NodeSearchPanel from "../../utils/NodeSearchPanel.js";
import ClickGuide from "../../utils/ClickGuide.js";
import ExceededGuide from "../../utils/ExceededGuide.js";

import {
Expand Down Expand Up @@ -50,6 +50,7 @@ const FileTreeCard = ({
onNodeSelect,
fileTypeFilter,
fileYaraFilter,
fileIocFilter,
fileNameFilter,
selectedNodeData,
setSelectedNodeData,
Expand Down Expand Up @@ -125,6 +126,13 @@ const FileTreeCard = ({
);
}

// Apply the ioc filter if it's set
if (fileIocFilter) {
nodesToFilter = nodesToFilter.filter((node) =>
node.data.nodeIocList?.includes(fileIocFilter)
);
}

// If the search term is not empty, further filter the nodes
if (searchTerm.trim()) {
nodesToFilter = nodesToFilter.filter(
Expand All @@ -134,7 +142,7 @@ const FileTreeCard = ({
);
}
return nodesToFilter;
}, [nodes, searchTerm, fileTypeFilter, fileYaraFilter, fileNameFilter]);
}, [nodes, searchTerm, fileTypeFilter, fileIocFilter, fileYaraFilter, fileNameFilter]);

// Reset View when filters are applied
useEffect(() => {
Expand Down Expand Up @@ -226,7 +234,7 @@ const FileTreeCard = ({
<div
style={{
width: "100%",
height: "50vh",
height: "80vh",
border: "5px solid rgba(0, 0, 0, 0.02)",
borderRadius: "10px",
}}
Expand Down
2 changes: 1 addition & 1 deletion ui/src/components/FileFlow/NodeTypes/EventNode.js
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ const ImageTooltip = styled(Tooltip)`
align-items: center;
justify-content: center;
padding: 0 !important;
overflow: hidden;
overflow: hidden;
}
.ant-tooltip-inner img {
pointer-events: auto;
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
import React, { useState } from "react";
import { Row, Col, Modal, Typography, List, Tag } from "antd";
import "../../styles/OcrOverviewCard.css";
import "../../../styles/OcrOverviewCard.css";

const EmailOverviewCard = ({ data }) => {

const [isModalVisible, setIsModalVisible] = useState(false);

const { Text } = Typography;
Expand Down Expand Up @@ -49,7 +48,7 @@ const EmailOverviewCard = ({ data }) => {
{ title: "Sender", description: from || "No Sender", tag: "Informational" },
{
title: "Recipients",
description: to.map((name) => ({ name })) || [],
description: to?.map((name) => ({ name })) || "No Recipients",
tag: "Informational",
},
{
Expand All @@ -62,8 +61,8 @@ const EmailOverviewCard = ({ data }) => {
description: message_id || "No Message ID",
tag: "Informational",
},
// Conditionally add "Attachment Names" row only if filenames is not empty
...(filenames.length > 0
// Check if filenames exist and is an array before trying to map over it
...(Array.isArray(filenames) && filenames.length > 0
? [
{
title: "Attachment Names",
Expand All @@ -74,7 +73,8 @@ const EmailOverviewCard = ({ data }) => {
: []),
{
title: "Domains in Header",
description: received_domain.map((domain) => ({ domain })) || [],
description:
received_domain.map((domain) => ({ domain })) || "No Domains",
tag: "Informational",
},
{ title: "Body", description: body || "No Body", tag: "Informational" },
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
import React, { useState, useEffect } from 'react';
import { Collapse, Typography, Tag } from 'antd';
import EmailOverviewCard from './EmailOverviewCard';

const { Text } = Typography;

const EmailOverviewLanding = ({ selectedNodeData, expandAll }) => {
const [activeKey, setActiveKey] = useState([]);

useEffect(() => {
setActiveKey(expandAll ? ["1"] : []);
}, [expandAll]);

if (!selectedNodeData || !selectedNodeData.scan || !selectedNodeData.scan.email) {
return null;
}

const emailData = selectedNodeData.scan.email;
const subjectPreview = emailData.subject && emailData.subject.length > 0
? `${emailData.subject.substring(0, 47)}...`
: "No Subject";
const attachmentsCount = emailData.total && emailData.total.attachments;

return (
<Collapse
activeKey={activeKey}
onChange={(keys) => setActiveKey(keys)}
style={{ width: "100%", marginBottom: "10px" }}
>
<Collapse.Panel
header={
<div
style={{
display: "flex",
justifyContent: "space-between",
alignItems: "center",
}}
>
<div style={{ display: "flex", alignItems: "center" }}>
<div>
<Text strong>Email</Text>
<div style={{ fontSize: "smaller", color: "#888" }}>{subjectPreview}</div>
</div>
</div>
<div>
<Tag color="default">
<b>Attachments: {attachmentsCount}</b>
</Tag>
<Tag
color={emailData.base64_thumbnail ? "success" : "error"}
>
<b>{emailData.base64_thumbnail ? "Preview Available" : "No Preview"}</b>
</Tag>
</div>
</div>
}
key="1"
>
<EmailOverviewCard data={selectedNodeData} />
</Collapse.Panel>
</Collapse>
);
};

export default EmailOverviewLanding;
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import React, { useState } from "react";
import { Input, Checkbox, Typography, Row, Col } from "antd";
import "../../styles/ExiftoolOverviewCard.css"
import "../../../styles/ExiftoolOverviewCard.css"

const { Text } = Typography;

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
import React, { useState, useEffect } from 'react';
import { Collapse, Typography } from 'antd';
import FileExiftoolCard from './ExiftoolOverviewCard';

const { Text } = Typography;

const ExiftoolOverviewLanding = ({ selectedNodeData, expandAll }) => {
const [activeKey, setActiveKey] = useState([]);

useEffect(() => {
setActiveKey(expandAll ? ["1"] : []);
}, [expandAll]);

if (!selectedNodeData || !selectedNodeData.scan || !selectedNodeData.scan.exiftool) {
return null;
}

const metadataCount = Object.keys(selectedNodeData.scan.exiftool).length;

return (
<Collapse
activeKey={activeKey}
onChange={(keys) => setActiveKey(keys)}
style={{ width: "100%", marginBottom: "10px" }}
>
<Collapse.Panel
header={
<div>
<Text strong>File Metadata</Text>
<div style={{ fontSize: "smaller", color: "#888" }}>
Metadata Count: {metadataCount}
</div>
</div>
}
key="1"
>
<FileExiftoolCard data={selectedNodeData} />
</Collapse.Panel>
</Collapse>
);
};

export default ExiftoolOverviewLanding;
Loading

0 comments on commit 65869be

Please sign in to comment.