This is a fork of:
https://github.com/roglew/puppy
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
313 lines
10 KiB
313 lines
10 KiB
import glob |
|
import imp |
|
import os |
|
import random |
|
import re |
|
import stat |
|
from jinja2 import Environment, FileSystemLoader |
|
from collections import namedtuple |
|
|
|
from .proxy import InterceptMacro |
|
|
|
class MacroException(Exception): |
|
pass |
|
|
|
class FileInterceptMacro(InterceptMacro): |
|
""" |
|
An intercepting macro that loads a macro from a file. |
|
""" |
|
def __init__(self, filename=''): |
|
InterceptMacro.__init__(self) |
|
self.file_name = '' # name from the file |
|
self.filename = filename or '' # filename we load from |
|
self.source = None |
|
|
|
if self.filename: |
|
self.load() |
|
|
|
def __repr__(self): |
|
s = self.name |
|
names = [] |
|
names.append(self.file_name) |
|
s += ' (%s)' % ('/'.join(names)) |
|
return "<InterceptingMacro %s>" % s |
|
|
|
def load(self): |
|
if self.filename: |
|
match = re.findall('.*int_(.*).py$', self.filename) |
|
if len(match) > 0: |
|
self.file_name = match[0] |
|
else: |
|
self.file_name = self.filename |
|
|
|
# yes there's a race condition here, but it's better than nothing |
|
st = os.stat(self.filename) |
|
if (st.st_mode & stat.S_IWOTH): |
|
raise MacroException("Refusing to load world-writable macro: %s" % self.filename) |
|
module_name = os.path.basename(os.path.splitext(self.filename)[0]) |
|
self.source = imp.load_source('%s'%module_name, self.filename) |
|
if self.source and hasattr(self.source, 'MACRO_NAME'): |
|
self.name = self.source.MACRO_NAME |
|
else: |
|
self.name = module_name |
|
else: |
|
self.source = None |
|
|
|
# Update what we can do |
|
if self.source and hasattr(self.source, 'mangle_request'): |
|
self.intercept_requests = True |
|
else: |
|
self.intercept_requests = False |
|
|
|
if self.source and hasattr(self.source, 'mangle_response'): |
|
self.intercept_responses = True |
|
else: |
|
self.intercept_responses = False |
|
|
|
if self.source and hasattr(self.source, 'mangle_websocket'): |
|
self.intercept_ws = True |
|
else: |
|
self.intercept_ws = False |
|
|
|
def init(self, args): |
|
if hasattr(self.source, 'init'): |
|
self.source.init(args) |
|
|
|
def mangle_request(self, request): |
|
if hasattr(self.source, 'mangle_request'): |
|
req = self.source.mangle_request(request) |
|
return req |
|
return request |
|
|
|
def mangle_response(self, request): |
|
if hasattr(self.source, 'mangle_response'): |
|
rsp = self.source.mangle_response(request, request.response) |
|
return rsp |
|
return request.response |
|
|
|
def mangle_websocket(self, request, message): |
|
if hasattr(self.source, 'mangle_websocket'): |
|
mangled_ws = self.source.mangle_websocket(request, request.response, message) |
|
return mangled_ws |
|
return message |
|
|
|
class MacroFile: |
|
""" |
|
A class representing a file that can be executed to automate actions |
|
""" |
|
|
|
def __init__(self, filename=''): |
|
self.name = '' # name from the file |
|
self.file_name = filename or '' # filename we load from |
|
self.source = None |
|
|
|
if self.file_name: |
|
self.load() |
|
|
|
def load(self): |
|
if self.file_name: |
|
match = re.findall('.*macro_(.*).py$', self.file_name) |
|
self.name = match[0] |
|
st = os.stat(self.file_name) |
|
if (st.st_mode & stat.S_IWOTH): |
|
raise PappyException("Refusing to load world-writable macro: %s" % self.file_name) |
|
module_name = os.path.basename(os.path.splitext(self.file_name)[0]) |
|
self.source = imp.load_source('%s'%module_name, self.file_name) |
|
else: |
|
self.source = None |
|
|
|
def execute(self, client, args): |
|
# Execute the macro |
|
if self.source: |
|
self.source.run_macro(client, args) |
|
|
|
MacroTemplateData = namedtuple("MacroTemplateData", ["filename", "description", "argdesc", "fname_fmt"]) |
|
|
|
class MacroTemplate(object): |
|
_template_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), |
|
"templates") |
|
_template_data = { |
|
'macro': MacroTemplateData('macro.py.tmpl', |
|
'Generic macro template', |
|
'[reqids]', |
|
'macro_{fname}.py'), |
|
|
|
'intmacro': MacroTemplateData('intmacro.py.tmpl', |
|
'Generic intercepting macro template', |
|
'', |
|
'int_{fname}.py'), |
|
} |
|
|
|
@classmethod |
|
def fill_template(cls, template, subs): |
|
loader = FileSystemLoader(cls._template_dir) |
|
env = Environment(loader=loader) |
|
template = env.get_template(cls._template_data[template].filename) |
|
return template.render(zip=zip, **subs) |
|
|
|
@classmethod |
|
def template_filename(cls, template, fname): |
|
return cls._template_data[template].fname_fmt.format(fname=fname) |
|
|
|
@classmethod |
|
def template_names(cls): |
|
for k, v in cls._template_data.iteritems(): |
|
yield k |
|
|
|
@classmethod |
|
def template_description(cls, template): |
|
return cls._template_data[template].description |
|
|
|
@classmethod |
|
def template_argstring(cls, template): |
|
return cls._template_data[template].argdesc |
|
|
|
## Other functions |
|
|
|
def load_macros(loc): |
|
""" |
|
Loads the macros stored in the location and returns a list of Macro objects |
|
""" |
|
macro_files = glob.glob(loc + "/macro_*.py") |
|
macro_objs = [] |
|
for f in macro_files: |
|
macro_objs.append(MacroFile(f)) |
|
|
|
# int_macro_files = glob.glob(loc + "/int_*.py") |
|
# int_macro_objs = [] |
|
# for f in int_macro_files: |
|
# try: |
|
# int_macro_objs.append(FileInterceptMacro(f)) |
|
# except PappyException as e: |
|
# print(str(e)) |
|
#return (macro_objs, int_macro_objs) |
|
return (macro_objs, []) |
|
|
|
def macro_from_requests(reqs): |
|
# Generates a macro that defines request objects for each of the requests |
|
# in reqs |
|
subs = {} |
|
|
|
req_lines = [] |
|
req_params = [] |
|
for req in reqs: |
|
lines = req.full_message().splitlines(True) |
|
#esclines = [line.encode('unicode_escape') for line in lines] |
|
esclines = [line for line in lines] |
|
req_lines.append(esclines) |
|
|
|
params = [] |
|
params.append('dest_host="{}"'.format(req.dest_host)) |
|
params.append('dest_port={}'.format(req.dest_port)) |
|
params.append('use_tls={}'.format(req.use_tls)) |
|
req_params.append(', '.join(params)) |
|
subs['req_lines'] = req_lines |
|
subs['req_params'] = req_params |
|
|
|
return MacroTemplate.fill_template('macro', subs) |
|
|
|
# @defer.inlineCallbacks |
|
# def mangle_request(request, intmacros): |
|
# """ |
|
# Mangle a request with a list of intercepting macros. |
|
# Returns a tuple that contains the resulting request (with its unmangled |
|
# value set if needed) and a bool that states whether the request was modified |
|
# Returns (None, True) if the request was dropped. |
|
|
|
# :rtype: (Request, Bool) |
|
# """ |
|
# # Mangle requests with list of intercepting macros |
|
# if not intmacros: |
|
# defer.returnValue((request, False)) |
|
|
|
# cur_req = request.copy() |
|
# for macro in intmacros: |
|
# if macro.intercept_requests: |
|
# if macro.async_req: |
|
# cur_req = yield macro.async_mangle_request(cur_req.copy()) |
|
# else: |
|
# cur_req = macro.mangle_request(cur_req.copy()) |
|
|
|
# if cur_req is None: |
|
# defer.returnValue((None, True)) |
|
|
|
# mangled = False |
|
# if not cur_req == request or \ |
|
# not cur_req.host == request.host or \ |
|
# not cur_req.port == request.port or \ |
|
# not cur_req.is_ssl == request.is_ssl: |
|
# # copy unique data to new request and clear it off old one |
|
# cur_req.unmangled = request |
|
# cur_req.unmangled.is_unmangled_version = True |
|
# if request.response: |
|
# cur_req.response = request.response |
|
# request.response = None |
|
# mangled = True |
|
# else: |
|
# # return the original request |
|
# cur_req = request |
|
# defer.returnValue((cur_req, mangled)) |
|
|
|
# @defer.inlineCallbacks |
|
# def mangle_response(request, intmacros): |
|
# """ |
|
# Mangle a request's response with a list of intercepting macros. |
|
# Returns a bool stating whether the request's response was modified. |
|
# Unmangled values will be updated as needed. |
|
|
|
# :rtype: Bool |
|
# """ |
|
# if not intmacros: |
|
# defer.returnValue(False) |
|
|
|
# old_rsp = request.response |
|
# for macro in intmacros: |
|
# if macro.intercept_responses: |
|
# # We copy so that changes to request.response doesn't mangle the original response |
|
# request.response = request.response.copy() |
|
# if macro.async_rsp: |
|
# request.response = yield macro.async_mangle_response(request) |
|
# else: |
|
# request.response = macro.mangle_response(request) |
|
|
|
# if request.response is None: |
|
# defer.returnValue(True) |
|
|
|
# mangled = False |
|
# if not old_rsp == request.response: |
|
# request.response.rspid = old_rsp |
|
# old_rsp.rspid = None |
|
# request.response.unmangled = old_rsp |
|
# request.response.unmangled.is_unmangled_version = True |
|
# mangled = True |
|
# else: |
|
# request.response = old_rsp |
|
# defer.returnValue(mangled) |
|
|
|
# @defer.inlineCallbacks |
|
# def mangle_websocket_message(message, request, intmacros): |
|
# # Mangle messages with list of intercepting macros |
|
# if not intmacros: |
|
# defer.returnValue((message, False)) |
|
|
|
# cur_msg = message.copy() |
|
# for macro in intmacros: |
|
# if macro.intercept_ws: |
|
# if macro.async_ws: |
|
# cur_msg = yield macro.async_mangle_ws(request, cur_msg.copy()) |
|
# else: |
|
# cur_msg = macro.mangle_ws(request, cur_msg.copy()) |
|
|
|
# if cur_msg is None: |
|
# defer.returnValue((None, True)) |
|
|
|
# mangled = False |
|
# if not cur_msg == message: |
|
# # copy unique data to new request and clear it off old one |
|
# cur_msg.unmangled = message |
|
# cur_msg.unmangled.is_unmangled_version = True |
|
# mangled = True |
|
# else: |
|
# # return the original request |
|
# cur_msg = message |
|
# defer.returnValue((cur_msg, mangled))
|
|
|