From de7124477690559c069389f57bfdc44a4b46017f Mon Sep 17 00:00:00 2001
From: Thorsten Sick <thorsten.sick@avast.com>
Date: Mon, 8 Nov 2021 09:51:01 +0100
Subject: [PATCH] Adding line numbers and file names to metasploit attacks

---
 app/attack_log.py                | 17 ++++++++++++++---
 plugins/base/plugin_base.py      | 11 +++++++++++
 templates/attack_description.rst |  1 +
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/app/attack_log.py b/app/attack_log.py
index efa9968..12c4b42 100644
--- a/app/attack_log.py
+++ b/app/attack_log.py
@@ -2,6 +2,7 @@
 
 """ Logger for the attack side. Output must be flexible, because we want to be able to feed it into many different processes. From ML to analysts """
 
+from inspect import currentframe, getsourcefile
 import json
 import datetime
 from random import randint
@@ -84,7 +85,11 @@ class AttackLog():
                 "697e8a432031075e47cccba24417013d": "Persistence"}
 
         ttp_data = {"t1547": "Persistence",
-                    "t1547.001": "Persistence"}
+                    "t1547.001": "Persistence",
+                    "t1547.004": "Persistence",
+                    "t1547.005": "Persistence",
+                    "t1547.009": "Persistence",
+                    "t1547.010": "Persistence"}
 
         if ability_id in data:
             return data[ability_id]
@@ -108,7 +113,11 @@ class AttackLog():
                 "697e8a432031075e47cccba24417013d": "TA0003"}
 
         ttp_data = {"t1547": "TA0003",
-                    "t1547.001": "TA0003"}
+                    "t1547.001": "TA0003",
+                    "t1547.004": "TA0003",
+                    "t1547.005": "TA0003",
+                    "t1547.009": "TA0003",
+                    "t1547.010": "TA0003"}
 
         if ability_id in data:
             return data[ability_id]
@@ -519,7 +528,9 @@ class AttackLog():
                 "description": kwargs.get("description", None),  # Generic description for this attack. Set by the attack
                 "situation_description": kwargs.get("situation_description", None),  # Description for the situation this attack was run in. Set by the plugin or attacker emulation
                 "countermeasure": kwargs.get("countermeasure", None),  # Set by the attack
-                "result": None
+                "result": None,
+                "sourcefile": kwargs.get("sourcefile", getsourcefile(currentframe().f_back)),
+                "sourceline": kwargs.get("sourceline", currentframe().f_back.f_lineno)
                 }
         self.__add_to_log__(data)
 
diff --git a/plugins/base/plugin_base.py b/plugins/base/plugin_base.py
index e8a510b..049a560 100644
--- a/plugins/base/plugin_base.py
+++ b/plugins/base/plugin_base.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 """ Base class for all plugin types """
 
+from inspect import currentframe
 import os
 from typing import Optional
 import yaml
@@ -26,6 +27,16 @@ class BasePlugin():
 
         self.default_config_name = "default_config.yaml"
 
+    def get_filename(self):
+        """ Returns the current filename.  """
+        cf = currentframe()
+        return cf.f_back.filename
+
+    def get_linenumber(self):
+        """ Returns the current linenumber.  """
+        cf = currentframe()
+        return cf.f_back.f_lineno
+
     def get_playground(self):
         """ Returns the machine specific playground
 
diff --git a/templates/attack_description.rst b/templates/attack_description.rst
index dfd8c80..3391147 100644
--- a/templates/attack_description.rst
+++ b/templates/attack_description.rst
@@ -75,6 +75,7 @@ Metasploit attack {{ e.name }}
 + Hunting Tag: {{ e.hunting_tag}}
 + At {{ e.timestamp }} a Metasploit command {{ e.name }} was used to attack {{ e.target }} from {{ e.source }}.
 + Description: {{ e.description }}
++ Code in {{ e.sourcefile }} / {{ e.sourceline }}
 {% if e.metasploit_command is string() %}
 + Metasploit command: {{ e.metasploit_command }}
 {% endif %}