twgs-proxy/tcp-proxy.py

#!/usr/bin/env python3

import sys
import re

from twisted.internet import defer
from twisted.internet import protocol
from twisted.internet import reactor
from twisted.python import log
from twisted.enterprise import adbapi

# Connect to:
HOST = "127.0.0.1"
PORT = 2002
# Listen on:
LISTEN_PORT = 9999
LISTEN_ON = "127.0.0.1"


class ProxyClientProtocol(protocol.Protocol):
    def __init__(self):
        self.user = None
        self.dbinit = False
        self.db = None
        self.buffer = ""
        self.lines = list()

    def connectionMade(self):
        log.msg("Client: connected to peer")
        self.cli_queue = self.factory.cli_queue
        self.cli_queue.get().addCallback(self.serverDataReceived)

    def serverDataReceived(self, chunk):
        # rlogin looks like this:  \x00 password \x00 username \x00 terminal/speed \x00
        # b'\x00up2lat3\x00bugz\x00ansi-bbs/115200\x00'
        # We're looking for 4 \x00!
        # TODO: Line processing, and line cleaning (remove ANSI color codes)
        if chunk is False:
            self.cli_queue = None
            log.msg("Client: disconnecting from peer")
            self.factory.continueTrying = False
            self.transport.loseConnection()
        else:
            self.buffer += chunk.decode("utf-8")
            if self.user is None:
                # Ok, process this
                # self.buffer += chunk.decode('utf-8')
                # We don't have the username yet
                parts = self.buffer.split("\x00")
                if len(parts) >= 5:
                    # Got it!
                    self.user = parts[1]
                    log.msg("User: {0}".format(self.user))
                    # Reset buffer -- remove everything before last \x00
                    zpos = self.buffer.rindex("\x00")
                    self.buffer = self.buffer[zpos + 1 :]
                    # init sqlite db using the username
            else:
                # process the buffer
                # Handle backspaces by deleting previous character.
                pos = self.buffer.find("\x08")
                while pos != -1:
                    # Ok, there's a backspace, so:
                    self.buffer = self.buffer[0 : pos - 1] + self.buffer[pos + 1 :]
                    pos = self.buffer.find("\x08")

                # There won't be ansi color codes from the user!
                # There will be \r\n  (If I remember correctly)
                pos = self.buffer.find("\r")
                while pos != -1:
                    line = self.buffer[0:pos]
                    self.lines.append(line)
                    log.msg("L> {0}".format(line))
                    self.buffer = self.buffer[pos + 1 :]
                    pos = self.buffer.find("\r")

                #
                #
                # Strip out ANSI color codes
                # self.buffer = re.sub(r'\x1b[\d;?\d+m', '', self.buffer)
                # Process lines ...

            self.transport.write(chunk)
            self.cli_queue.get().addCallback(self.serverDataReceived)

        # elif b"$" == chunk:
        #     self.factory.svr_queue.put(b"HELLO.\r\n")
        #     self.cli_queue.get().addCallback(self.serverDataReceived)
        # elif self.cli_queue:
        #     log.msg("Client: writing %d bytes to peer" % len(chunk))
        #     log.msg(">>", repr(chunk))
        #     self.transport.write(chunk)
        #     self.cli_queue.get().addCallback(self.serverDataReceived)
        # else:
        #     self.factory.cli_queue.put(chunk)

    def dataReceived(self, chunk):
        # log.msg("Client: %d bytes received from peer" % len(chunk))
        # clean, strip ANSI, etc.

        log.msg("<<", repr(chunk))
        self.factory.svr_queue.put(chunk)

    def connectionLost(self, why):
        log.msg("connectionLost because: %s" % why)
        # if self.cli_queue:
        #     self.cli_queue = None
        #     log.msg("Client: peer disconnect unexpectedly")
        self.factory.svr_queue.put(False)
        self.cli_queue = None
        self.transport.loseConnection()


class ProxyClientFactory(protocol.ClientFactory):
    maxDelay = 10
    protocol = ProxyClientProtocol

    def __init__(self, svr_queue, cli_queue):
        self.svr_queue = svr_queue
        self.cli_queue = cli_queue


# ProxyServer is created for each connection


class ProxyServer(protocol.Protocol):
    def connectionMade(self):
        self.svr_queue = defer.DeferredQueue()
        self.cli_queue = defer.DeferredQueue()
        self.svr_queue.get().addCallback(self.clientDataReceived)

        factory = ProxyClientFactory(self.svr_queue, self.cli_queue)
        reactor.connectTCP(HOST, PORT, factory)

    def clientDataReceived(self, chunk):
        if chunk is False:
            self.transport.loseConnection()
        else:
            log.msg("Server: writing %d bytes to original client" % len(chunk))
            self.transport.write(chunk)
            self.svr_queue.get().addCallback(self.clientDataReceived)

    def dataReceived(self, chunk):
        log.msg("Server: %d bytes received" % len(chunk))
        self.cli_queue.put(chunk)

    def connectionLost(self, why):
        log.msg("lost connection %s" % why)
        self.cli_queue.put(False)


if __name__ == "__main__":
    log.startLogging(sys.stdout)
    factory = protocol.Factory()
    factory.protocol = ProxyServer
    reactor.listenTCP(LISTEN_PORT, factory, interface=LISTEN_ON)
    reactor.run()