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.
150 lines
4.3 KiB
150 lines
4.3 KiB
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'), |
|
])
|
|
|