Upload 5 files

.dockerignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+.env

Dockerfile

@@ -0,0 +1,20 @@
+# Use an official Python runtime as a parent image
+FROM python:3.9-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy the current directory contents into the container at /app
+COPY . /app
+
+# Install any needed packages specified in requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Make port 5000 available to the world outside this container
+EXPOSE 5000
+
+# Define environment variable
+ENV FLASK_APP=app.py
+
+# Run app.py when the container launches
+CMD ["flask", "run", "--host", "0.0.0.0"]

app.py

@@ -0,0 +1,97 @@
+from flask import Flask, render_template, request, jsonify
+import torch
+import torch.nn as nn
+import numpy as np
+from transformers import AutoTokenizer, AutoModel
+import nvdlib
+
+# Flask app initialization
+app = Flask(__name__)
+
+# Define the model architecture
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+        self.transformer_model = AutoModel.from_pretrained('jackaduma/SecRoBERTa')
+        self.dropout = nn.Dropout(0.3)
+        self.output = nn.Linear(768, 14)
+
+    def forward(self, input_ids, attention_mask=None):
+        _, o2 = self.transformer_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            return_dict=False
+        )
+        x = self.dropout(o2)
+        out = self.output(x)
+        return out
+
+# Function to predict MITRE ATT&CK techniques
+def predict_techniques(model, tokenizer, cve_description, device):
+    tokenized_input = tokenizer.encode_plus(
+        cve_description,
+        max_length=320,
+        padding='max_length',
+        truncation=True,
+        return_attention_mask=True,
+        return_tensors='pt'
+    )
+    input_ids = tokenized_input['input_ids'].to(device)
+    attention_mask = tokenized_input['attention_mask'].to(device)
+    with torch.no_grad():
+        logits = model(input_ids, attention_mask)
+        probs = torch.sigmoid(logits).cpu().numpy()
+
+    predicted_techniques = np.round(probs)
+    return predicted_techniques
+
+# Global variables for model and tokenizer
+global_model = None
+global_tokenizer = None
+
+# Lazy loading function to get the model and tokenizer
+def get_model_and_tokenizer(device='cpu'):
+    global global_model, global_tokenizer
+    if global_model is None or global_tokenizer is None:
+        global_model = Model()
+        global_model.load_state_dict(torch.load('tactic_predict.pt', map_location=device, weights_only=True))
+        global_model.to(device)
+        global_model.eval()
+        global_tokenizer = AutoTokenizer.from_pretrained('jackaduma/SecRoBERTa')
+    return global_model, global_tokenizer
+
+# Route for the home page
+@app.route('/')
+def home():
+    return render_template('index.html')
+
+# Route to handle form submission and return results
+@app.route('/predict', methods=['POST'])
+def predict():
+    cve_id = request.form['cve_id']
+    r = nvdlib.searchCVE(cveId=cve_id)[0]
+    desc_list = r.descriptions
+    cve_data = next(desc.value for desc in desc_list if desc.lang == "en")
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # Load model and tokenizer lazily
+    model, tokenizer = get_model_and_tokenizer(device)
+    
+    predicted_techniques = predict_techniques(model, tokenizer, cve_data, device)
+
+    tactic_names = [
+        "Reconnaissance", "Resource Development", "Initial Access", "Execution", 
+        "Persistence", "Privilege Escalation", "Defense Evasion", 
+        "Credential Access", "Discovery", "Lateral Movement", "Collection", 
+        "Command and Control", "Exfiltration", "Impact"
+    ]
+
+    predicted_tactic_names = [tactic_names[i] for i, val in enumerate(predicted_techniques[0]) if val == 1]
+    
+    return render_template('result.html', tactics=predicted_tactic_names, cve_id=cve_id, cve_desc=cve_data)
+
+# Run the app
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)
+

requirements.txt

@@ -0,0 +1,5 @@
+Flask
+torch
+numpy
+transformers
+nvdlib

tactic_predict.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdbe380766df249f0c50081b3a482e6eab588552603c23329087342ed4190870
+size 333893590