A Netcat-like implementation of the socket module to easily write and read strings and data to and from a socket. Particularly useful for CTFs.

netcat.py

# encoding: utf-8

"""
netcat.py
by GioBonvi
https://gist.github.com/GioBonvi/2cb9080039f5f5abfdcc44ee81ef102c

A Netcat-like implementation of the socket module to easily write and read strings and data
to and from a connection. Particularly useful for CTFs.

Credits to Leon Jacobs (@leonjza) for the original implementation (https://gist.github.com/leonjza/f35a7252babdf77c8421)
which I adapted slightly.
"""

import socket

class Netcat:
  """ Python 'netcat like' module """

  def __init__(self, ip, port, encoding = 'UTF-8', newline = '\n'):
    self.buff = b''
    self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self.socket.connect((ip, port))
    self.encoding = encoding
    self.newline = newline

  def read(self, length = 1024):
    """ Read some bytes off the socket """

    return self.socket.recv(length)

  def reads(self, length = 1024, encoding = None):
    """ Read some bytes off the socket and decode them into a string """

    encoding = encoding or self.encoding

    return self.read(length).decode(encoding)

  def reaadln(self, n = 1, include_newline = True, encoding = None, newline = None):
    """ Read some lines as strings """

    encoding = encoding or self.encoding
    newline = newline or self.newline

    lines = []
    for _ in range(n):
      lines.append(self.reads_until(newline, encoding) if include_newline else self.reads_before(newline, encoding))

    if n == 1:
      return lines[0]
    else:
      return lines

  def read_until(self, marker):
    """ Read data into the buffer until we have read the marker """

    while not marker in self.buff:
        self.buff += self.socket.recv(1024)

    pos = self.buff.find(marker)
    rval = self.buff[:pos + len(marker)]
    self.buff = self.buff[pos + len(marker):]

    return rval

  def reads_until(self, str_marker, encoding = None):
    """ Read data into the buffer until we have read the string marker and return the data parsed as a string """

    encoding = encoding or self.encoding

    return self.read_until(str.encode(str_marker)).decode(encoding)

  def readln_until(self, str_marker, encoding = None, newline = None):
    """ Read data into the buffer until we have read the string marker and return the data parsed as an array of lines """

    encoding = encoding or self.encoding
    newline = newline or self.newline

    return self.reads_until(str_marker, encoding).split(newline)

  def read_before(self, marker, encoding = None):
    """ Read data into the buffer until we reach the marker and return all the data before the marker """

    encoding = encoding or self.encoding

    return self.read_until(marker)[:-len(marker)]

  def reads_before(self, str_marker, encoding = None):
    """ Read data into the buffer until we reach the marker string and return all the data before the marker as a string """

    encoding = encoding or self.encoding

    return self.reads_until(str_marker)[:-len(str_marker)]

  def readln_before(self, str_marker, encoding = None, newline = None):
    """ Read data into the buffer until we reach the marker string and return all the data before the marker as an array of lines """

    encoding = encoding or self.encoding
    newline = newline or self.newline

    return self.reads_before(str_marker, encoding).split(newline)

  def write(self, data):
    """ Write some bytes to the socket """

    self.socket.send(data)

  def writes(self, str_data):
    """ Write a string to the socket """

    self.socket.send(str.encode(str_data))

  def writeln(self, line, newline = None):
    """ Write a string and end with a newline character """

    newline = newline or self.newline

    self.writes(line + newline)

  def write_after(self, data, marker):
    """ Write data after reaching a marker """

    self.read_until(marker)
    self.write(data)

  def writes_after(self, str_data, str_marker):
    """ Write a string after reaching a marker string """

    self.reads_until(str_marker)
    self.writes(str_data)

  def writeln_after(self, line, str_marker):
    """ Write a string followed by a newline after reaching a marker string """

    self.reads_until(str_marker)
    self.writeln(line)

  def close(self):
    """ Close the connection to the socket """

    self.socket.close()

sample.py

#!/usr/bin/env python3
# encoding: utf-8

from netcat import Netcat

# Start a new Netcat() instance.
nc = Netcat('127.0.0.1', 1234)

# Get to the prompt (discard the read text).
nc.read_until_str('> ')

# Send a command.
nc.write_str('whoami' + '\n')

# Extract the output of the command.
output = nc.read_str_before('\n> ')
print(f"You are {output}") # E.g. "You are testuser"

# Close the connection.
nc.close()