PyCat: a netcat implementation in Python3.x v2

Question

This is a follow up question from this question

Intro

netcat is an all-round tool used with many applicable features

My last try felt a bit rushed, and should've improved before posting here. But this time, I am happy with the result.

CHANGELOG

• kwargs
• @use_ssl decorator
• Multi-Platform (Posix, *nix, Windows)
• Improved code structure
• Download
• Upload

I tried adding a context manager, but couldn't really make it work in an elegant way.

Any and all reviews are welcome.

Example

server

$ python pycat.py -lsp 8080
[*] Incoming connection from 127.0.0.1:53391
username@hostame PyCat C:\dev\Pycat
> echo hooooooi
hooooooi
username@hostame PyCat C:\dev\PyCat
> cd ../

username@hostame PyCat C:\dev
> exit

client

python pycat.py -si localhost -p 8080

Code

import argparse
import datetime
from functools import wraps
import socket
from ssl import wrap_socket, create_default_context, CERT_NONE
import sys
import subprocess
import tempfile
import os
import re

from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.backends import default_backend

SUCCES_RESPONSE = b"Command succesfully completed"

def ssl_server(func):
    @wraps(func)
    def wrapper(inst, *args):
        inst.socket.bind((inst.host, inst.port))
        inst.socket.listen(0)
        if inst.ssl:
            inst.context = create_default_context()
            inst.key, inst.cert = inst.generate_temp_cert()
            inst.socket = wrap_socket(
                inst.socket,
                server_side=True, 
                certfile=inst.cert, 
                keyfile=inst.key
            )
        func(inst, *args)
    return wrapper

def ssl_client(func):
    @wraps(func)
    def wrapper(inst, *args):
        inst.socket.connect((inst.host, inst.port))
        if inst.ssl:
            inst.context = create_default_context()
            inst.context.check_hostname = False
            inst.context.verify_mode = CERT_NONE
            inst.socket = wrap_socket(inst.socket)
        func(inst, *args)
    return wrapper

class PyCatBase():
    def __init__(self, **kwargs):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.port = kwargs['port']
        self.host = kwargs['host'] or "0.0.0.0"
        self.operating_system = os.name == "nt"
        self.upload = kwargs['upload']
        self.download = kwargs['download']
        self.timeout = kwargs['timeout']
        self.ssl = kwargs['ssl']

    def exit(self):
        self.socket.close()
        sys.exit(0)

    def read(self, connection, length=1024):
        response = b""
        while True:
            data = connection.recv(length)
            response += data
            if len(data) < length:
                break
        return response.decode("utf-8").rstrip()

    def upload_file(self, connection, file):
        with open(file, "rb") as f:
            connection.send(f.read())

    def download_file(self, connection, file):
        recieved = self.read(connection)
        with open(file, "wb") as f:
            f.write(recieved)

    @staticmethod
    def generate_temp_cert():
        _, key_path = tempfile.mkstemp()
        _, cert_path = tempfile.mkstemp()
        name_attributes = [
            x509.NameAttribute(NameOID.COUNTRY_NAME, "OK"),
            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "OK"),
            x509.NameAttribute(NameOID.LOCALITY_NAME, "OK"),
            x509.NameAttribute(NameOID.ORGANIZATION_NAME, "OK"),
            x509.NameAttribute(NameOID.COMMON_NAME, "PyCat")
        ]

        key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=default_backend()
        )

        with open(key_path, "wb") as f:
            f.write(
                key.private_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PrivateFormat.TraditionalOpenSSL,
                    encryption_algorithm=serialization.NoEncryption()
                )
            )

        subject = issuer = x509.Name(name_attributes)

        cert = x509.CertificateBuilder()\
                    .subject_name(subject)\
                    .issuer_name(issuer)\
                    .public_key(key.public_key())\
                    .serial_number(x509.random_serial_number())\
                    .not_valid_before(datetime.datetime.utcnow())\
                    .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))

        cert = cert.sign(key, hashes.SHA256(), default_backend())
        with open(cert_path, "wb") as f:
            f.write(
                cert.public_bytes(serialization.Encoding.PEM)
            )

        return key_path, cert_path

