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.

172 lines
5.6 KiB

import argparse
import sys
from ..util import copy_to_clipboard, confirm, printable_data
from ..console import CommandError
from ..proxy import InterceptMacro
from ..colors import url_formatter, verb_color, Colors, scode_color
class WatchMacro(InterceptMacro):
def __init__(self):
InterceptMacro.__init__(self)
self.name = "WatchMacro"
def mangle_request(self, request):
printstr = "> "
printstr += verb_color(request.method) + request.method + Colors.ENDC + " "
printstr += url_formatter(request, colored=True)
print(printstr)
return request
def mangle_response(self, request, response):
printstr = "< "
printstr += verb_color(request.method) + request.method + Colors.ENDC + ' '
printstr += url_formatter(request, colored=True)
printstr += " -> "
response_code = str(response.status_code) + ' ' + response.reason
response_code = scode_color(response_code) + response_code + Colors.ENDC
printstr += response_code
print(printstr)
return response
def mangle_websocket(self, request, response, message):
printstr = ""
if message.to_server:
printstr += ">"
else:
printstr += "<"
printstr += "ws(b={}) ".format(message.is_binary)
printstr += printable_data(message.message)
print(printstr)
return message
def message_address(client, args):
msg_addr = client.maddr
if msg_addr is None:
print("Client has no message address")
return
print(msg_addr)
if len(args) > 0 and args[0] == "-c":
try:
copy_to_clipboard(msg_addr.encode())
print("Copied to clipboard!")
except:
print("Could not copy address to clipboard")
def cpinmem(client, args):
req = client.req_by_id(args[0])
client.save_new(req, client.inmem_storage.storage_id)
def ping(client, args):
print(client.ping())
def watch(client, args):
macro = WatchMacro()
macro.intercept_requests = True
macro.intercept_responses = True
macro.intercept_ws = True
with client.new_conn() as conn:
conn.intercept(macro)
print("Watching requests. Press <Enter> to quit...")
input()
def submit(client, cargs):
"""
Resubmit some requests, optionally with modified headers and cookies.
Usage: submit [-h] [-m] [-u] [-p] [-o REQID] [-c [COOKIES [COOKIES ...]]] [-d [HEADERS [HEADERS ...]]]
"""
#Usage: submit reqids [-h] [-m] [-u] [-p] [-o REQID] [-c [COOKIES [COOKIES ...]]] [-d [HEADERS [HEADERS ...]]]
parser = argparse.ArgumentParser(prog="submit", usage=submit.__doc__)
#parser.add_argument('reqids')
parser.add_argument('-m', '--inmem', action='store_true', help='Store resubmitted requests in memory without storing them in the data file')
parser.add_argument('-u', '--unique', action='store_true', help='Only resubmit one request per endpoint (different URL parameters are different endpoints)')
parser.add_argument('-p', '--uniquepath', action='store_true', help='Only resubmit one request per endpoint (ignoring URL parameters)')
parser.add_argument('-c', '--cookies', nargs='*', help='Apply a cookie to requests before submitting')
parser.add_argument('-d', '--headers', nargs='*', help='Apply a header to requests before submitting')
parser.add_argument('-o', '--copycookies', help='Copy the cookies used in another request')
args = parser.parse_args(cargs)
headers = {}
cookies = {}
clear_cookies = False
if args.headers:
for h in args.headers:
k, v = h.split('=', 1)
headers[k] = v
if args.copycookies:
reqid = args.copycookies
req = client.req_by_id(reqid)
clear_cookies = True
for k, v in req.cookie_iter():
cookies[k] = v
if args.cookies:
for c in args.cookies:
k, v = c.split('=', 1)
cookies[k] = v
if args.unique and args.uniquepath:
raise CommandError('Both -u and -p cannot be given as arguments')
# Get requests to submit
#reqs = [r.copy() for r in client.in_context_requests()]
reqs = client.in_context_requests()
# Apply cookies and headers
for req in reqs:
if clear_cookies:
req.headers.delete("Cookie")
for k, v in cookies.items():
req.set_cookie(k, v)
for k, v in headers.items():
req.headers.set(k, v)
conf_message = "You're about to submit %d requests, continue?" % len(reqs)
if not confirm(conf_message):
return
# Filter unique paths
if args.uniquepath or args.unique:
endpoints = set()
new_reqs = []
for r in reqs:
if unique_path_and_args:
s = r.url.geturl()
else:
s = r.url.geturl(include_params=False)
if not s in endpoints:
new_reqs.append(r)
endpoints.add(s)
reqs = new_reqs
# Tag and send them
for req in reqs:
req.tags.add('resubmitted')
sys.stdout.write(client.prefixed_reqid(req) + " ")
sys.stdout.flush()
storage = client.disk_storage.storage_id
if args.inmem:
storage = client.inmem_storage.storage_id
client.submit(req, storage=storage)
sys.stdout.write("\n")
sys.stdout.flush()
def load_cmds(cmd):
cmd.set_cmds({
'maddr': (message_address, None),
'ping': (ping, None),
'submit': (submit, None),
'cpim': (cpinmem, None),
'watch': (watch, None),
})