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.
314 lines
10 KiB
314 lines
10 KiB
7 years ago
|
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))
|