class PyCatServer(PyCatBase):
    def __init__(self, **kwargs):
        super(PyCatServer, self).__init__(**kwargs)

    def create_prompt_string(self):
        self.client.send(b"cd") if self.operating_system else self.client.send(b"pwd")
        pwd = self.read(self.client)
        self.client.send(b"whoami")
        whoami = self.read(self.client)
        self.client.send(b"hostname")
        hostname = self.read(self.client)
        return f"{whoami}@{hostname} PyCat {pwd}\n> "

    @ssl_server
    def main(self):
        if self.timeout > 0:
            self.socket.settimeout(self.timeout)
        self.client, addr = self.socket.accept()
        print(f"[*] Incomming connection from {':'.join(map(str, addr))}")
        self.handle_client()

    def handle_client(self):
        if self.upload is not None:
            self.upload_file(self.client, self.upload)
        elif self.download is not None:
            self.download_file(self.client, self.download)
        else:
            while True:
                prompt_string = self.create_prompt_string()
                buf = input(prompt_string)
                self.client.send(buf.encode("utf-8"))
                if buf == "exit":
                    break
                print(self.read(self.client))
        self.exit()

class PyCatClient(PyCatBase):
    def __init__(self, **kwargs):
        super(PyCatClient, self).__init__(**kwargs)

    def change_dir(self, path):
        try:
            os.chdir(path)
            return SUCCES_RESPONSE
        except FileNotFoundError as e:
            return str(e).encode("utf-8")

    def exec_command(self, command):
        try:
            return subprocess.check_output(command, stderr=subprocess.STDOUT, shell=True)
        except Exception as e:
            return str(e).encode("utf-8")

    def handle_command(self, command):
        if command == "exit":
            self.exit()
        change_dir = re.match(r'cd(?:\s+|$)(.*)', command)   
        if change_dir and change_dir.group(1):
            return self.change_dir(change_dir.group(1))
        return self.exec_command(command)

    @ssl_client
    def main(self):
        if self.timeout > 0:
            self.socket.settimeout(self.timeout)

        if self.upload is not None:
            self.upload_file(self.socket, self.upload)
        elif self.download is not None:
            self.download_file(self.socket, self.download)
        else:
            while True:
                cmd = self.read(self.socket)
                response = self.handle_command(cmd)
                if len(response) > 0:
                    self.socket.send(response)

def parse_arguments():
    parser = argparse.ArgumentParser(usage='%(prog)s [options]',
                                     description='PyCat @Ludisposed',
                                     formatter_class=argparse.RawDescriptionHelpFormatter,
                                     epilog='Examples:\npython3 pycat.py -lvp 443\npython3 pycat.py -i localhost -p 443')
    parser.add_argument('-l', '--listen', action="store_true", help='Listen')
    parser.add_argument('-s', '--ssl', action="store_true", help='Encrypt connection')
    parser.add_argument('-p', '--port', type=int, help='Port to listen on')
    parser.add_argument('-i', '--host', type=str, help='Ip/host to connect to')
    parser.add_argument('-d', '--download', type=str, help='download file')
    parser.add_argument('-u', '--upload', type=str, help='upload file')
    parser.add_argument('-t', '--timeout', type=int, default=0, help='timeout')
    args = parser.parse_args()

    if (args.listen or args.host) and not args.port:
        parser.error('Specify which port to connect to')
    elif not args.listen and not args.host:
        parser.error('Specify --listen or --host')
    return args

if __name__ == '__main__':
    args = parse_arguments()
    pycat_class = PyCatServer if args.listen else PyCatClient
    pycat = pycat_class(**vars(args))
    pycat.main()

Answer 1

I tried adding a context manager, but couldn't really make it work in an elegant way.

class PyCatBase():
    def __enter__(self):
        return self
    def __exit__(exc_type, exc_val, exc_tb):
        self.socket.close()
        return False

# ...
if __name__ == '__main__':
    args = parse_arguments()
    pycat_class = PyCatServer if args.listen else PyCatClient
    pycat = pycat_class(**vars(args))
    with pycat:
        pycat.main()

Other things. SUCCES_RESPONSE should be spelled SUCCESS_RESPONSE. Similarly, there's a typo in that string.

This:

self.operating_system = os.name == "nt"

suggests one of two things. Either operating_system should be named is_windows, or you need to change it to simply self.operating_system = os.name.

This:

def __init__(self, **kwargs):

is only a good idea in other, limited, contexts (for example, if you're extending a class with a highly complex initializer). Don't do that, here. Spell out your args. Having implicit kwargs hurts you and your users in a number of ways, including kneecapping your IDE's static analysis efforts.

Here:

cert = x509.CertificateBuilder()\
                .subject_name(subject)\
                .issuer_name(issuer)\
                .public_key(key.public_key())\
                .serial_number(x509.random_serial_number())\
                .not_valid_before(datetime.datetime.utcnow())\
                .not_valid_after(datetime.datetime.utcnow() +datetime.timedelta(days=365))

the generally accepted thing to do rather than a handful of newline continuations is to surround the thing in parens.

This:

change_dir = re.match(r'cd(?:\s+|$)(.*)', command)  

should have its regex pre-compiled in __init__, since you call it for every command.