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

parser = argparse.ArgumentParser(description='Send and receive stuff using python sockets (ala netcat)')
parser.add_argument("-l", "--listen", help="enable listener mode (specify -p)", action="store_true")
parser.add_argument("-e", "--execute", help="execute command upon receiving connection", default="")
parser.add_argument("-c", "--command", help="initialize a command shell", action="store_true")
parser.add_argument("-u", "--upload_destination", help="write incoming data to file", default="")
parser.add_argument("-t", "--target", help="target host", default="")
parser.add_argument("-p", "--port", help="target port", type=int, default="")
args = parser.parse_args()

def main():
    global args

    # TX or RX?
    if not args.listen and len(args.target) and args.port > 0:
        buffer = sys.stdin.read()
        client_sender(buffer)
    if args.listen:
        server_loop()

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

    try:
        client.connect((args.target,args.port))
        if len(buffer):
            client.send(buffer)
        while True:
            recv_len = 1
            response = ""

            while recv_len:
                data = client.recv(4096)
                recv_len = len(data)
                response+=data
                if recv_len < 4096:
                    break
            print response
            buffer = raw_input("")
            buffer += "\n"
            client.send(buffer)

    except Exception as e:
        #print "[*] Exception: %s" % e
        print "[*] Exception! Exited."
        client.close()

def server_loop():
    if not len(args.target):
        target="0.0.0.0"
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind((args.target, args.port))
    server.listen(5)

    while True:
        client_socket, addr = server.accept()
        client_thread = threading.Thread(target=client_handler,
                                         args=(client_socket,))
        client_thread.start()

def run_command(command):

    #trim newline
    command = command.rstrip()
    try:
        output = subprocess.check_output(command,stderr=subprocess.STDOUT, shell=True)
    except Exception as e:
        print e
        output = "Failed to execute command.\r\n"
    return output

def client_handler(client_socket):
    #if 'upload_destination' option is specified, listen for data
    if len(args.upload_destination):
        file_buffer = ""
        while True:
            data = client_socket.recv(1024)

            if not data:
                break

            else:
                file_buffer += data

        #write to file
        try:
            fd = open(args.upload_destination,"wb")
            fd.write(file_buffer)
            fd.close()
        except Exception as e:
            client_socket.send("Failed to save file to %s\r\n" % args.upload_destination)
            print "[*] Exception: %s" % e

    #if 'execute' option specified, execute command
    if len(args.execute):
            output = run_command(args.execute)
            client_socket.send(output)

    #if 'command' option specified, enter REPL
    if args.command:
        while True:
            client_socket.send("$> ")
            cmd_buffer  = ""
            while "\n" not in cmd_buffer:
                cmd_buffer += client_socket.recv(1024)
            response = run_command(cmd_buffer)
            client_socket.send(response)

main()