First order of business is updating the netcat replacement Seitz has written so it works again. My code is here
The issues
Simply put our biggest problem here is types.
The socket.recv
and socket.send
no longer deal with strings and
expect a bytes
type, so our updated script will need to be aware
of and work with those, meaning we can’t really rely on newline
characters to denote the end of a message. Our new implementation
will need to change this up to utilise those new types
General Niceties
First things first, I wanted to fix up that argument parsing to
be a little less finnicky, so let’s back it onto the built in
argparse
library instead
def parse_args():
parser = argparse.ArgumentParser(prog='nettool.py',description='Connect to a TCP server or create a server on a port')
parser.add_argument('-t', '--target', dest='target', metavar='host', type=str,
help='IP target or address to bind to')
parser.add_argument('-p', '--port', dest='port', metavar='port', type=int,
help='Target port or port to bind to')
parser.add_argument('-l', '--listen', dest='listen', action='store_true',
help='Initialise a listener on {target}:{port}')
parser.add_argument('-c', '--command', dest='command', action='store_true',
help='Attach a command listener to a server. Cannot be used with -u')
parser.add_argument('-e', '--echo', dest='echo', action='store_true',
help='Attach an echo listener to a server')
parser.add_argument('-u', '--upload', dest='upload', metavar='upload_location', type=str,
help='Start an upload server and upload to {upload_location}. Cannot be used with -c')
parser.set_defaults(listen=False, command=False, echo=False)
args = parser.parse_args(sys.argv[1:])
arg_problems = arg_sanity_check(args)
if len(arg_problems) > 0:
for p in arg_problems:
print("[*] {0}".format(p))
parser.print_help()
sys.exit(1)
return args
Here we take the system.argv
values and push them into a parser with
the following arguments set up:
-t
or--target
for host names. Naming this-h
conflicts with
the built in help command
-p
or--port
for ports-l
or--listen
to start servers, also set to be a boolean flag-c
or--command
to attach a command listener to a server-e
or--echo
to attach an echo command listener to a server. Not
really needed, but nice, simple functionality to help test everything
-u
or--upload
to attach an upload listener and specify upload
location
We then set default values for the boolean flags to false.
Little cleaner than the book’s function, and also handles the usage string for us.
TCP Client Class
My general strategy now that I’m going through this book will be
to use the select
library instead of relying on looking for
newlines in my received data. Looking through Seitz’ work he seems
to be going for the quick wins, which is fair enough, but any real
world, networked application is going to use this call to help keep
things clean. In theory, I could open up multiple connections and
use select
in a single thread to handle all of them and cut down
on inter-process communication.
I’ve also improved error handling, encapsulated it all into a class, and shut down sockets cleanly to limit other issues we might see
class Client:
def __init__(self, target, port):
self.__target = target
self.__port = port
self.__socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.__readerThread = threading.Thread(target=self.__reader, args=())
self.__stop = False
self.__target_disconnect = False
def run(self):
try:
self.__socket.connect((self.__target, self.__port))
self.__readerThread.start()
while not self.__stop:
t = input()
msg = "{0}\r\n".format(t)
sent = self.__socket.send(bytes(msg, 'utf-8'))
if sent == 0:
print("[*] Collection closed")
break
except EOFError as e:
print('[*] End of file reached')
except InterruptedError as e:
print('[*] Interrupted. Exiting')
except Exception as e:
print('[*] Exception thrown')
print(e)
finally:
if not self.__target_disconnect:
self.__stop = True
self.__socket.shutdown(socket.SHUT_RDWR)
self.__socket.close()
self.__readerThread.join(100)
def __reader(self):
while not self.__stop:
try:
buffer = ""
while True:
(readylist, x, y) = select.select([self.__socket], [], [], 0.01)
if len(readylist) == 0:
# Target is probably done writing
break
data = self.__socket.recv(1024)
if len(data) == 0 or data == b'\xff\xf4\xff\xfd\x06':
# Socket closed
self.__stop = True
self.__target_disconnect = True
break
elif data:
buffer += data.decode('utf-8')
else:
break
if len(buffer) > 0:
print(buffer)
except Exception as e:
print("[*] Exception thrown: {0}".format(e))
if self.__target_disconnect:
print("[!!] Target machine shut down")
# Interrupt the input
os.kill(os.getpid(), signal.SIGINT)
You can see here that the client class splits up the reader and
writer routines, also putting the reader on a background thread.
If the reader finds that the socket is closed, it will send an
interrupt to itself to rip out the call to input()
and close
the program down gracefully.
The Server Class
class Server:
def __init__(self, target, port, handlers):
self.__socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.__target = target
self.__port = port
self.__handlers = handlers
self.__stop = False
def listen(self):
print('[*] Listening on {0}:{1}'.format(self.__target, self.__port))
self.__socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self.__socket.bind((self.__target, self.__port))
self.__socket.listen(5)
try:
while True:
connection, addr = self.__socket.accept()
print("[*] Connection from {0}".format(addr))
proc = threading.Thread(target=self.__handle, args=(connection, addr))
proc.start()
except Exception as e:
print("[*] Exception caught, shutting down")
print(e)
except KeyboardInterrupt as e:
print("[!!] Server was interupted. Shutting down")
finally:
self.__stop = True
self.__socket.shutdown(socket.SHUT_RDWR)
self.__socket.close()
def __handle(self, client_conn, addr):
close = False
for handler in self.__handlers:
handler.init_connection(client_conn)
try:
while not close and not self.__stop:
raw_buffer = bytearray()
while True:
(sock_ready,x,y) = select.select([client_conn],[], [], 0.01)
if len(sock_ready) == 0:
break
data = client_conn.recv(1028)
if len(data) == 0 or data == b'\xff\xf4\xff\xfd\x06':
# Connection was probably closed
close = True
break
elif data:
raw_buffer.extend(data)
else:
break
if len(raw_buffer) > 0:
for handler in self.__handlers:
try:
close = handler.handle_msg(raw_buffer,client_conn)
if close:
break
except Exception as e:
print("[*] Caught an exception")
print(e)
except BrokenPipeError as e:
print("[*] Connection closed")
finally:
print("[*] Closing connection from {0}".format(addr))
client_conn.shutdown(socket.SHUT_RDWR)
client_conn.close()
Here’s the fun part. The code in the book goes looking for a newline
before dealing with the message, so I’ve switched that around
and put in a select
with a timeout instead to avoid the possible
typing issues.
We could tidy this class up further by using a thread pool and worker pattern instead of spinning up a new thread for every incoming connection, but for a quick script this will do.
Message Handlers
This is just a way to make it easier for me to add functionality later if I want to. I could add more handler classes quite easily, and just add a flag or even directly into my initialisation routine later on.
class Handler:
def init_connection(self, connection):
pass
def handle_msg(self, msg, connection):
return False
init_connection
is called when a client first connects, while
handle_msg
is called when a client sends us a message. Implementing
a new handler is pretty easy
class EchoHandler(Handler):
def init_connection(self, connection):
connection.send(b'Echo enabled\r\n')
def handle_msg(self, msg, connection):
if len(msg) == 0:
return False
strmsg = msg.decode('utf-8')
strmsg.rstrip()
connection.send(bytes("{0}\r\n".format(strmsg), 'utf-8'))
return False
Which are then called in the __handler
method of the server class
...
for handler in self.__handlers:
handler.init_connection(client_conn)
...
for handler in self.__handlers:
try:
close = handler.handle_msg(raw_buffer,client_conn)
if close:
break
except Exception as e:
print("[*] Caught an exception")
print(e)
The handler can return True
to tell the server to close off the
connection.
Pulling all this together
The main
function brings all this together for use. It looks
at the args passed in and spins up the appropriate client or
server.
Here we can see the call to parse_args
, and how the script builds
up our handler list if we have passed in the -l
flag
if args.listen:
handlers = []
if args.command:
handlers.append(CommandHandler())
if args.echo:
handlers.append(EchoHandler())
if args.upload:
handlers.append(UploadHandler(args.upload))
s = Server(args.target, args.port, handlers)
s.listen()
Or how it creates a client if that flag isn’t passed
if not args.listen:
client = Client(args.target, args.port)
client.run()
Final thoughts
I’m not a Python programmer (well, not an expert anyway), so there may be more Pythonic way of doing this, as such as I work through the book I may go back and update the repo. The basic structure won’t change, but if I can find ways to make the code cleaner I will.
This script is definitely overkill, but I always like making things easier to understand and leave them a bit cleaner than when I found them. It was interesting seeing how python 3 has changed a few things, since the last time I really used it I was using Python 2.7. I do really appreciate how fast it was to get something up and running, having been using Java and golang a lot it was refreshing to be able to do something short and sweet without having to think too hard about it.
Having read through the TCP proxy script now, I think it’s going to need a lot of similar modifications as well as a general cleanup just to help readability. No biggie, this is a fun little project for now. Should hopefully have it translated pretty quickly with a matching blog to go with it