From eddb0008901a5751335df3288fa4f27d341ea752 Mon Sep 17 00:00:00 2001
From: Terrtia <or1994@hotmail.fr>
Date: Wed, 12 Feb 2020 17:12:17 +0100
Subject: [PATCH] chg: [MISP import-export] screenshots with relationships

---
 bin/export/MispImport.py    |  4 ++-
 bin/lib/Correlate_object.py |  2 +-
 bin/lib/Screenshot.py       | 55 +++++++++++++++++++++++++++++++++----
 3 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/bin/export/MispImport.py b/bin/export/MispImport.py
index 4c40be86..a1a08727 100755
--- a/bin/export/MispImport.py
+++ b/bin/export/MispImport.py
@@ -159,7 +159,7 @@ def unpack_file(map_uuid_global_id, misp_obj):
         if obj_id and io_content:
             obj_meta = get_object_metadata(misp_obj)
             if obj_type == 'screenshot':
-                #Screenshot.create_screenshot(obj_id, obj_meta, io_content)
+                Screenshot.create_screenshot(obj_id, obj_meta, io_content)
                 pass
             else: #decoded
                 Decoded.create_decoded(obj_id, obj_meta, io_content)
@@ -214,5 +214,7 @@ if __name__ == '__main__':
 
     import_objs_from_file('test_import_item.json')
 
+    #Screenshot.delete_screenshot('d4065d632a232a323d33a30144924763ae229a972c363e12abf48009017ec10c')
+
     #Decoded.delete_correlation('23a44cc266880d26386a0a77318afbe09696f935')
     #Pgp.pgp.delete_correlation('key', '0xA4BB02A75E6AF448')
diff --git a/bin/lib/Correlate_object.py b/bin/lib/Correlate_object.py
index 45706d6f..2c971f07 100755
--- a/bin/lib/Correlate_object.py
+++ b/bin/lib/Correlate_object.py
@@ -207,7 +207,7 @@ def create_obj_relationship(obj1_type, obj1_id, obj2_type, obj2_id, obj1_subtype
     if obj1_type == 'domain':
         pass
     elif obj1_type == 'item':
-        pass
+        pass # son/father + duplicate
     elif obj1_type == 'pgp':
         Pgp.pgp.save_obj_relationship(obj1_subtype, obj1_id, obj2_type, obj2_id)
     elif obj1_type == 'cryptocurrency':
diff --git a/bin/lib/Screenshot.py b/bin/lib/Screenshot.py
index 326e5368..7491efb9 100755
--- a/bin/lib/Screenshot.py
+++ b/bin/lib/Screenshot.py
@@ -114,7 +114,7 @@ def get_screenshot_correlated_object(sha256_string, correlation_objects=[]):
     :return: a dict of all correlation for a given sha256
     :rtype: dict
     '''
-    if correlation_objects is None:
+    if not correlation_objects:
         correlation_objects = Correlate_object.get_all_correlation_objects()
     decoded_correlation = {}
     for correlation_object in correlation_objects:
@@ -128,6 +128,19 @@ def get_screenshot_correlated_object(sha256_string, correlation_objects=[]):
             decoded_correlation[correlation_object] = res
     return decoded_correlation
 
+def save_item_relationship(obj_id, item_id):
+    r_serv_metadata.hset('paste_metadata:{}'.format(item_id), 'screenshot', obj_id)
+    r_serv_onion.sadd('screenshot:{}'.format(obj_id), item_id)
+
+def save_domain_relationship(obj_id, domain):
+    r_serv_onion.sadd('domain_screenshot:{}'.format(domain), obj_id)
+    r_serv_onion.sadd('screenshot_domain:{}'.format(obj_id), domain)
+
+def save_obj_relationship(obj_id, obj2_type, obj2_id):
+    if obj2_type == 'domain':
+        save_domain_relationship(obj_id, obj2_id)
+    elif obj2_type == 'item':
+        save_item_relationship(obj_id, obj2_id)
 
 def get_screenshot_file_content(sha256_string):
     filepath = get_screenshot_filepath(sha256_string)
@@ -145,14 +158,44 @@ def save_screenshot_file(sha256_string, io_content):
         f.write(io_content.getvalue())
     return True
 
-def create_screenshot(sha256_string, obj_meta, io_content):
-    # check if sha256
-    res = save_screenshot_file(sha256_string, io_content)
+def delete_screenshot_file(obj_id):
+    filepath = get_screenshot_filepath(obj_id)
+    if not os.path.isfile(filepath):
+        return False
+    Tag.delete_obj_tags(obj_id, 'image', Tag.get_obj_tag(obj_id))
+    os.remove(filepath)
+    return True
+
+def create_screenshot(obj_id, obj_meta, io_content):
+    print(obj_id)
+    # # TODO: check if sha256
+    res = save_screenshot_file(obj_id, io_content)
     if res:
         # creata tags
-        if 'tags' in obj_metadata:
+        if 'tags' in obj_meta:
             # # TODO: handle mixed tags: taxonomies and Galaxies
-            Tag.api_add_obj_tags(tags=obj_metadata['tags'], object_id=obj_id, object_type="image")
+            Tag.api_add_obj_tags(tags=obj_meta['tags'], object_id=obj_id, object_type="image")
         return True
 
     return False
+
+def delete_screenshot(obj_id):
+    if not exist_screenshot(obj_id):
+        return False
+
+    res = delete_screenshot_file(obj_id)
+    if not res:
+        return False
+
+    obj_correlations = get_screenshot_correlated_object(obj_id)
+    if 'domain' in obj_correlations:
+        for domain in obj_correlations['domain']:
+            r_serv_onion.srem('domain_screenshot:{}'.format(domain), obj_id)
+        r_serv_onion.delete('screenshot_domain:{}'.format(obj_id))
+
+    if 'paste' in obj_correlations: # TODO: handle item
+        for item_id in obj_correlations['paste']:
+            r_serv_metadata.hdel('paste_metadata:{}'.format(item_id), 'screenshot')
+        r_serv_onion.sadd('screenshot:{}'.format(obj_id), item_id)
+
+    return True