Merge pull request #30 from deovan/feature/add-headers

.gitignore

@@ -7,3 +7,6 @@ dupsko.yaml
 dist
 /.mypy_cache/
 .DS_Store
+.idea
+flow*
+swagger

README.md

@@ -109,6 +109,7 @@ To create a specification by inspecting HTTP traffic you will need to:
     Please note that mitmproxy2swagger will not overwrite existing endpoint descriptions, if you want to overwrite them, you can delete them before running the second pass.
 
     Passing `--examples` will add example data to requests and responses. Take caution when using this option, as it may add sensitive data (tokens, passwords, personal information etc.) to the schema.
+    Passing `--headers` will add headers data to requests and responses. Take caution when using this option, as it may add sensitive data (tokens, passwords, personal information etc.) to the schema.
 
 ### HAR
 

mitmproxy2swagger/mitmproxy_to_swagger.py

@@ -2,23 +2,21 @@
 """
 Converts a mitmproxy dump file to a swagger schema.
 """
+import argparse
+import json
 import os
+import re
 import sys
 import traceback
-from mitmproxy.exceptions import FlowReadException
-import json
-import argparse
-import ruamel.yaml
-import re
-from . import swagger_util
-from .har_capture_reader import HarCaptureReader, har_archive_heuristic
-from .mitmproxy_capture_reader import (
-    MitmproxyCaptureReader,
-    mitmproxy_dump_file_huristic,
-)
-from . import console_util
 import urllib
 
+import ruamel.yaml
+from mitmproxy.exceptions import FlowReadException
+
+from mitmproxy2swagger import console_util, swagger_util
+from mitmproxy2swagger.har_capture_reader import har_archive_heuristic, HarCaptureReader
+from mitmproxy2swagger.mitmproxy_capture_reader import mitmproxy_dump_file_huristic, MitmproxyCaptureReader
+
 
 def path_to_regex(path):
     # replace the path template with a regex
@@ -85,6 +83,12 @@ def main():
         action="store_true",
         help="Include examples in the schema. This might expose sensitive information.",
     )
+    parser.add_argument(
+        "-hd",
+        "--headers",
+        action="store_true",
+        help="Include headers in the schema. This might expose sensitive information.",
+    )
     parser.add_argument(
         "-f",
         "--format",
@@ -107,7 +111,10 @@ def main():
 
     # try loading the existing swagger file
     try:
-        with open(args.output, "r") as f:
+        base_dir = os.getcwd()
+        relative_path = args.output
+        abs_path = os.path.join(base_dir, relative_path)
+        with open(abs_path, "r") as f:
             swagger = yaml.load(f)
     except FileNotFoundError:
         print("No existing swagger file found. Creating new one.")
@@ -192,7 +199,14 @@ def main():
             )
 
             params = swagger_util.url_to_params(url, path_template_to_set)
-
+            if args.headers:
+                headers_request = swagger_util.request_to_headers(f.get_request_headers())
+                if headers_request is not None and len(headers_request) > 0:
+                    set_key_if_not_exists(
+                        swagger["paths"][path_template_to_set][method],
+                        "parameters",
+                        headers_request
+                    )
             if params is not None and len(params) > 0:
                 set_key_if_not_exists(
                     swagger["paths"][path_template_to_set][method], "parameters", params
@@ -261,6 +275,7 @@ def main():
                 if response_json is not None:
                     resp_data_to_set = {
                         "description": f.get_response_reason(),
+                        "headers": None,
                         "content": {
                             "application/json": {
                                 "schema": swagger_util.value_to_schema(response_json)
@@ -271,6 +286,9 @@ def main():
                         resp_data_to_set["content"]["application/json"][
                             "example"
                         ] = swagger_util.limit_example_size(response_json)
+                    if args.headers:
+                        resp_data_to_set["headers"] = swagger_util.response_to_headers(f.get_response_headers())
+
                     set_key_if_not_exists(
                         swagger["paths"][path_template_to_set][method]["responses"],
                         str(status),

mitmproxy2swagger/swagger_util.py

@@ -78,6 +78,39 @@ def url_to_params(url, path_template):
     return params
 
 
+# when given an url and its path template, generates the parameters section of the request
+def request_to_headers(headers):
+    header = []
+    if headers:
+        for key in headers:
+            header.append(
+                {
+                    "name": key,
+                    "value": headers[key][0],
+                    "default": headers[key][0],
+                    "in": "header",
+                    "required": True,
+                    "schema": {
+                        "type": "number" if headers[key][0].isdigit() else "string"
+                    },
+                }
+            )
+    return header
+
+
+def response_to_headers(headers):
+    header = {}
+    if headers:
+        for key in headers:
+            header[key] = {
+                "description": headers[key][0],
+                "schema": {
+                    "type": "number" if headers[key][0].isdigit() else "string"
+                },
+            }
+    return header
+
+
 def value_to_schema(value):
     # check if value is a number
     if type(value) == int or type(value) == float: