Adding logic for MITRE enrichment

README.md

@@ -4,12 +4,13 @@ A simple linter for Sigma rules
 
 ## Description
 
-sigmalint is a command line interface for validating Sigma rules against the Sigma schema.
+sigmalint is a command line interface for validating and enriching Sigma rules against the Sigma schema.
 
 The available arguments are:
 * `--sigmainput` - Path to a directory that comtains Sigma files or to a single Sigma file.
 * `--directory` - Flag for if sigmainput is a directory
 * `--method` - The schema validator that you wish to use (Default: rx)
+* `--mitre` - Enrich Sigma file with MITRE content based off of any MITRE references in `tags`
 
 The available methods are:
 * `rx` - uses PyRx and the Rx schema from the Sigma repo
@@ -32,6 +33,7 @@ Options:
                                [required]
  --directory                  Flag for if sigmainput is a directory
  --method [rx|jsonschema|s2]  Validation method.
+ --mitre                      Enrich Sigma file with MITRE content
  --help                       Show this message and exit.
 ```
 

requirements.txt

@@ -4,3 +4,4 @@ pyrx==0.*,>=0.3.0
 pytest==5.*,>=5.4.3
 pytest-cov==2.*,>=2.10.0
 pyyaml==5.*,>=5.3.1
+beautifulsoup4==4.*,>=4.9.3

setup.py

@@ -20,12 +20,12 @@
     author='Ryan Plas',
     author_email='ryan.plas@stage2sec.com',
     entry_points={"console_scripts": ["sigmalint = sigmalint.sigmalint:cli"]},
-    packages=['sigmalint', 'sigmalint.schema'],
+    packages=['sigmalint', 'sigmalint.schema', 'sigmalint.modules'],
     package_dir={"": "."},
     package_data={},
     install_requires=[
         'click==7.*,>=7.1.2', 'jsonschema==3.*,>=3.2.0', 'pyrx==0.*,>=0.3.0',
-        'pyyaml==5.*,>=5.3.1'
+        'pyyaml==5.*,>=5.3.1', 'beautifulsoup4==4.*,>=4.9.3'
     ],
     extras_require={
         "dev": ["pytest==5.*,>=5.4.3", "pytest-cov==2.*,>=2.10.0"]},

sigmalint/modules/__init__.py

@@ -0,0 +1 @@
+from .mitre import mitre_pull

sigmalint/modules/mitre.py

@@ -0,0 +1,37 @@
+#!/usr/bin/python3
+from bs4 import BeautifulSoup
+import requests
+import re
+
+def mitre_pull(technique_id):
+  tactics = []
+
+  # Redirect
+  try:
+      url = "https://attack.mitre.org/techniques/{}/".format(str(technique_id))
+      redirect = requests.get(url,allow_redirects=True)
+      redirect_url = re.search('url=/(.*)"',redirect.text) 
+
+      # Technique URL
+      mitre_url = "https://attack.mitre.org/{}/".format(str(redirect_url.group(1)))
+      mitre_request = requests.get(mitre_url,allow_redirects=False)
+      mitre_soup = BeautifulSoup(mitre_request.text,"html.parser")
+
+      # Tactic ID's
+      tactic_id = mitre_soup.find('div', class_='card-data', id='card-tactics').get_text()
+      tactic_id = tactic_id.replace('Tactics:','').replace(' ','').replace('Tactic:\n','').strip("\n") 
+      tactic_id = tactic_id.split(',') 
+      tactics = tactic_id
+
+      # Sub Techniques
+      sub_techniques_regex = re.search('Sub-technique of(.*?)</a>',mitre_request.text,re.M | re.DOTALL)
+      sub_techniques = sub_techniques_regex.group(1).replace(':&nbsp;','').replace(' ','').strip("\n")
+      sub_techniques = re.sub('<[^>]+>','',sub_techniques)
+      sub_techniques = sub_techniques.replace('\n','')
+
+      # Reference
+      references = "https://attack.mitre.org/techniques/"+str(technique_id)+"/"
+      # Make sure we return technique_id if valid values were returned
+      return tactics, sub_techniques, technique_id, references
+  except:
+      return None, None, None, None 

sigmalint/sigmalint.py

@@ -4,7 +4,9 @@
 import yaml
 import pyrx
 import jsonschema
+import re
 
+from .modules.mitre import mitre_pull
 from .schema import rx_schema, json_schema, s2_schema
 
 rx = pyrx.Factory({'register_core_types': True})
@@ -15,7 +17,8 @@
 @click.option('--sigmainput', type=click.Path(exists=True, file_okay=True, readable=True, resolve_path=True), help='Path to a directory that comtains Sigma files or to a single Sigma file.', required=True)
 @click.option('--directory', is_flag=True, help="Flag for if sigmainput is a directory")
 @click.option('--method', type=click.Choice(['rx', 'jsonschema', 's2'], case_sensitive=False), default='rx', help='Validation method.')
-def cli(sigmainput, directory, method):
+@click.option('--mitre', is_flag=True, help='Enrich Sigma file with MITRE content based off of any MITRE references in tags.  This will append to the Sigma file.', required=False)
+def cli(sigmainput, directory, method, mitre):
     results = []
     filepaths = []
 
@@ -56,6 +59,39 @@ def cli(sigmainput, directory, method):
                         result = False if len(errors) > 0 else True
                         results.append({'result': result, 'reasons': errors, 'filename': filename})
 
+                    if mitre:
+                        ## Most mentions of MITRE are currently found in the 'tags' field for Sigma rules
+                        if 'tags' in sigma_yaml_list[0]:
+                            try:
+                                for tag in sigma_yaml_list[0]['tags']: 
+                                    if len(re.findall(r'(?i)t\d{4}',tag)) > 0:
+                                        mitre_id = re.findall(r'(?i)t\d{4}',tag)
+                                        mitre_id = mitre_id[0].upper()
+                                        tactics, sub_techniques, technique_id, references = mitre_pull(mitre_id) 
+                                        ## Prevent multiple writes to file if `mitre` key already exists
+                                        if 'mitre' not in sigma_yaml_list[0]:
+                                            sigma_yaml_list[0]['mitre'] = {} 
+                                            sigma_yaml_list[0]['mitre']['tactics'] = [] 
+                                            sigma_yaml_list[0]['mitre']['subTechniques'] = []
+                                            sigma_yaml_list[0]['mitre']['techniqueIds'] = []
+                                            sigma_yaml_list[0]['mitre']['references'] = []
+                                        if tactics is not None and sub_techniques is not None and technique_id is not None and references is not None: 
+                                            for tactic in tactics:
+                                                ## Make sure we aren't duplicating
+                                                if tactic not in sigma_yaml_list[0]['mitre']['tactics']:
+                                                    sigma_yaml_list[0]['mitre']['tactics'].append(tactic) 
+                                                if sub_techniques not in sigma_yaml_list[0]['mitre']['subTechniques']:
+                                                    sigma_yaml_list[0]['mitre']['subTechniques'].append(sub_techniques)
+                                                if technique_id not in sigma_yaml_list[0]['mitre']['techniqueIds']:
+                                                    sigma_yaml_list[0]['mitre']['techniqueIds'].append(technique_id)
+                                                if references not in sigma_yaml_list[0]['mitre']['references']:
+                                                    sigma_yaml_list[0]['mitre']['references'].append(references)
+                                with open(os.path.join(sigmainput, filename), 'w') as f:
+                                    yaml.dump(sigma_yaml_list[0], f)
+                            except:
+                                click.secho('Unable to parse {}'.format(filename))
+
+
     click.echo('Results:')
 
     for result in results:
@@ -72,4 +108,4 @@ def cli(sigmainput, directory, method):
 
     click.echo('Total Valid Rule Files: {}'.format(str(len(results) - invalid_count) + "/" + str(len(results))))
     click.echo('Total Invalid Rule Files: {}'.format(str(invalid_count) + "/" + str(len(results))))
-    click.echo('Total Unsupported Rule Files (Multi-document): {}'.format(str(unsupported_count) + "/" + str(len(results))))
+    click.echo('Total Unsupported Rule Files (Multi-document): {}'.format(str(unsupported_count) + "/" + str(len(results))))