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