twgs-proxy/examples/tcp-proxy2-rev0.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.python.logfile import DailyLogFile

# from twisted.enterprise import adbapi
import pendulum
from subprocess import check_output

# This isn't the best configuration, but it's simple
# and works.  Mostly.

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()

# Cleans all ANSI
cleaner = re.compile(r"\x1b\[[0-9;]*[A-Zmh]")
# Looks for ANSI (that should be considered to be a newline)
makeNL = re.compile(r"\x1b\[[0-9;]*[J]")


def treatAsNL(line):
    """ Replace any ANSI codes that would be better understood as newlines. """
    global makeNL
    return makeNL.sub("\n", line)


def cleanANSI(line):
    """ Remove all ANSI codes. """
    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 Game(protocol.Protocol):
    def __init__(self):
        # user is rlogin username
        self.user = None
        # buffer is used to capture the rlogin username
        self.buffer = ""

    def connectionMade(self):
        log.msg("Client: connected to peer")
        self.queue_player = self.factory.queue_player
        self.queue_game = self.factory.queue_game
        self.setPlayerReceived()
        # self.queue_twgs.get().addCallback(self.serverDataReceived)

    def setPlayerReceived(self):
        """ Get deferred from client queue, callback clientDataReceived. """
        self.queue_player.get().addCallback(self, playerDataReceived)

    def playerDataReceived(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.setPlayerReceived()
            # 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)
        # self.queue_game.put(chunk)
        self.queue_player.put(chunk)

    def connectionLost(self, why):
        log.msg("Game connectionLost because: %s" % why)
        self.queue_game.put(False)

        # self.factory.queue_client.put(False)
        # self.queue_twgs = None
        self.transport.loseConnection()


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

    def __init__(self, twgs):
        self.twgs = twgs
        self.queue_player = twgs.queue_player
        self.queue_game = twgs.queue_game

    def closeIt(self):
        log.msg("closeIt")
        self.queue_player.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("connection to game failed: %s" % why)
        self.queue_player.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)


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

    def connectionMade(self):
        """ connected, setup queues. """
        self.queue_player = defer.DeferredQueue()
        self.queue_game = defer.DeferredQueue()
        self.setGameReceived()
        # self.action = ProxyAction(self)

        factory = GlueFactory(self)
        # Make connection to the game server
        reactor.connectTCP(HOST, PORT, factory, 5)

    def setGameReceived(self):
        """ Get deferred from client queue, callback clientDataReceived. """
        self.queue_game.get().addCallback(self.gameDataReceived)

    def logUser(self, user):
        """ We have the username. """
        now = pendulum.now()
        self.username = user
        filename = now.format("YYYY-MM-DD_HHmm") + "-" + user.lower()

        # Are we saving RAW output?
        if RAW:
            self.fpRaw = open(filename + ".raw", "ab")
        self.fpLines = open(filename + ".lines", "a")
        print("Log created:", now.to_rss_string(), "\n", file=self.fpLines)

    def setGame(self, game):
        """ We have the game (A-P) they are playing. """
        if self.game != game:
            log.msg("USER {0} ENTERED {1}".format(self.username, self.game))
            self.data = {}
        self.game = game

    def gotLine(self, line):
        """ We got a line from the server.

        This is ANSI filtered.  
        Backspaces have removed the character from the line.
        The line is unicode.  We don't need to decode it. ;)

        """
        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 gameDataReceived(self, chunk):
        """ Data received from the client/player. """
        if chunk is False:
            self.transport.loseConnection()
        else:
            if type(chunk) is tuple:
                # Special case where we want to send something to the user
                # but not have it logged.
                self.transport.write(chunk[0])
                self.setGameReceived()
            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.setGameReceived()

    def dataReceived(self, chunk):
        if self.action and self.action.isActive():
            # Do something completely different here
            self.action.received(chunk)
        else:
            # Did player activate hotkey?
            if chunk == b"~":
                self.action.activate(self.buffer)
            else:
                self.queue_game.put(chunk)

    def connectionLost(self, why):
        log.msg("lost connection %s" % why)
        self.queue_game.put(False)

        # Close log files, if open
        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(DailyLogFile("proxy.log", "."))
    log.msg("This is version: %s" % version)
    factory = protocol.Factory()
    factory.protocol = Player
    reactor.listenTCP(LISTEN_PORT, factory, interface=LISTEN_ON)
    reactor.run()