From c894e079b8f7d099ecba910cbf7025bf0d22d654 Mon Sep 17 00:00:00 2001 From: Omar Abdelhafith Date: Sun, 28 Jun 2015 05:55:24 +0100 Subject: [PATCH] Create proxy mapper to extract the mapping logic to its own class Adding watchdog to watch the folder and rebuild the mapping index --- mockpy/core/cherrypy_mapper.py | 20 ++--- mockpy/core/proxy_mapper.py | 90 +++++++++++++++++++ mockpy/core/proxy_server.py | 89 ++---------------- mockpy/models/mapping_items_manager.py | 50 +++++++++-- mockpy/status/status.py | 2 +- mockpy/utils/cherrypy_extensions.py | 18 ++++ mockpy/utils/log.py | 2 +- .../{extensions.py => proxy_extensions.py} | 8 +- requirements.txt | 3 +- setup.py | 3 +- tests/test_mapping.py | 8 -- 11 files changed, 177 insertions(+), 116 deletions(-) create mode 100644 mockpy/core/proxy_mapper.py create mode 100644 mockpy/utils/cherrypy_extensions.py rename mockpy/utils/{extensions.py => proxy_extensions.py} (99%) diff --git a/mockpy/core/cherrypy_mapper.py b/mockpy/core/cherrypy_mapper.py index 96295e4..84f6acb 100644 --- a/mockpy/core/cherrypy_mapper.py +++ b/mockpy/core/cherrypy_mapper.py @@ -4,6 +4,7 @@ from mockpy.utils import log from mockpy.utils.config import * from mockpy.status.status import Status +import mockpy.utils.cherrypy_extensions class CherryPyMapper(object): @@ -15,14 +16,15 @@ def __init__(self, mapping_handler=None, cherrypy=None): def handle_request(self): - if Status.is_status(self.cherrypy.url()): + log.log_url(self.cherrypy.url()) + + if self.status.is_status(self.cherrypy.url()): info("Accessing Satus") - self.print_seperator() + log.print_seperator() return self.status.html_response() - request = MappingRequest(self.cherry_request_dict()) + request = self.cherrypy.to_mapper_request() items = self.mapping_handler.mapping_item_for_mapping_request(request) - log.log_url(self.cherrypy.url()) if len(items) == 0: self.cherrypy.response.status = 500 @@ -43,16 +45,6 @@ def handle_request(self): return response.body_response() - def cherry_request_dict(self): - dic = {"method": self.cherrypy.request.method, - "url": self.cherrypy.url(), - "headers": self.cherrypy.request.headers} - - if self.cherrypy.request.process_request_body: - dic["body"] = self.cherrypy.request.body.read() - - return dic - def fill_headers(self, headers): if type({}) is not type(headers): return diff --git a/mockpy/core/proxy_mapper.py b/mockpy/core/proxy_mapper.py new file mode 100644 index 0000000..215e82a --- /dev/null +++ b/mockpy/core/proxy_mapper.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import httplib +from threading import Thread + +from libmproxy.protocol.http import HTTPRequest, HTTPResponse +from netlib.odict import ODictCaseless + +import mockpy.utils.proxy_extensions +from ..utils.config import * +from mockpy.status.status import Status +from ..utils import log +from ..models.mapping_items_manager import * + +class ProxyMapper(object): + + def __init__(self, mapping_handler, http_proxy): + self.http_proxy = http_proxy + self.mapping_handler = mapping_handler + self.status = Status(self.mapping_handler) + + success("Proxy server started") + + def handle_request(self, flow): + log.log_url(flow.request.url) + + if self.status.is_status(flow.request.url): + info("Accessing Satus") + flow.reply(HTTPResponse.with_html(self.status.html_response())) + log.print_seperator() + return + + request = flow.request.to_mapper_request() + mapping_items = self.mapping_handler.mapping_item_for_mapping_request(request) + + if len(mapping_items) > 1: + log.log_multiple_matches(mapping_items) + + if len(mapping_items) == 0: + self.perform_http_request(flow) + else: + self.perform_mapping_request(flow, mapping_items[0]) + + log.print_seperator() + + def perform_mapping_request(self, flow, mapping_item): + response, request = mapping_item.response, mapping_item.request + + log.log_request(request) + log.log_response(response) + + response = HTTPResponse.from_intercepted_response(response) + flow.reply(response) + + def perform_http_request(self, flow): + if self.http_proxy is None: + flow.reply() + else: + thread = Thread(target=self.threaded_perform_http_request, + args=(flow, self.http_proxy)) + thread.start() + + def threaded_perform_http_request(self, flow, proxy_settings): + response = self.perform_http_connection(flow.request, proxy_settings[0], proxy_settings[1]) + flow.reply(response) + + @staticmethod + def perform_http_connection(request, url, port): + try: + conn = httplib.HTTPConnection(url, port) + headers = dict(request.headers.items()) + + conn.request(request.method, request.url, + body=request.content, headers=headers) + httplib_response = conn.getresponse() + + headers = ODictCaseless.from_httplib_headers(httplib_response.getheaders()) + response = HTTPResponse(code=httplib_response.status, + content=httplib_response.read(), + msg="", + httpversion=(1, 1), + headers=headers) + return response + except Exception as ex: + error("Error Happened") + error(ex) + error("method: %s\nurl: %s\nbody: --\nheaders: --" % + (request.method, request.url)) + return None diff --git a/mockpy/core/proxy_server.py b/mockpy/core/proxy_server.py index 1a98af7..5040a86 100755 --- a/mockpy/core/proxy_server.py +++ b/mockpy/core/proxy_server.py @@ -8,99 +8,28 @@ from libmproxy import controller, proxy from libmproxy.proxy.server import ProxyServer -from libmproxy.protocol.http import HTTPRequest, HTTPResponse -from netlib.odict import ODictCaseless -import mockpy.utils.extensions -from ..utils.config import * -from mockpy.status.status import Status from ..utils import log -from ..models.mapping_items_manager import * - +from .proxy_mapper import * class MITMProxy(controller.Master): - def __init__(self, server, inout_path, res_path, http_proxy): + def __init__(self, server, proxy_mapper): controller.Master.__init__(self, server) - - self.http_proxy = http_proxy - self.handler = MappingItemsManager(inout_path, res_path) - self.status = Status(self.handler) - - success("Proxy server started") + self.proxy_mapper = proxy_mapper def handle_request(self, flow): - request = flow.request.to_mapper_request() - mapping_items = self.handler.mapping_item_for_mapping_request(request) - - log.log_url(flow.request.url) - - if Status.is_status(flow.request.url): - info("Accessing Satus") - flow.reply(HTTPResponse.with_html(self.status.html_response())) - log.print_seperator() - return - - if len(mapping_items) > 1: - log.log_multiple_matches(mapping_items) - - if len(mapping_items) == 0: - self.perform_http_request(flow) - else: - self.perform_mapping_request(flow, mapping_items[0]) - - log.print_seperator() - - def perform_mapping_request(self, flow, mapping_item): - response, request = mapping_item.response, mapping_item.request - - log.log_request(request) - log.log_response(response) - - response = HTTPResponse.from_intercepted_response(response) - flow.reply(response) - - def perform_http_request(self, flow): - if self.http_proxy is None: - flow.reply() - else: - thread = Thread(target=self.threaded_perform_http_request, - args=(flow, self.http_proxy)) - thread.start() - - def threaded_perform_http_request(self, flow, proxy_settings): - response = self.perform_request(flow.request, proxy_settings[0], proxy_settings[1]) - flow.reply(response) - - @staticmethod - def perform_request(request, url, port): - try: - conn = httplib.HTTPConnection(url, port) - headers = dict(request.headers.items()) - - conn.request(request.method, request.url, - body=request.content, headers=headers) - httplib_response = conn.getresponse() - - headers = ODictCaseless.from_httplib_headers(httplib_response.getheaders()) - response = HTTPResponse(code=httplib_response.status, - content=httplib_response.read(), - msg="", - httpversion=(1, 1), - headers=headers) - return response - except Exception as ex: - error("Error Happened") - error(ex) - error("method: %s\nurl: %s\nbody: --\nheaders: --" % - (request.method, request.url)) - return None + self.proxy_mapper.handle_request(flow) def start_proxy_server(port, inout_path, res_path, http_proxy): config = proxy.ProxyConfig(port=port) server = ProxyServer(config) - m = MITMProxy(server, inout_path, res_path, http_proxy) + + mapping_handler = MappingItemsManager(inout_path, res_path) + + proxy_mapper = ProxyMapper(mapping_handler, http_proxy) + m = MITMProxy(server, proxy_mapper) def signal_handler(signal, frame): info("\nShutting down proxy server") diff --git a/mockpy/models/mapping_items_manager.py b/mockpy/models/mapping_items_manager.py index a7f111c..cbc3984 100644 --- a/mockpy/models/mapping_items_manager.py +++ b/mockpy/models/mapping_items_manager.py @@ -1,15 +1,29 @@ __author__ = 'omarsubhiabdelhafith' from .mapping_item import * +from mockpy.utils import log from functools import partial import re +from watchdog.observers import Observer +from watchdog.events import FileSystemEventHandler class MappingItemsManager(object): def __init__(self, inout_path, res_path): - self.yaml_files = get_yaml_files(inout_path) - create_mapping_item_with_yaml = partial(create_mapping_item, inout_path, res_path) - self.mappings = list(map(create_mapping_item_with_yaml, self.yaml_files)) + self.inout_path = inout_path + self.res_path = res_path + self.parse_inout_and_res() + self.install_watchers() + + def parse_inout_and_res(self): + self.yaml_files = get_yaml_files(self.inout_path) + self.mappings = list(map(self.create_mapping_item, self.yaml_files)) + + def create_mapping_item(self, yml_file): + full_path = self.inout_path + "/" + yml_file + + with open(full_path, "r") as file: + return MappingItem(yaml.load(file), full_path, self.res_path) def response_for_mapping_request(self, request): return [item.response for item in self.mappings if item.handles_mapping_request(request)] @@ -17,15 +31,35 @@ def response_for_mapping_request(self, request): def mapping_item_for_mapping_request(self, request): return [item for item in self.mappings if item.handles_mapping_request(request)] + def install_watchers(self): + self.event_handler = MapperDirectoryListener(self) + self.install_watcher(self.inout_path) + self.install_watcher(self.res_path) + + def install_watcher(self, path): + observer = Observer() + observer.schedule(self.event_handler, path, recursive=True) + observer.start() + + +class MapperDirectoryListener(FileSystemEventHandler): + + def __init__(self, mapping_manager): + self.mapping_manager = mapping_manager + + def on_any_event(self, event): + log.info("Inout or Res directory changed, rebuilding mapping settings\n" + "File path: %s" % event.src_path + + "\nEvent type: %s" % event.event_type) + + self.mapping_manager.parse_inout_and_res() + log.success("Mapping settings rebuilt successfully") + log.print_seperator() + def get_yaml_files(path): files = os.listdir(path) return filter(lambda file: re.match(".*\.yml$", file), files) -def create_mapping_item(inout_path, res_path, yml_file): - file_to_open = inout_path + "/" + yml_file - with open(file_to_open, "r") as file: - return MappingItem(yaml.load(file), file_to_open, res_path) - diff --git a/mockpy/status/status.py b/mockpy/status/status.py index 45d21bc..f013464 100644 --- a/mockpy/status/status.py +++ b/mockpy/status/status.py @@ -14,6 +14,7 @@ def html_response(self): string += "Parsed interceptors:
" string += "_" * 80 string += "
" + for item in self.mapping_handler.mappings: string += " - " + item.file_name + "
" @@ -27,7 +28,6 @@ def html_response(self): string += "_" * 80 string += "
" - string += "" string += "" return string diff --git a/mockpy/utils/cherrypy_extensions.py b/mockpy/utils/cherrypy_extensions.py new file mode 100644 index 0000000..8bad922 --- /dev/null +++ b/mockpy/utils/cherrypy_extensions.py @@ -0,0 +1,18 @@ +import cherrypy +from mockpy.models.mapping_request import * + + +""" + Extensions +""" +def to_mapper_request(): + dic = {"method": cherrypy.request.method, + "url": cherrypy.url(), + "headers": cherrypy.request.headers} + + if cherrypy.request.process_request_body: + dic["body"] = cherrypy.request.body.read() + + return MappingRequest(dic) + +setattr(cherrypy, "to_mapper_request", to_mapper_request) diff --git a/mockpy/utils/log.py b/mockpy/utils/log.py index 21f89dc..9766bde 100644 --- a/mockpy/utils/log.py +++ b/mockpy/utils/log.py @@ -19,7 +19,7 @@ def log_multiple_matches(items): def log_url(url): - info("\nRequest with url '%s'" % url) + info("Request with url '%s'" % url) def log_request(request): diff --git a/mockpy/utils/extensions.py b/mockpy/utils/proxy_extensions.py similarity index 99% rename from mockpy/utils/extensions.py rename to mockpy/utils/proxy_extensions.py index c08c184..d2edebb 100644 --- a/mockpy/utils/extensions.py +++ b/mockpy/utils/proxy_extensions.py @@ -17,6 +17,8 @@ def to_mapper_request(self): flow_request = MappingRequest(params) return flow_request +HTTPRequest.to_mapper_request = to_mapper_request + def from_httplib_headers(cls, headers): odict = cls() @@ -31,6 +33,9 @@ def from_httplib_headers(cls, headers): return odict +ODictCaseless.from_httplib_headers = classmethod(from_httplib_headers) + + def from_intercepted_response(cls, response): headers = ODictCaseless.from_httplib_headers(response.headers) response = cls(code=response.status, @@ -48,8 +53,7 @@ def with_html(cls, html): httpversion=(1, 1)) return response -HTTPRequest.to_mapper_request = to_mapper_request -ODictCaseless.from_httplib_headers = classmethod(from_httplib_headers) HTTPResponse.from_intercepted_response = classmethod(from_intercepted_response) HTTPResponse.with_html = classmethod(with_html) + diff --git a/requirements.txt b/requirements.txt index 6a5794b..03a249f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ netlib -wheel==0.23.0 \ No newline at end of file +watchdog +wheel==0.23.0 diff --git a/setup.py b/setup.py index 15295fd..115c2f3 100755 --- a/setup.py +++ b/setup.py @@ -18,7 +18,8 @@ "mock", "cherrypy", "netlib", - "termcolor" + "termcolor", + "watchdog" ] test_requirements = [ diff --git a/tests/test_mapping.py b/tests/test_mapping.py index 373cd80..c561f0b 100644 --- a/tests/test_mapping.py +++ b/tests/test_mapping.py @@ -86,13 +86,5 @@ def test_can_set_headers(self): assert mock_cherry.response.headers == {"header_key": "header_value"} - def test_return_correct_dict_for_cherry(self): - mock_cherry = Mock() - mock_cherry.response = Mock(status=1) - mock_cherry.url = Mock(return_value="ss") - - cherry_mapper = CherryPyMapper(cherrypy=mock_cherry) - assert cherry_mapper.cherry_request_dict()["url"] == "ss" - if __name__ == '__main__': unittest.main()