From 065865ef3657871d88ce469faa1c212b013c3123 Mon Sep 17 00:00:00 2001
From: alufers <alufers@wp.pl>
Date: Sat, 20 Aug 2022 11:25:14 +0200
Subject: [PATCH] fix: add error message and tune the mitmproxy/har detection
 logic

---
 mitmproxy2swagger/console_util.py             |  1 +
 mitmproxy2swagger/har_capture_reader.py       |  6 ++++--
 mitmproxy2swagger/mitmproxy2swagger.py        | 20 +++++++++++++++++--
 mitmproxy2swagger/mitmproxy_capture_reader.py |  2 ++
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/mitmproxy2swagger/console_util.py b/mitmproxy2swagger/console_util.py
index 62456f8..f1dda76 100644
--- a/mitmproxy2swagger/console_util.py
+++ b/mitmproxy2swagger/console_util.py
@@ -2,6 +2,7 @@ import sys
 
 ANSI_RGB = "\033[38;2;{};{};{}m"
 ANSI_RGB_BG = "\033[48;2;{};{};{}m"
+ANSI_RED = "\033[31m"
 ANSI_RESET = "\033[0m"
 
 RAINBOW_COLORS = [
diff --git a/mitmproxy2swagger/har_capture_reader.py b/mitmproxy2swagger/har_capture_reader.py
index 1277ec5..4dff62c 100644
--- a/mitmproxy2swagger/har_capture_reader.py
+++ b/mitmproxy2swagger/har_capture_reader.py
@@ -14,10 +14,10 @@ def har_archive_heuristic(file_path: str) -> int:
         data = f.read(2048)
         # if file contains only ascii characters
         if data.decode('utf-8', 'ignore').isprintable() is True:
-            val += 40
+            val += 25
         # if first character is a '{'
         if data[0] == '{':
-            val += 15
+            val += 23
         # if it contains the word '"WebInspector"'
         if b'"WebInspector"' in data:
             val += 15
@@ -89,3 +89,5 @@ class HarCaptureReader:
                 if self.progress_callback:
                     self.progress_callback(f.tell() / har_file_size)
                 yield HarFlowWrapper(entry)
+    def name(self):
+        return 'har'
diff --git a/mitmproxy2swagger/mitmproxy2swagger.py b/mitmproxy2swagger/mitmproxy2swagger.py
index 6ff429a..80968fe 100755
--- a/mitmproxy2swagger/mitmproxy2swagger.py
+++ b/mitmproxy2swagger/mitmproxy2swagger.py
@@ -2,6 +2,9 @@
 """
 Converts a mitmproxy dump file to a swagger schema.
 """
+import os
+import sys
+import traceback
 from mitmproxy.exceptions import FlowReadException
 import json
 import argparse
@@ -37,7 +40,12 @@ def progress_callback(progress):
 
 
 def detect_input_format(file_path):
-    if har_archive_heuristic(file_path) > mitmproxy_dump_file_huristic(file_path):
+    har_score = har_archive_heuristic(file_path)
+    mitmproxy_score =  mitmproxy_dump_file_huristic(file_path)
+    if 'MITMPROXY2SWAGGER_DEBUG' in os.environ:
+        print('har score: ' + str(har_score))
+        print('mitmproxy score: ' + str(mitmproxy_score))
+    if har_score > mitmproxy_score:
         return HarCaptureReader(file_path, progress_callback)
     return MitmproxyCaptureReader(file_path, progress_callback)
 
@@ -59,7 +67,7 @@ def main():
     yaml = ruamel.yaml.YAML()
 
     caputre_reader = None
-    if args.format == 'flow':
+    if args.format == 'flow' or args.format == 'mitmproxy':
         caputre_reader = MitmproxyCaptureReader(args.input, progress_callback)
     elif args.format == 'har':
         caputre_reader = HarCaptureReader(args.input, progress_callback)
@@ -204,6 +212,14 @@ def main():
 
     except FlowReadException as e:
         print(f"Flow file corrupted: {e}")
+    except ValueError as e:
+        print(f"ValueError: {e}")
+        # print stack trace
+        traceback.print_exception(*sys.exc_info())
+        print(f"{console_util.ANSI_RED}Failed to parse the input file as '{caputre_reader.name()}'. ")
+        if not args.format:
+            print(f"It might happen that the input format as incorrectly detected. Please try using '--fromat flow' or '--format har' to specify the input format.{console_util.ANSI_RESET}")
+        sys.exit(1)
 
     new_path_templates.sort()
 
diff --git a/mitmproxy2swagger/mitmproxy_capture_reader.py b/mitmproxy2swagger/mitmproxy_capture_reader.py
index a2b3abb..a356d45 100644
--- a/mitmproxy2swagger/mitmproxy_capture_reader.py
+++ b/mitmproxy2swagger/mitmproxy_capture_reader.py
@@ -86,3 +86,5 @@ class MitmproxyCaptureReader:
                         yield MitmproxyFlowWrapper(f)
             except FlowReadException as e:
                 print(f"Flow file corrupted: {e}")
+    def name(self):
+        return "flow"