HarshV1315/CVE2TTP

1from flask import Flask, render_template, request, jsonify
2import torch
3import torch.nn as nn
4import numpy as np
5from transformers import AutoTokenizer, AutoModel
6import nvdlib
7
8# Flask app initialization
9app = Flask(__name__)
10
11# Define the model architecture
12class Model(nn.Module):
13    def __init__(self):
14        super(Model, self).__init__()
15        self.transformer_model = AutoModel.from_pretrained('jackaduma/SecRoBERTa')
16        self.dropout = nn.Dropout(0.3)
17        self.output = nn.Linear(768, 14)
18
19    def forward(self, input_ids, attention_mask=None):
20        _, o2 = self.transformer_model(
21            input_ids=input_ids,
22            attention_mask=attention_mask,
23            return_dict=False
24        )
25        x = self.dropout(o2)
26        out = self.output(x)
27        return out
28
29# Function to predict MITRE ATT&CK techniques
30def predict_techniques(model, tokenizer, cve_description, device):
31    tokenized_input = tokenizer.encode_plus(
32        cve_description,
33        max_length=320,
34        padding='max_length',
35        truncation=True,
36        return_attention_mask=True,
37        return_tensors='pt'
38    )
39    input_ids = tokenized_input['input_ids'].to(device)
40    attention_mask = tokenized_input['attention_mask'].to(device)
41    with torch.no_grad():
42        logits = model(input_ids, attention_mask)
43        probs = torch.sigmoid(logits).cpu().numpy()
44
45    predicted_techniques = np.round(probs)
46    return predicted_techniques
47
48# Global variables for model and tokenizer
49global_model = None
50global_tokenizer = None
51
52# Lazy loading function to get the model and tokenizer
53def get_model_and_tokenizer(device='cpu'):
54    global global_model, global_tokenizer
55    if global_model is None or global_tokenizer is None:
56        global_model = Model()
57        global_model.load_state_dict(torch.load('tactic_predict.pt', map_location=device, weights_only=True))
58        global_model.to(device)
59        global_model.eval()
60        global_tokenizer = AutoTokenizer.from_pretrained('jackaduma/SecRoBERTa')
61    return global_model, global_tokenizer
62
63# Route for the home page
64@app.route('/')
65def home():
66    return render_template('index.html')
67
68# Route to handle form submission and return results
69@app.route('/predict', methods=['POST'])
70def predict():
71    cve_id = request.form['cve_id']
72    r = nvdlib.searchCVE(cveId=cve_id)[0]
73    desc_list = r.descriptions
74    cve_data = next(desc.value for desc in desc_list if desc.lang == "en")
75
76    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
77    
78    # Load model and tokenizer lazily
79    model, tokenizer = get_model_and_tokenizer(device)
80    
81    predicted_techniques = predict_techniques(model, tokenizer, cve_data, device)
82
83    tactic_names = [
84        "Reconnaissance", "Resource Development", "Initial Access", "Execution", 
85        "Persistence", "Privilege Escalation", "Defense Evasion", 
86        "Credential Access", "Discovery", "Lateral Movement", "Collection", 
87        "Command and Control", "Exfiltration", "Impact"
88    ]
89
90    predicted_tactic_names = [tactic_names[i] for i, val in enumerate(predicted_techniques[0]) if val == 1]
91    
92    return render_template('result.html', tactics=predicted_tactic_names, cve_id=cve_id, cve_desc=cve_data)
93
94# Run the app
95if __name__ == "__main__":
96    app.run(host="0.0.0.0", port=7860)
97
98