#!/usr/bin/env python
import sys
import socket
import threading
import subprocess
import argparse


def main():
    if not args.listen:
        buffer = input()
        buffer += '\n'
        client_sender(buffer)
    else:
        server_loop()


def client_sender(buf):
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    try:
        # connect to our target host
        client.connect((args.target, args.port))
        if len(buf):
            client.sendall(buf.encode())

        # now wait for the data back
        while True:
            recv_len = 1
            response = ""

            while recv_len:
                data = client.recv(4096).decode()
                recv_len = len(data)
                response += data
                if recv_len < 4096:
                    break

            # wait for more input
            buf = input()
            buf += '\n'

            # send if off
            client.sendall(buf.encode())

    except Exception as e:
        print(f"[*] Exception! Exiting. {e}")

        # tear down the connection
        client.close()


def server_loop():

    # if no target defined we listen on all interfaces
    if not args.target:
        args.target = '0.0.0.0'

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind((args.target, args.port))

    # maximum 5 concurrent connections
    server.listen(5)

    while True:
        client_socket, addr = server.accept()

        # spin off a thread to handle our new client
        client_thread = threading.Thread(
            target=client_handler, args=(client_socket,))
        client_thread.start()


def run_command(command):

    # trim the new line
    command = command.rstrip()

    # run the command and get the output back
    try:
        output = subprocess.check_output(
            command, stderr=subprocess.STDOUT, shell=True).decode()
    except:
        output = f"Failed to execute command {command}"

    # send the output back to the client
    return output


def client_handler(client_socket):

    print(f"Handling client {client_socket}")

    if args.upload:
        # read in all the bytes and write to our destination
        file_buffer = ""

        # keep reading data until none is available
        while True:
            data = client_socket.recv(1024).decode()

            if not data:
                break
            else:
                file_buffer += data

            # now we take these bytes and try to write them out
            try:
                fd = open(args.upload, "wb")
                fd.write(file_buffer)
                fd.close()

                # acknowledge that we wrote the file out
                client_socket.send(
                    f"Successfully saved file to {args.upload}\r\n".encode())
            except:
                client_socket.send(
                    f"Failed saved file to {args.upload}\r\n".encode())

    # check for command execution
    if args.execute:
        # run this command
        output = run_command(args.execute)
        client_socket.send(output.encode())

    if args.command:

        while True:
            # show a simple prompt

            # now we receive until we see a linefeed (enter key)
            cmd_buffer = ""
            while "\n" not in cmd_buffer:
                cmd_buffer += client_socket.recv(1024).decode()

            # send back the command output, response is bytes
            response = run_command(cmd_buffer).rstrip('\n')
            response += '\n> Command:'
            client_socket.sendall(response.encode())


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-t", "--target", help="target host", type=str)
    parser.add_argument("-p", "--port", help="target port",
                        required=True, type=int)
    parser.add_argument("-v", "--verbose", help="increase output verbosity",
                        action="count", default=0)
    parser.add_argument("-l",
                        "--listen", help="listen on [host]:[port] for incoming connections", action='store_true')
    parser.add_argument("-e",
                        "--execute", help="execute the given file upon receiving a connection", type=str)
    parser.add_argument('-c', '--command',
                        help='initialize a command shell', action='store_true')
    parser.add_argument(
        '-u', '--upload', help='upon receiving connection upload a file and write to [UPLOAD]', type=str)
    args = parser.parse_args()
    print(args)
    main()
