From 52e079fea2923d6a5bff583dea2f5efa9bfda177 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Sat, 9 Dec 2017 13:12:04 +0100
Subject: [PATCH] chg: Allow to pass a pseudofile to LIEF

---
 pymisp/tools/create_misp_object.py |  7 +++++--
 tests/test_offline.py              | 29 +++++++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/pymisp/tools/create_misp_object.py b/pymisp/tools/create_misp_object.py
index 95a43f1..c6a4dea 100644
--- a/pymisp/tools/create_misp_object.py
+++ b/pymisp/tools/create_misp_object.py
@@ -49,9 +49,12 @@ def make_macho_objects(lief_parsed, misp_file):
 
 def make_binary_objects(filepath=None, pseudofile=None, filename=None):
     misp_file = FileObject(filepath=filepath, pseudofile=pseudofile, filename=filename)
-    if HAS_LIEF and filepath:
+    if HAS_LIEF and filepath or (pseudofile and filename):
         try:
-            lief_parsed = lief.parse(filepath)
+            if filepath:
+                lief_parsed = lief.parse(filepath=filepath)
+            else:
+                lief_parsed = lief.parse(raw=pseudofile.getvalue(), name=filename)
             if isinstance(lief_parsed, lief.PE.Binary):
                 return make_pe_objects(lief_parsed, misp_file)
             elif isinstance(lief_parsed, lief.ELF.Binary):
diff --git a/tests/test_offline.py b/tests/test_offline.py
index 992fb80..b4c1644 100644
--- a/tests/test_offline.py
+++ b/tests/test_offline.py
@@ -5,6 +5,7 @@ import unittest
 import requests_mock
 import json
 import os
+from io import BytesIO
 
 import pymisp as pm
 from pymisp import PyMISP
@@ -210,9 +211,12 @@ class TestOffline(unittest.TestCase):
         p.add_internal_other(evt, 'foobar')
         p.add_attachment(evt, "testFile")
 
-    def make_objects(self, path):
+    def make_objects(self, path=None, pseudofile=None, filename=None):
         to_return = {'objects': [], 'references': []}
-        fo, peo, seos = make_binary_objects(path)
+        if path:
+            fo, peo, seos = make_binary_objects(path)
+        else:
+            fo, peo, seos = make_binary_objects(pseudofile=pseudofile, filename=filename)
 
         if seos:
             for s in seos:
@@ -229,8 +233,29 @@ class TestOffline(unittest.TestCase):
             to_return['objects'].append(fo)
             if fo.ObjectReference:
                 to_return['references'] += fo.ObjectReference
+
+        # Remove UUIDs for comparing the objects.
+        for o in to_return['objects']:
+            o.pop('uuid')
+        for o in to_return['references']:
+            o.pop('referenced_uuid')
+            o.pop('object_uuid')
         return json.dumps(to_return, cls=MISPEncode)
 
+    def test_objects_pseudofile(self, m):
+        paths = ['cmd.exe', 'tmux', 'MachO-OSX-x64-ls']
+        try:
+            for path in paths:
+                with open(os.path.join('tests', 'viper-test-files', 'test_files', path), 'rb') as f:
+                    pseudo = BytesIO(f.read())
+                    json_blob = self.make_objects(pseudofile=pseudo, filename=path)
+                # Compare pseudo file / path
+                filepath_blob = self.make_objects(os.path.join('tests', 'viper-test-files', 'test_files', path))
+                self.assertEqual(json_blob, filepath_blob)
+        except IOError:  # Can be replaced with FileNotFoundError when support for python 2 is dropped
+            return unittest.SkipTest()
+        print(json_blob)
+
     def test_objects(self, m):
         paths = ['cmd.exe', 'tmux', 'MachO-OSX-x64-ls']
         try: