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 "" % 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))