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 subprocess import check_output

RAW = True

try:
    from config_dev import *
except ModuleNotFoundError as E:
    from config import *

# from config import *
# from config_dev import *

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

version = check_output(
    [
        "git",
        "describe",
        "--long",
        "--tags",
        #        "--dirty",
        "--always",
        "--match",
        "v[0-9]\.[0-9]\.[0-9]",
    ],
    universal_newlines=True,
).strip()

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)

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

    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

    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 ProxyAction:
    def __init__(self, queue_client, queue_twgs):
        self.active = False
        self.queue_client = queue_client
        self.queue_twgs = queue_twgs
        self.buffer = ""
        self.prompt = ""

    def isActive(self):
        return self.active

    def keepAlive(self):
        if self.active:
            self.queue_twgs.put(b" ")
            reactor.callLater(30, self.keepAlive)

    def activate(self, prompt):
        self.active = True
        self.prompt = prompt
        self.send("\r\n**********\r\nTWGS Proxy ready...\r\n")
        self.send("Prompt seen: [{0}] ...\r\n".format(self.prompt))
        self.send("(T) Display Time\r\n")
        self.send("(Q) Quit\r\n")
        self.send("  --==> ")
        reactor.callLater(30, self.keepAlive)

    def server(self, line):
        pass

    def send(self, text):
        self.queue_client.put((text.encode(),))

    def received(self, chunk):
        # self.buffer += chunk.encode('UTF-8', 'ignore')
        text = chunk.decode("utf-8", "ignore").upper()

        if text == "T":
            now = pendulum.now('America/New_York')
            self.send("\r\nThe time is: {0}.\r\n".format(now.to_rss_string()))
        if text == "Q":
            self.send("\r\nReturning to TWGS.\r\n{0}".format(self.prompt))
            self.active = False


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

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

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

    def logUser(self, user):
        now = pendulum.now()
        self.user = user
        filename = now.format("YYYY-MM-DD_HHmm") + "-" + user.lower()
        if RAW:
            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)
        if "TWGS v2.20b" in line:
            if "www.eisonline.com" in line:
                # Must not be unicode

                # Is there a way to NOT have this logged?
                self.queue_client.put(
                    (
                        b"TWGS Proxy build "
                        + version.encode()
                        + b" is active. \x1b[1;34m~\x1b[0m to activate.\n\r\n\r",
                    )
                )
        if "TradeWars Game Server" in line:
            if "Copyright (C) EIS" in line:
                # We are not in a game
                self.game = ""
        if "Selection (? for menu): " in line:
            game = line[-1]
            if game >= "A" and game < "Q":
                self.game = game
                log.msg("USER {0} ENTERED {1}".format(self.user, self.game))

    def clientDataReceived(self, chunk):
        if chunk is False:
            self.transport.loseConnection()
        else:
            if type(chunk) is tuple:
                self.transport.write(chunk[0])
                self.queue_client.get().addCallback(self.clientDataReceived)
            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)
                # I think I need something else in here.  When something enters or leaves the sector
                # The message isn't shown on it's own line.  I think they are using ANSI codes to
                # clear the line.  (Which certainly would be faster!)

                # 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))
        if self.action and self.action.isActive():
            # Do something completely different here
            self.action.received(chunk)
        else:
            if chunk == b"~":
                self.action.activate(self.buffer)
            else:
                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)
    log.msg("This is version: %s" % version)
    factory = protocol.Factory()
    factory.protocol = TWGSServer
    reactor.listenTCP(LISTEN_PORT, factory, interface=LISTEN_ON)
    reactor.run()