From e84dcd98f1188bb89baa70654babf0383e445f01 Mon Sep 17 00:00:00 2001
From: Alexandre Dulaunoy <a@foo.be>
Date: Sat, 29 May 2021 11:38:41 +0200
Subject: [PATCH] new: [backend] add reference to namespace id towards UUID

Such as CAPEC, MITRE-ATTACK-ID and what ever you like.

Data structure updated and MITRE ATT&CK crawler support
---
 backend/data-structure.md         | 12 ++++++++++++
 crawler/mitre-cti/cti-importer.py | 18 ++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/backend/data-structure.md b/backend/data-structure.md
index 65a5bdc..a177806 100644
--- a/backend/data-structure.md
+++ b/backend/data-structure.md
@@ -57,6 +57,18 @@ The relationship data from an UUID tuple.
 
 - `rd:<source UUID>:<destination UUID>` -> {`value`, `value`}
 
+# id:<NAMESPACE>:<NAMESPACE_ID> (set)
+
+A namespace id referenced in an UUID
+
+- `id:<NAMESPACE>:<NAMESPACE_ID>` -> {`UUID`, `UUID`}
+
+# idk:<NAMESPACE> (set)
+
+Known id per namespace
+
+- `idk:<NAMESPACE>' -> {`NAMESPACE_ID`, `NAMESPACE_ID`}
+
 # Statistics
 
 ## Automatic API statistics
diff --git a/crawler/mitre-cti/cti-importer.py b/crawler/mitre-cti/cti-importer.py
index 0b0e083..5ea9f63 100644
--- a/crawler/mitre-cti/cti-importer.py
+++ b/crawler/mitre-cti/cti-importer.py
@@ -28,6 +28,10 @@ def additem(uuidref=None, data=None, project=None):
     if project is not None:
         rdb.sadd("parent:{}".format(uuidref), project)
         rdb.sadd("child:{}".format(project), uuidref)
+    if 'capec' in data:
+        addexternalid(uuidsource=uuidref, namespace='capec', namespaceid=data['capec'])
+    if 'mitre-attack-id' in data:
+        addexternalid(uuidsource=uuidref, namespace='mitre-attack-id', namespaceid=data['mitre-attack-id'])
     return True
 
 def addrelationship(uuidsource=None, uuiddest=None, data=None):
@@ -37,6 +41,14 @@ def addrelationship(uuidsource=None, uuiddest=None, data=None):
     rdb.sadd("rd:{}:{}".format(uuidsource, uuiddest), data)
     return True
 
+def addexternalid(uuidsource=None, namespace=None, namespaceid=None):
+    if uuidsource is None or namespace is None or namespaceid is None:
+        return None
+    k = "id:{}:{}".format(namespace.lower(), namespaceid)
+    rdb.sadd(k, uuidsource)
+    k = "idk:{}".format(namespace)
+    rdb.sadd(k, namespaceid)
+
 models = ['enterprise-attack', 'mobile-attack', 'ics-attack', 'pre-attack']
 
 for model in models:
@@ -53,6 +65,12 @@ for model in models:
                 data['mitre-cti:description'] = obj['description']
             if 'name' in obj:
                 data['mitre-cti:name'] = obj['name']
+            if 'external_references' in obj:
+                for ref in obj['external_references']:
+                    if ref['source_name'] == 'mitre-attack':
+                        data['mitre-attack-id'] = ref['external_id']
+                    if ref['source_name'] == 'capec':
+                        data['capec'] = ref['external_id']
             additem(uuidref=obj_id, project=projectuuid, data=data)
         elif obj_type == 'relationship':
             (source_type, source_id) = obj['source_ref'].split('--')