fixed file fuzzing routine

Daniele Linguaglossa · Daniele Linguaglossa · commit f5cbf5a8acff · 2017-03-30T01:43:16.000+02:00
diff --git a/pyjfuzz/core/errors/__init__.py b/pyjfuzz/core/errors/__init__.py
@@ -56,6 +56,12 @@ class PJFInvalidArgument(PJFBaseException):
     """
     err_type = "INVALID ARGUMENT"
 
+class PJFInvalidJSON(PJFBaseException):
+    """
+    Invalid argument passed to PyJFuzz
+    """
+    err_type = "INVALID JSON"
+
 class PJFSocketError(PJFBaseException):
     """
     Socket issue
@@ -90,4 +96,13 @@ class PJFSocketPortInUse(PJFSocketError):
 class PJFProcessExecutionError(PJFProcessError):
     """
     Error during process execution
-    """
+    """
+
+class PJFMalformedJSON(PJFInvalidJSON):
+    """
+    Invalid argument passed to PyJFuzz
+    """
+    err_type = "MALFORMED JSON"
+
+    def __init__(self):
+        self.message = "Invalid JSON object"
diff --git a/pyjfuzz/core/pjf_worker.py b/pyjfuzz/core/pjf_worker.py
@@ -28,10 +28,10 @@
 from pjf_factory import PJFFactory
 from pjf_process_monitor import PJFProcessMonitor
 from pjf_external_fuzzer import PJFExternalFuzzer
+from errors import PJFMalformedJSON
 from argparse import Namespace
-from ast import literal_eval
 import tempfile
-import json
+import json as json_eval
 import netifaces
 import time
 import sys
@@ -127,23 +127,22 @@ def start_process_monitor(self):
             raise PJFBaseException(e.message)
 
     def start_file_fuzz(self):
-        try:
             with open(self.config.json_file, "rb") as json_file:
-                if not self.config.strong_fuzz:
-                    setattr(self, "json", literal_eval(json_file.read()))
-                    json = PJFFactory(self.config)
-                else:
-                    try:
-                        setattr(self, "json", literal_eval(json_file.read()))
-                    except:
-                        json_file.seek(0)
-                        setattr(self, "json", json_file.read())
-                    json = PJFFactory(self.config)
-                    json_file.close()
-            with open(self.config.json_file, "wb") as json_file:
-                json_file.write(json.fuzzed)
-        except Exception as e:
-            raise PJFBaseException(e.message)
+                j = json_file.read()
+                json = None
+                try:
+                    if not self.config.strong_fuzz:
+                        setattr(self.config, "json", json_eval.loads(j))
+                        json = PJFFactory(self.config)
+                    else:
+                        setattr(self.config, "json", json_eval.loads(j))
+                        json = PJFFactory(self.config)
+                except:
+                    raise PJFMalformedJSON()
+                json_file.close()
+            if json:
+                with open(self.config.json_file, "wb") as json_file:
+                    json_file.write(json.fuzzed)
 
     def start_http_server(self):
         try:
@@ -160,7 +159,7 @@ def start_http_server(self):
     def fuzz_command_line(self):
         try:
             with tempfile.NamedTemporaryFile(delete=False) as temp_file:
-                temp_file.write(json.dumps(self.config.json))
+                temp_file.write(json_eval.dumps(self.config.json))
                 temp_file.close()
                 setattr(self, "temp_file_name", temp_file.name)
                 if self.config.debug:
@@ -174,7 +173,7 @@ def fuzz_command_line(self):
 
     def fuzz_stdin(self):
         try:
-            result = PJFExternalFuzzer(self.config).execute(json.dumps(self.config.json))
+            result = PJFExternalFuzzer(self.config).execute(json_eval.dumps(self.config.json))
             if result:
                 sys.stdout.write(result)
             else:
