Lightml ML Training
Continuing from the initial proof-of-concept in the development write-up, there has been some new development and changes in the performance of the model. The initial sample was no where close to anything production ready because it was trained on about 10 sample logs, which is not realistic at all. So, the main focus from that point was finding the right dataset to use to train the model. The model needs to be able to predict the techniques given live and native Linux based logs. I haven’t been able to find the perfect dataset I can use for this task, which means I have to create the dataset myself. Considering how much work and time this would take, I immediately thought of more efficient ways to achieve this task, but they all lead to me generating the detection samples for training. The MITRE ATT&CK framework has over 500 techniques and sub-techniques, and each would need at least 10 sample detection logs for training, which means over 10K sample logs combined. It’s also worth noting that not all the techniques are Linux related or technical in nature, but we still need some form of detection samples for better training coverage of the model. After some research and surfing the great internet, I realized that MITRE has Excel sheets of all the enterprise related techniques and their associated tactics.
Creating Training Data
After downloading the Excel-sheet, it needs to be converted to CSV, and this can be done in multiple ways. If you’re not sure how to do so via the Python or the command line, you can upload the document to Google Sheets and download it as CSV or open it with the Excel app and click Save as, then you can change the format. For the purpose of the work and end-goal, we don’t need all the data in the CSV, so I wrote a python script to parse through the data and create a new dataframe from the initial CSV, which was then converted to JSON format. Just make sure the CSV file is in the working directory, or specify the absolute path in the script below.
1
2
3
4
5
6
7
| import panadas as pd
original_df = pd.read_csv('enterprise-attack-techniques.csv')
columns_to_keep = ['ID', 'name', 'tactics', 'description', 'detection']
new_df = original_df[columns_to_keep]
new_df.to_csv('lightml-train-01.csv', index=False)
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
| import pandas as pd
import json
df = pd.read_csv('lightml-train-01.csv')
mitre_data = {}
for _, row in df.iterrows():
mitre_id = row['ID']
mitre_data[mitre_id] = {
"name": row['name'],
"tactic": row['tactics'],
"description": row['description'] if 'description' in row else "",
"detection": [],
}
with open('lightml-train-02.json', 'w') as json_file:
json.dump(mitre_data, json_file, indent=4)
print("JSON file has been saved as 'lightml-train-02.json'.")
|
From the initial JSON template, I was able to go through and add detection examples for each technique, which should look something like the following.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
| "T1548": {
"name": "Abuse Elevation Control Mechanism",
"tactic": "Defense Evasion, Privilege Escalation",
"description": "Adversaries may circumvent mechanisms designed to control elevate privileges to gain higher-level permissions. Most modern systems contain native elevation control mechanisms that are intended to limit privileges that a user can perform on a machine. Authorization has to be granted to specific users in order to perform tasks that can be considered of higher risk.(Citation: TechNet How UAC Works)(Citation: sudo man page 2018) An adversary can perform several methods to take advantage of built-in control mechanisms in order to escalate privileges on a system.(Citation: OSX Keydnap malware)(Citation: Fortinet Fareit)",
"detection": [
"Mar 8 10:15:32 localhost sudo: user : TTY=pts/0 ; PWD=/home/user ; COMMAND=/bin/bash",
"Mar 8 10:16:01 localhost polkitd[1234]: Registered Authentication Agent for unix-session:1",
"Mar 8 10:16:45 localhost sudo: user : TTY=pts/0 ; PWD=/home/user ; COMMAND=/usr/bin/systemctl start myservice",
"Mar 8 10:17:10 localhost kernel: [12345.678901] audit: type=USER_AUTH msg=audit(1615196580.123:456): user pid=7890 uid=1000 auid=1000 ses=1",
"Mar 8 10:17:35 localhost sudo: user : TTY=pts/0 ; PWD=/home/user ; COMMAND=/usr/bin/dnf update",
"Mar 8 10:18:12 localhost gnome-session[1234]: User 'user' authenticated successfully to perform elevated operations",
"Mar 8 10:18:45 localhost dbus-daemon[1234]: Activating service name='org.freedesktop.login1' requested by ':1.23' (uid=1000 pid=7890)",
"Mar 8 10:19:01 localhost polkitd[1234]: Unregistered Authentication Agent for unix-session:1",
"Mar 8 10:19:30 localhost gnome-session[1234]: User 'user' has elevated privileges for session actions",
"Mar 8 10:20:05 localhost kernel: [12350.123456] audit: type=USER_AUTH msg=audit(1615196800.123:461): user pid=7896 uid=1000 auid=1000 ses=1"
]
},
|
This was repeated for the rest of the techniques and sub-techniques, which is time consuming and it is hard to differentiate the techniques based on the logs because there are slight differences and depending on the context, one log sample could fall under multiple tactics. This is why I am willing to settle for at least 80% prediction accuracy. Without getting to deep into the design of OffStrike EDR, it’s safe to compromise here with the model score because the LLM will close the gap. Based on the initial training script, I wrote another script to clean and format the final training data accordingly.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
| import json
import csv
def json_to_csv(json_file, csv_file):
"""
Converts a JSON file with security techniques and their corresponding detection logs to a CSV file.
Args:
json_file (str): Path to the JSON file.
csv_file (str): Path to the CSV file to be created.
"""
with open(json_file, 'r') as f:
data = json.load(f)
with open(csv_file, 'w', newline='') as f:
writer = csv.writer(f)
writer.writerow(['technique', 'detection'])
for technique, details in data.items():
technique = technique.split('.')[0]
for detection in details['detection']:
writer.writerow([technique, detection])
json_to_csv('lightml-train-03.json', 'lightml-train-04.csv')
|
Updated Training Script - Test
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
| #!/bin/python3
import pandas as pd
from sklearn.model_selection import train_test_split, GridSearchCV, StratifiedKFold
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline
from sklearn.metrics import classification_report
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
import string, nltk
nltk.download('punkt_tab')
nltk.download('stopwords')
df = pd.read_csv('lightml-train-04.csv')
print("Training Data:\n", df)
def preprocess_text(text):
tokens = word_tokenize(text.lower())
stop_words = set(stopwords.words('english'))
tokens = [word for word in tokens if word not in stop_words and word not in string.punctuation]
return " ".join(tokens)
df["processed_log"] = df["detection"].apply(preprocess_text)
pipeline = Pipeline([
('tfidf', TfidfVectorizer()),
('classifier', MultinomialNB())
])
X_train, X_test, y_train, y_test = train_test_split(df["processed_log"], df["technique"], test_size=0.2, random_state=42)
stratified_kfold = StratifiedKFold(n_splits=2)
param_grid = {
'classifier__alpha': [0.1, 1.0, 10.0]
}
grid_search = GridSearchCV(pipeline, param_grid, cv=stratified_kfold)
grid_search.fit(X_train, y_train)
y_pred = grid_search.predict(X_test)
print("ML Classification Report:\n", classification_report(y_test, y_pred, zero_division=1))
def predict_mitre_technique(logs):
processed_logs = [preprocess_text(log) for log in logs]
predictions = grid_search.predict(processed_logs)
return predictions
plain_security_logs = [
"Mar 9 05:00:01 localhost security.log: Command 'net localgroup Administrators user1 /add' executed to add user 'attacker' to local admins.",
"Mar 8 10:37:05 localhost sudo: user : TTY=pts/0 ; PWD=/home/user ; COMMAND=/usr/bin/sudo -n whoami",
"Mar 8 10:20:05 localhost kernel: [12350.123456] audit: type=USER_AUTH msg=audit(1615196800.123:461): user pid=7896 uid=1000 auid=1000 ses=1",
"/var/log/application.log: User-Agent 'curl/7.58.0' used in suspicious access attempt to gather host info at 2025-03-11 22:40:00",
"Mar 9 07:00:01 localhost security.log: Modification detected in '/home/user/.ssh/authorized_keys' by user 'attacker'.",
"Mar 10 00:00:01 localhost dnsmasq: query[TXT] secret-data.example.com from 192.0.2.10"
]
predicted_techniques = predict_mitre_technique(plain_security_logs)
for log, technique in zip(plain_security_logs, predicted_techniques):
print(f"Log: {log} \nPredicted MITRE Technique: {technique}\n")
|
Results - Test
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
| Training Data:
technique detection
0 T1548 Mar 8 10:15:32 localhost sudo: user : TTY=...
1 T1548 Mar 8 10:16:01 localhost polkitd[1234]: Regis...
2 T1548 Mar 8 10:16:45 localhost sudo: user : TTY=...
3 T1548 Mar 8 10:17:10 localhost kernel: [12345.67890...
4 T1548 Mar 8 10:17:35 localhost sudo: user : TTY=...
... ... ...
3800 T1220 2023-11-08 16:25:00 INFO: Network traffic dete...
3801 T1220 2023-11-08 16:30:00 WARNING: Unexpected file e...
3802 T1220 2023-11-08 16:35:00 INFO: User 'admin' opened ...
3803 T1220 2023-11-08 16:40:00 ALERT: Execution of 'wget'...
3804 T1220 2023-11-08 16:45:00 INFO: Anomalous user behav...
[3805 rows x 2 columns]
ML Classification Report:
precision recall f1-score support
T1001 0.67 1.00 0.80 8
T1005 1.00 0.00 0.00 1
T1006 1.00 0.00 0.00 2
T1008 1.00 0.00 0.00 1
T1010 1.00 0.00 0.00 3
T1011 1.00 0.50 0.67 4
T1020 1.00 0.40 0.57 5
T1025 1.00 0.00 0.00 3
T1030 1.00 0.00 0.00 4
T1037 0.92 0.92 0.92 12
T1039 1.00 0.00 0.00 2
T1041 1.00 0.25 0.40 4
T1047 1.00 0.00 0.00 2
T1048 0.71 0.62 0.67 8
T1052 1.00 0.00 0.00 2
T1059 0.54 0.97 0.69 36
T1068 1.00 0.00 0.00 5
T1071 0.75 0.88 0.81 17
...
accuracy 0.73 761
macro avg 0.92 0.45 0.44 761
weighted avg 0.82 0.73 0.67 761
Log: Mar 9 05:00:01 localhost security.log: Command 'net localgroup Administrators user1 /add' executed to add user 'attacker' to local admins.
Predicted MITRE Technique: T1098
Log: Mar 8 10:37:05 localhost sudo: user : TTY=pts/0 ; PWD=/home/user ; COMMAND=/usr/bin/sudo -n whoami
Predicted MITRE Technique: T1548
Log: Mar 8 10:20:05 localhost kernel: [12350.123456] audit: type=USER_AUTH msg=audit(1615196800.123:461): user pid=7896 uid=1000 auid=1000 ses=1
Predicted MITRE Technique: T1548
Log: /var/log/application.log: User-Agent 'curl/7.58.0' used in suspicious access attempt to gather host info at 2025-03-11 22:40:00
Predicted MITRE Technique: T1592
Log: Mar 9 07:00:01 localhost security.log: Modification detected in '/home/user/.ssh/authorized_keys' by user 'attacker'.
Predicted MITRE Technique: T1098
Log: Mar 10 00:00:01 localhost dnsmasq: query[TXT] secret-data.example.com from 192.0.2.10
Predicted MITRE Technique: T1071
|