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
import pendulum

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

cleaner = re.compile(r"\x1b\[[0-9;]*[A-Zm]")
makeNL = re.compile(r"\x1b\[[0-9;]*[J]")


def treatAsNL(line):
    global makeNL
    return makeNL.sub("\n", line)


def cleanANSI(line):
    global cleaner
    return cleaner.sub("", line)
    # return re.sub(r'\x1b\[([0-9,A-Z]{1,2}(;[0-9]{1,2})?(;[0-9]{3})?)?[m|K]?', '', line)


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

    def connectionMade(self):
        log.msg("Client: connected to peer")
        self.queue_twgs = self.factory.queue_twgs
        self.queue_twgs.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.queue_twgs = None
            log.msg("Client: disconnecting from peer")
            self.factory.continueTrying = False
            self.transport.loseConnection()
        else:
            # self.buffer += chunk.decode("utf-8", 'ignore')
            if self.user is None:
                self.buffer += chunk.decode("utf-8", "ignore")

                # 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 :]
                    self.buffer = ""
                    # init sqlite db using the username
                    self.factory.getUser(self.user)

            # else:
            # process the buffer
            # Handle backspaces by deleting previous character.

            # Send the received data into the linereader for "automatic" line processing.

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

            self.transport.write(chunk)
            self.queue_twgs.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("<<", chunk.decode("utf-8", "ignore"))
        # log.msg("<<", repr(chunk))
        self.factory.queue_client.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.queue_client.put(False)
        self.queue_twgs = None
        self.transport.loseConnection()


class GlueFactory(protocol.ClientFactory):
    # class GlueFactory(protocol.Factory):
    maxDelay = 10
    protocol = PlayerProtocol

    def __init__(self, queue_client, queue_twgs, twgs):
        self.queue_client = queue_client
        self.queue_twgs = queue_twgs
        self.twgs = twgs
        self.fpRaw = None
        self.fpLines = None

    def closeIt(self):
        log.msg("closeIt")
        self.queue_client.put(False)

    def getUser(self, user):
        log.msg("getUser( %s )" % user)
        self.twgs.logUser(user)

    # This was needed when I replaced ClientFactory with Factory.
    # def clientConnectionLost(self, connector, why):
    #     log.msg("clientconnectionlost: %s" % why)
    #     self.queue_client.put(False)

    def clientConnectionFailed(self, connector, why):
        log.msg("connectionFailed: %s" % why)
        self.queue_client.put(b"Sorry!  I'm Unable to connect to the game server.\r\n")
        # syncterm gets cranky/locks up if we close this here.
        # (Because it is still sending rlogin information?)
        reactor.callLater(2, self.closeIt)


# ProxyServer is created for each connection


class TWGSServer(protocol.Protocol):
    def __init__(self):
        self.buffer = ""
        self.fpRaw = None
        self.fpLines = None

    def connectionMade(self):
        self.queue_twgs = defer.DeferredQueue()
        self.queue_client = defer.DeferredQueue()
        self.queue_client.get().addCallback(self.clientDataReceived)

        factory = GlueFactory(self.queue_client, self.queue_twgs, self)
        reactor.connectTCP(HOST, PORT, factory, 5)

    def logUser(self, user):
        now = pendulum.now()
        filename = now.format("YYYY-MM-DD_HHmm") + "-" + user.lower()
        self.fpRaw = open(filename + ".raw", "ab")
        self.fpLines = open(filename + ".lines", "a")
        # print("Log created:", now.to_rss_string(), "\n", file=self.fpRaw)
        print("Log created:", now.to_rss_string(), "\n", file=self.fpLines)

    def gotLine(self, line):
        # log.msg(">>> [{0}]".format(line.decode("utf-8", "ignore")))
        log.msg(">>> [{0}]".format(line))
        if self.fpLines is not None:
            print(line, file=self.fpLines)

    def clientDataReceived(self, chunk):
        if chunk is False:
            self.transport.loseConnection()
        else:
            if self.fpRaw is not None:
                self.fpRaw.write(chunk)
            self.buffer += chunk.decode("utf-8", "ignore")

            # Process any backspaces in the buffer

            while "\x08" in self.buffer:
                part = self.buffer.partition("\x08")
                self.buffer = part[0][:-1] + part[2]

            # Treat some ANSI codes as a newline (for purposes of Lines)
            self.buffer = treatAsNL(self.buffer)

            # Break the buffer into lines

            while "\n" in self.buffer:
                part = self.buffer.partition("\n")
                line = part[0].replace("\r", "")
                # Clean ANSI codes from the line
                line = cleanANSI(line)
                self.gotLine(line)
                self.buffer = part[2]

            # log.msg("Server: writing %d bytes to original client" % len(chunk))
            self.transport.write(chunk)
            self.queue_client.get().addCallback(self.clientDataReceived)

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

    def connectionLost(self, why):
        log.msg("lost connection %s" % why)
        self.queue_twgs.put(False)
        if self.fpRaw is not None:
            self.fpRaw.close()
            self.fpRaw = None
        if self.fpLines is not None:
            if self.buffer != "":
                print(self.buffer, file=self.fpLines)
            self.fpLines.close()
            self.fpLines = None

    def connectionFailed(self, why):
        log.msg("connectionFailed: %s" % why)


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