Python implementation of netcat

Question

I'm trying to write a simple python script that acts like the nc program (netcat). I've got it working on OSX (the client that is) by testing against a netcat server, but I'm unsure about the way I've implemented it.

I'd appreciate it if anyone could give me some advice and/or let me know how you would implement it as I still consider myself a beginner. Maybe I should study more on threads and sockets before taking this task on?

By the way, I chose threads over polling for the reason that files in select don't work on Windows.

import sys
import socket
import select
import time
import threading

def _stdout_write(data):
    sys.stdout.write(data)
    sys.stdout.flush()

class NetSnake(object):
    def __init__(self):
        self._close_connection = threading.Event()
        self._stdin_handler = threading.Thread(target=self._handle_stdin)
        self._stdin_handler.daemon = True

    def connect(self, rhost, rport):
        # handle connection refused...
        self._active_sock = socket.create_connection((rhost, rport))
        self._active_sock.setblocking(0)
        self._stdin_handler.start()
        self._handle_active_sock()

    def reverse(self, *args, **kwargs):
        self.connect(*args, **kwargs)

    def _handle_stdin(self):
        while True:
            line = sys.stdin.readline()
            if not line:
                # eof
                break
            self._active_sock.sendall(line)
        # let other thread know that we want to close the connection
        # because eof has been reached.
        self._close_connection.set()

    def _handle_active_sock(self):
        # NOTE: select on windows may not be interrupted by Ctrl-C
        # if not, try time.sleep or event.wait instead.
        self._close_connection.clear()
        timeout = 0.1
        try:
            while not self._close_connection.is_set()
                ready = select.select([self._active_sock], [], [], timeout)
                if ready[0]:
                    data = self._active_sock.recv(4096)
                    if not data:
                        break
                    _stdout_write(data)
        except KeyboardInterrupt:
            _stdout_write('\n')
        finally:
            self._clean_up()

    def _clean_up(self):
        self._active_sock.close()

Edit: I've now realized that I can get rid of the event that signals the main thread that EOF has been reached and just simply close the socket and handle select's error instead.

def _handle_stdin(self):
    try:
        while True:
            line = sys.stdin.readline()
            if not line:
                # eof
                break
            self._active_sock.sendall(line)
    finally:
        self._clean_up()


def _handle_active_sock(self):
    try:
        while True:
            ready = select.select([self._active_sock], [], [])
            if ready[0]:
                data = self._active_sock.recv(4096)
                if not data:
                    break
                _stdout_write(data)
    except KeyboardInterrupt:
        _stdout_write('\n')
    except select.error:
        pass
    finally:
        self._clean_up()

Answer 1

In case you don't know, sockets in Pythons can have an associated timeout. This will take care of the setblocking(0) + select part for you. Just catch the socket.timeout exception and decide what to do when the operation was canceled.

You can also set the timeout directly as the second argument to create_connection.

You can thus read from the socket using:

def print_recv(sock):
    while True:
        try:
            msg = sock.recv(4096)
        except socket.timeout:
            pass
        else:
            sys.stdout.write(msg)
            sys.stdout.flush()

Granted that the socket already has an associated timeout (or you can call sock.settimeout(0.1) as the first instruction of the function).

As a side note, this function does not account for shutdowns of the connexion from the server. If you want to handle this case, you can detect it when recv return an empty message. So the basic thing to add could be if not msg: break before sys.stdout.write(msg).

Using this function as a thread, you can close the socket from the outside, somehow like you did in your second version, but the socket will, at some point, trigger an OSError on sock.recv. Catching that could be a bad idea, as it may hide more serious issues. Instead, I’d still rely on threading.Event:

def print_recv(sock, event):
    while not event.is_set():
        try:
            msg = sock.recv(4096)
        except socket.timeout:
            pass
        else:
            if not msg:
                break
            sys.stdout.write(msg)
            sys.stdout.flush()

So you’ll have to event.set() before socket.close() outside of the thread. Note, however, that due to scheduling randomness, the order of execution could be:

while not event.is_set(): # Thread-1
event.set()               # Main thread
sock.close()              # Main thread
    msg = sock.recv(4096) # Thread-1

and still lead to an OSError. It will not happen so often, but it will. So you may want to come up with better synchronization if you want a cleaner termination.

---

You’ll also note that I’m using a function instead of a method to print the messages from the server. It feels cleaner as it separate concerns. In fact, there is no real need to create a class here, appart from keeping around a bunch of related objects. But your setup + teardown logic had me wondering: this is a job for a context manager:

class NetcatManager(object):
    def __init__(self, rhost, rport):
        self._event = threading.Event()
        self._socket = socket.create_connection((rhost, rport))
        self._socket.settimeout(0.1)

    def __enter__(self):
        self._reading_thread = threading.Thread(target=print_recv, args=(self._socket, self._event))
        self._reading_thread.start()
        return self._socket

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._event.set()
        # Waiting for thread termination so there is no OSError when closing the socket
        self._reading_thread.join()
        self._socket.close()

Now all what is left to do is to send messages from stdin to the netcat server. Using the above context manager, it should look like the following:

def netcat_session(rhost, rport):
    with NetcatManager(rhost, rport) as netcat_socket:
        # Process stdin and use netcat_socket.send

And that’s it.

---

Now, on the topic of processing stdin, be aware that Ctrl+C (which is translated to KeyboardInterrupt) does not mean EOF: Ctrl+D does. So catching the KeyboardInterrupt is meant to recover from an unusual termination of the program; it may or may not be what you want.

Putting everything together, we can come to:

import sys
import socket
import threading


class NetcatManager(object):
    def __init__(self, rhost, rport, timeout=0.1):
        self._event = threading.Event()
        self._socket = socket.create_connection((rhost, rport))
        self._socket.settimeout(timeout)

    def __enter__(self):
        self._reading_thread = threading.Thread(target=print_recv, args=(self._socket, self._event))
        self._reading_thread.start()
        return self._socket

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._event.set()
        # Waiting for thread termination so there is no OSError when closing the socket
        self._reading_thread.join()
        self._socket.close()


def print_recv(sock, event):
    while not event.is_set():
        try:
            msg = sock.recv(4096)
        except socket.timeout:
            pass
        else:
            if not msg:
                break
            sys.stdout.write(msg)
            sys.stdout.flush()


def netcat_session(rhost, rport, input_file=sys.stdin):
    with NetcatManager(rhost, rport) as netcat_socket:
        # We do not iterate (for line in input_file) because stdin does not behave well in Python 2
        while True:
            line = input_file.readline()
            if not line:
                # EOF reached
                break
            netcat_socket.sendall(line)
    # Leave with the screen in a clean state
    print

Note the use of default values. You can now provide any file-like object as the third parameter of netcat_session and it will send instructions from there (so you can save a test session in a file and replay it over and over again using with open('the_file.txt') as f: netcat_session('localhost', 5555, f)).