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

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 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)
        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",
                    )
                )

    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)

                # 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 chunk == b"~":
            self.queue_client.put(
                (b"\r\n**********\r\nTWGS Proxy is almost awake...\r\n",)
            )
        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()