diff --git a/bin/http_request_translator b/bin/http_request_translator index 0c683f0..61fa08f 100755 --- a/bin/http_request_translator +++ b/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') diff --git a/http_request_translator/base.py b/http_request_translator/base.py index 3fdd306..827ddfd 100644 --- a/http_request_translator/base.py +++ b/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)) diff --git a/http_request_translator/plugin_manager.py b/http_request_translator/plugin_manager.py index 234496d..c346e8d 100644 --- a/http_request_translator/plugin_manager.py +++ b/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() diff --git a/http_request_translator/translator.py b/http_request_translator/translator.py index 30f07c9..f903300 100644 --- a/http_request_translator/translator.py +++ b/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: diff --git a/http_request_translator/util.py b/http_request_translator/util.py index 12d3096..1a66e13 100644 --- a/http_request_translator/util.py +++ b/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) \ No newline at end of file