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

from config import *

# Connect to:
# HOST = "twgs"
# PORT = 2002
# Listen on:
# LISTEN_PORT = 2002
# LISTEN_ON = "0.0.0.0"

cleaner = re.compile(r"\x1b\[[0-9;]*[A-Zmh]")
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()