PappyProxy/pappyproxy/interface/macros.py

from ..util import load_reqlist
from ..macros import macro_from_requests, MacroTemplate, load_macros
from ..colors import Colors

macro_dict = {}
int_macro_dict = {}
int_conns = {}

def generate_macro(client, args):
    if len(args) == 0:
        print("usage: gma [name] [reqids]")
        return
    macro_name = args[0]

    reqs = [r for r in load_reqlist(client, ','.join(args[1:]))]
    script_string = macro_from_requests(reqs)
    fname = MacroTemplate.template_filename('macro', macro_name)
    with open(fname, 'w') as f:
        f.write(script_string)
    print("Macro written to {}".format(fname))

def generate_int_macro(client, args):
    if len(args) == 0:
        print("usage: gima [name] [reqids]")
        return
    macro_name = args[0]

    reqs = [r for r in load_reqlist(client, ','.join(args[1:]))]

    script_string = macro_from_requests(reqs, template='intmacro')
    fname = MacroTemplate.template_filename('intmacro', macro_name)
    with open(fname, 'w') as f:
        f.write(script_string)
    print("Macro written to {}".format(fname))
    
def load_macros_cmd(client, args):
    global macro_dict

    load_dir = '.'
    if len(args) > 0:
        load_dir = args[0]

    _stop_all_int_macros()

    loaded_macros, loaded_int_macros = load_macros(load_dir, client)
    for macro in loaded_macros:
        macro_dict[macro.name] = macro
        print("Loaded {} ({})".format(macro.name, macro.file_name))
    for macro in loaded_int_macros:
        int_macro_dict[macro.name] = macro
        print("Loaded {} ({})".format(macro.name, macro.file_name))

def complete_run_macro(text, line, begidx, endidx):
    from ..util import autocomplete_starts_with

    global macro_dict
    strs = macro_dict.keys()
    return autocomplete_startswith(text, strs)
        
def run_macro(client, args):
    global macro_dict
    if len(args) == 0:
        print("usage: rma [macro name]")
        return
    macro = macro_dict[args[0]]
    macro.execute(client, args[1:])

def complete_run_int_macro(text, line, begidx, endidx):
    from ..util import autocomplete_starts_with

    global int_macro_dict
    strs = int_macro_dict.keys()
    return autocomplete_startswith(text, strs)

def run_int_macro(client, args):
    global int_macro_dict
    global int_conns
    if len(args) == 0:
        print("usage: rim [macro name]")
        return
    if args[0] in int_conns:
        print("%s is already running!" % args[0])
        return
    macro = int_macro_dict[args[0]]
    macro.init(args[1:])
    conn = client.new_conn()
    int_conns[args[0]] = conn
    conn.intercept(macro)
    print("Started %s" % args[0])

def complete_stop_int_macro(text, line, begidx, endidx):
    from ..util import autocomplete_starts_with

    global int_conns
    strs = int_conns.keys()
    return autocomplete_startswith(text, strs)

def stop_int_macro(client, args):
    global int_conns
    if len(args) > 0:
        conn = int_conns[args[0]]
        conn.close()
        del int_conns[args[0]]
        print("Stopped %s" % args[0])
    else:
        _stop_all_int_macros()

def _stop_all_int_macros():
    global int_conns
    for k, conn in int_conns.items():
        conn.close()
        del int_conns[k]
        print("Stopped %s" % k)

def list_macros(client, args):
    global macro_dict
    global int_macro_dict
    global int_conns
    if len(macro_dict) > 0:
        print('Loaded Macros:')
    for k, m in macro_dict.items():
        print('  '+k)

    if len(int_macro_dict) > 0:
        print('Loaded Intercepting Macros:')
        for k, m in int_macro_dict.items():
            pstr = '  '+k
            if k in int_conns:
                pstr += ' (' + Colors.GREEN + 'RUNNING' + Colors.ENDC + ')'
            print(pstr)

def load_cmds(cmd):
    cmd.set_cmds({
        'generate_macro': (generate_macro, None),
        'generate_int_macro': (generate_int_macro, None),
        'load_macros': (load_macros_cmd, None),
        'run_macro': (run_macro, complete_run_macro),
        'run_int_macro': (run_int_macro, complete_run_int_macro),
        'stop_int_macro': (stop_int_macro, complete_stop_int_macro),
        'list_macros': (list_macros, None),
    })
    cmd.add_aliases([
        ('generate_macro', 'gma'),
        ('generate_int_macro', 'gima'),
        ('load_macros', 'lma'),
        ('run_macro', 'rma'),
        ('run_int_macro', 'rim'),
        ('stop_int_macro', 'sim'),
        ('list_macros', 'lsma'),
    ])