Step 2: Train and Read the Tree

Training the Tree

from sklearn.tree import DecisionTreeClassifier

model = DecisionTreeClassifier(
    max_depth=4,          # limits depth to prevent overfitting
    criterion='gini',     # splitting criterion (or 'entropy')
    random_state=42
)
model.fit(X_train, y_train)

Visualising the Tree

plot_tree() renders the tree as a matplotlib figure where you can see every decision:

from sklearn.tree import plot_tree
import matplotlib.pyplot as plt

plt.figure(figsize=(20, 10))
plot_tree(model,
          feature_names=FEATURES,
          class_names=CLASS_NAMES,
          filled=True,           # colour by majority class
          rounded=True,
          fontsize=10)
plt.tight_layout()
plt.show()

Each node shows: the split rule, Gini impurity, sample count, and class distribution.

Reading the Tree as Text

from sklearn.tree import export_text

rules = export_text(model, feature_names=FEATURES)
print(rules)

Output looks like:

|--- connection_rate <= 50.50
|   |--- class: benign
|--- connection_rate >  50.50
|   |--- unique_dest_ports <= 20.50
|   |   |--- bytes_sent <= 100000.00
|   |   |   |--- class: benign
|   |   |--- bytes_sent >  100000.00
|   |   |   |--- class: exfiltration
|   |--- unique_dest_ports >  20.50
|   |   |--- class: port_scan

You can turn these rules directly into firewall policies or SIEM alerts.