Add custom template option

bin/http_request_translator

@@ -38,7 +38,10 @@ def take_arguments():
         help="Interactive mode: read raw HTTP request from keyboard.")
     parser.add_argument(
         '--data', '-d',
-        help='Add the data that you want to send along with the header')
+        help='Add the data that you want to send along with the header.')
+    parser.add_argument(
+        '--custom-template', '-t',
+        help='Add the path of the custom template you want to provide.')
     request_group.add_argument(
         '--request', '-r',
         help='Input the HTTP request')

http_request_translator/base.py

@@ -11,6 +11,7 @@
 from importlib import import_module
 
 from .url import get_url, check_valid_url
+from .util import exists_global_custom_template, get_global_custom_template
 
 
 class AbstractScript(object):
@@ -27,16 +28,17 @@ class AbstractScript(object):
     code_search = ''
     code_nosearch = ''
 
-    def __init__(self, headers=None, details=None, search=None):
+    def __init__(self, headers=None, details=None, search=None, template=None):
         """Initialize the script generation.
 
         :param list headers: Headers list containing fields like 'Host', 'User-Agent', etc.
         :param dict details: Request specific details dictionary like body and method of the request.
         :param str search: String to search for in the response to the request.
+        :param str template: Custom template path in dotted module format.
 
         :raises ValueError: When url is invalid.
         """
-        self.load_attributes(self.__class__)
+        self.load_attributes(self.__class__, template)
         self._script = ''
         self.headers = headers
         self.details = details
@@ -190,21 +192,27 @@ def encode_url(self, url):
         return encoded_url
 
     @staticmethod
-    def load_attributes(cls):
+    def load_attributes(cls, template=None):
         """Loads attributes to Script class from a given script's template
 
         Imports the template file/module, assigns all the attributes defined in the template file to the given class.
 
         :param class cls: Script class to which template is to be loaded.
+        :param str template: Custom template path in dotted module format.
 
         :raises AttributeError: When __language__ attribute is not present.
         """
-        templates_path = "{}.templates".format(__name__.split('.', 1)[0])
-        if not hasattr(cls, '__language__'):
-            raise AttributeError("__language__ not found in class: {}, attributes cannot be loaded".format(cls.__name__))
-        template = import_module("{templates_path}.{class_template}".format(
-            templates_path=templates_path,
-            class_template=cls.__language__))
+        if template:
+            template = import_module(template)
+        elif exists_global_custom_template(cls.__language__):
+            template = get_global_custom_template(cls.__language__)
+        else:
+            templates_path = "{}.templates".format(__name__.split('.', 1)[0])
+            if not hasattr(cls, '__language__'):
+                raise AttributeError("__language__ not found in class: {}, attributes cannot be loaded".format(cls.__name__))
+            template = import_module("{templates_path}.{class_template}".format(
+                templates_path=templates_path,
+                class_template=cls.__language__))
         attributes = (var for var in vars(template) if var.startswith('code_'))
         for attr in attributes:
             setattr(cls, attr, getattr(template, attr))

http_request_translator/plugin_manager.py

@@ -22,7 +22,7 @@ def get_script_class(script_name):
     raise ValueError("The {} language is not supported.".format(script_name))
 
 
-def generate_script(script, headers, details, search_string=None):
+def generate_script(script, headers, details, search_string=None, template=None):
     """Returns the script code for the HTTP request passed in script language
 
     :param str script: Name of the language for which script is to be generated
@@ -34,4 +34,4 @@ def generate_script(script, headers, details, search_string=None):
     :rtype: `str`
     """
     class_script = get_script_class(script.strip().lower())
-    return class_script(headers=headers, details=details, search=search_string).generate_script()
+    return class_script(headers=headers, details=details, search=search_string, template=template).generate_script()

http_request_translator/translator.py

@@ -72,7 +72,7 @@ def process_arguments(args):
         elif args.search_regex:
             arg_option = args.search_regex
         if len(script_list) == 0:
-            generated_code = generate_script('bash', headers, details, arg_option)
+            generated_code = generate_script('bash', headers, details, arg_option, args.custom_template)
             print(generated_code)
         else:
             for script in script_list:

http_request_translator/util.py

@@ -1,4 +1,7 @@
+from importlib import import_module
+import os
 import re
+import sys
 
 
 # Blindy copied from: https://gist.github.com/mnordhoff/2213179
@@ -7,3 +10,15 @@
 re_ipv6_address = re.compile('^(?:(?:[0-9A-Fa-f]{1,4}:){6}(?:[0-9A-Fa-f]{1,4}:[0-9A-Fa-f]{1,4}|(?:(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){3}(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5]))|::(?:[0-9A-Fa-f]{1,4}:){5}(?:[0-9A-Fa-f]{1,4}:[0-9A-Fa-f]{1,4}|(?:(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){3}(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5]))|(?:[0-9A-Fa-f]{1,4})?::(?:[0-9A-Fa-f]{1,4}:){4}(?:[0-9A-Fa-f]{1,4}:[0-9A-Fa-f]{1,4}|(?:(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){3}(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5]))|(?:[0-9A-Fa-f]{1,4}:[0-9A-Fa-f]{1,4})?::(?:[0-9A-Fa-f]{1,4}:){3}(?:[0-9A-Fa-f]{1,4}:[0-9A-Fa-f]{1,4}|(?:(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){3}(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5]))|(?:(?:[0-9A-Fa-f]{1,4}:){,2}[0-9A-Fa-f]{1,4})?::(?:[0-9A-Fa-f]{1,4}:){2}(?:[0-9A-Fa-f]{1,4}:[0-9A-Fa-f]{1,4}|(?:(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){3}(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5]))|(?:(?:[0-9A-Fa-f]{1,4}:){,3}[0-9A-Fa-f]{1,4})?::[0-9A-Fa-f]{1,4}:(?:[0-9A-Fa-f]{1,4}:[0-9A-Fa-f]{1,4}|(?:(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){3}(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5]))|(?:(?:[0-9A-Fa-f]{1,4}:){,4}[0-9A-Fa-f]{1,4})?::(?:[0-9A-Fa-f]{1,4}:[0-9A-Fa-f]{1,4}|(?:(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){3}(?:[0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5]))|(?:(?:[0-9A-Fa-f]{1,4}:){,5}[0-9A-Fa-f]{1,4})?::[0-9A-Fa-f]{1,4}|(?:(?:[0-9A-Fa-f]{1,4}:){,6}[0-9A-Fa-f]{1,4})?::)$')
 # Homebrew
 re_domain = re.compile(r'^(?:(?:[A-Z](?:[A-Z-]{0,61}[A-Z])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}\.?)|)$', re.IGNORECASE)
+
+DEFAULT_CUSTOM_TEMPLATE_PATH = os.environ.get('DEFAULT_CUSTOM_TEMPLATE_PATH', '~/.config/translator/templates')
+DEFAULT_CUSTOM_TEMPLATE_PATH = DEFAULT_CUSTOM_TEMPLATE_PATH.replace('~', os.environ.get('HOME', ''))
+
+
+def exists_global_custom_template(language):
+    return os.path.exists(os.path.join(DEFAULT_CUSTOM_TEMPLATE_PATH, "{}.py".format(language)))
+
+
+def get_global_custom_template(language):
+    sys.path.insert(0, DEFAULT_CUSTOM_TEMPLATE_PATH)
+    return import_module(language)