socklib.py

import socket
import time

class Socket:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.buf = b""

        self.s.connect((host, port))

    def __enter__(self):
        return self

    def __exit__(self, typ, value, traceback):
        self.s.close()

    def recv(self, n=4096):
        return self.s.recv(n)

    def recvuntil(self, target):
        if isinstance(target, str):
            target = target.encode()
        while target not in self.buf:
            self.buf += self.s.recv(1024)

        pos = self.buf.find(target)
        rslt = self.buf[:pos + len(target)]
        self.buf = self.buf[pos + len(target):]
        return rslt

    def recvline(self):
        return self.recvuntil(b'\n')

    def send(self, data):
        if isinstance(data, str):
            data = data.encode()
        self.s.send(data)

    def sendline(self, data):
        if isinstance(data, str):
            data = data.encode()
        self.s.send(data + b'\n')

    def sendafter(self, target, data):
        self.recvuntil(target)
        self.send(data)

    def sendlineafter(self, target, data):
        self.recvuntil(target)
        self.sendline(data)

def parse_nc_cmd(s: str):
    if not s.startswith("nc"):
        raise ValueError("invalid argument")
    cmd = s.split()
    host = cmd[1]
    port = int(cmd[2])
    return Socket(host, port)