twgs-proxy/examples/tcp-proxy3.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:
    from config import *

# from config 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:
            if self.user is None:
                # Decode the rlogin data
                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)

            # Pass received data to the server
            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, twgs):
        self.active = False
        self.twgs = twgs
        self.queue_client = twgs.queue_client
        self.queue_twgs = twgs.queue_twgs
        self.buffer = ""
        self.prompt = ""
        self.rptstate = 0

    def isActive(self):
        return self.active

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

    def menu(self):
        self.send("\r\n**********\r\nTWGS Proxy ready...\r\n")
        self.send("(T) Display Time\r\n")
        self.send("(P) CIM Port Report\r\n")
        self.send("(Q) Quit\r\n")
        self.send("  --==> ")

    def activate(self, prompt):
        cleaned = cleanANSI(prompt)
        if cleaned.startswith("Command [TL=0"):
            self.active = True
            self.prompt = prompt
            self.menu()
            reactor.callLater(30, self.keepAlive)
        else:
            self.send("\a")

    def sendtwgs(self, text):
        self.queue_twgs.put(text.encode())

    def server(self, line):
        # 2019-11-20 18:37:28-0500 [PlayerProtocol,client] >>> [: ]
        # 2019-11-20 18:37:28-0500 [PlayerProtocol,client] >>> [ 436   2870 100% - 1520 100% - 2820 100% ]
        # 2019-11-20 18:37:28-0500 [PlayerProtocol,client] >>> []
        # 2019-11-20 18:37:42-0500 [PlayerProtocol,client] >>> [: ENDINTERROG]
        # 2019-11-20 18:37:42-0500 [PlayerProtocol,client] >>> []

        if self.rptstate == 1:
            if line.startswith(":"):
                self.sendtwgs("R")
                self.rptstate == 2
        if self.rptstate == 2:
            if line.startswith(":"):
                self.sendtwgs("Q")
                self.rptstate == 3
            else:
                log.msg("[[{0}]]".format(line))
        if self.rptstate == 3:
            if line == ": ENDINTERROG":
                self.rptstate == 3
                self.twgs.passon = True
                self.menu()

    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 == "P":
            # Port Report
            self.sendtwgs("^")
            self.twgs.passon = False
            self.rptstate = 1

        if text == "Q":
            self.send("\r\nReturning to TWGS.\r\n{0}".format(self.prompt))
            self.active = False
            self.twgs.passon = True


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

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

        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 setGame(self, game):
        if self.game != game:
            log.msg("USER {0} ENTERED {1}".format(self.user, self.game))
            self.data = {}
        self.game = game

    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 and "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.setGame(game)
        # If we're not passing it on to the user, we better be looking at it.
        if self.action and not self.passon:
            self.action.server(line)

    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))
                if self.passon:
                    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()