Adding VT Lookup

README.md

@@ -81,11 +81,12 @@ Please reference `./app/example.env` for environment variable setup.
 The following detail the configuration items in `./app/config/config.py`.
 
 | Field Name                              | Value                                                                   | Required |
-| --------------------------------------- | ----------------------------------------------------------------------- | -------- |
+|-----------------------------------------|-------------------------------------------------------------------------| -------- |
 | STRELKA_HOST                            | Strelka hostname (e.g., `0.0.0.0`)                                      | Yes      |
 | STRELKA_PORT                            | Strelka port number (e.g., `57314`)                                     | Yes      |
 | STRELKA_CERT                            | Path to certificate for Strelka, if needed (e.g., `/path/to/cert.pem`)  | No       |
 | CA_CERT_PATH                            | Path to CA certificates for LDAP, if needed (e.g., `/path/to/ca_certs`) | No       |
+| VIRUSTOTAL_API_KEY                      | API Key for VirusTotal Hash Lookup                                      | Yes      |
 | LDAP_URL                                | URL to LDAP server (e.g., `ldaps://ldap.example.com:636`)               | No       |
 | LDAP_SEARCH_BASE                        | Search base for LDAP queries (e.g., `DC=example,DC=com`)                | No       |
 | LDAP_USERNAME_ORGANIZATION              | Username organization for LDAP queries (e.g., `org//`)                  | No       |

app/blueprints/auth.py

@@ -102,6 +102,7 @@ def login():
         # without a lookup on the user submitting it.
         session["user_id"] = dbUser.id
         session["logged_in"] = True
+
     except Exception as err:
         # current_app.logger.error("Failed connection to database: %s", err)
         return jsonify({"error": "Failed to connect to database"}), 400

app/blueprints/strelka.py

@@ -1,4 +1,6 @@
 import datetime
+import logging
+import os
 from typing import Any, Dict, Tuple
 
 from flask import Blueprint, current_app, jsonify, request, session, Response
@@ -8,6 +10,7 @@
 from models import FileSubmission, User
 from services.auth import auth_required
 from services.strelka import get_db_status, get_frontend_status, submit_data
+from services.virustotal import get_virustotal_positives
 
 strelka = Blueprint("strelka", __name__, url_prefix="/strelka")
 
@@ -117,6 +120,23 @@ def submit_file(user: User) -> Tuple[Response, int]:
             # Get the submitted file object from the analysis results.
             submitted_file = response[0]
 
+            # If VirusTotal API key provided, get positives
+            # -1    = VirusTotal Lookup Error
+            # -2    = VirusTotal API Key Not Provided
+            # >= 0  = Response Positives from VirusTotal
+            virustotal_positives = -2
+
+            if os.environ.get("VIRUSTOTAL_API_KEY"):
+                try:
+                    virustotal_positives = get_virustotal_positives(
+                        api_key=os.environ.get("VIRUSTOTAL_API_KEY"),
+                        file_hash=file["response"]["scan"]["hash"],
+                    )
+                except Exception as e:
+                    logging.warning(
+                        f"Could not process VirusTotal search with error: {e} "
+                    )
+
             # Create a new submission object and add it to the database.
             new_submission = FileSubmission(
                 get_request_id(submitted_file),
@@ -133,6 +153,7 @@ def submit_file(user: User) -> Tuple[Response, int]:
                 submitted_description,
                 submitted_at,
                 get_request_time(submitted_file),
+                virustotal_positives,
             )
 
             db.session.add(new_submission)

app/config/config.py

@@ -35,6 +35,7 @@ class Config(object):
     DATABASE_DBNAME = os.environ.get("DATABASE_DBNAME", "strelka_ui")
 
     # API Details
+    VIRUSTOTAL_API_KEY = os.environ.get("VIRUSTOTAL_API_KEY", "")
     API_KEY_EXPIRATION = os.environ.get("API_KEY_EXPIRATION", "999")
 
     # LDAP Details

app/models.py

@@ -28,6 +28,7 @@ class FileSubmission(db.Model):
         user (User): The user who submitted the file.
         submitted_at (datetime.datetime): The date and time the file was submitted.
         processed_at (datetime.datetime): The date and time the file was processed.
+        virustotal_positives (int): Amount of positives returned by VirusTotal.
     """
 
     __tablename__ = "file_submission"
@@ -60,6 +61,9 @@ class FileSubmission(db.Model):
     )
     processed_at: datetime.datetime = db.Column(db.DateTime())
 
+    # Enrichments
+    virustotal_positives: int = db.Column(db.Integer)
+
     def __init__(
         self,
         file_id: str,
@@ -76,6 +80,7 @@ def __init__(
         submitted_description: str,
         submitted_at: datetime.datetime,
         processed_at: datetime.datetime,
+        virustotal_positives: int,
     ):
         self.file_id = file_id
         self.file_name = file_name
@@ -91,6 +96,7 @@ def __init__(
         self.submitted_description = submitted_description
         self.submitted_at = submitted_at
         self.processed_at = processed_at
+        self.virustotal_positives = virustotal_positives
 
     def __repr__(self):
         return "<id {}>".format(self.id)