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.internet import task
from twisted.python import log
from twisted.python.logfile import DailyLogFile
import pendulum
from subprocess import check_output

from colorama import Fore, Back, Style

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

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

# Extract the version information from git.
# The match gives us only tags starting with v[0-9]*  Using anything else trips up on double digits.
version = check_output(
    [
        "git",
        "describe",
        "--abbrev=8",
        "--long",
        "--tags",
        "--dirty",
        "--always",
        "--match",
        "v[0-9]*",
    ],
    universal_newlines=True,
).strip()


def merge(color_string):
    """ Given a string of colorama ANSI, merge them if you can. """
    return color_string.replace("m\x1b[", ";")


# https://en.wikipedia.org/wiki/ANSI_escape_code

# Cleans all ANSI
cleaner = re.compile(r"\x1b\[[0-9;]*[A-Zmh]")

# Looks for ANSI (that should be considered to be a newline)
# This needs to see what is send when something enters / leaves
# the player's current sector.  (That doesn't work/isn't
# detected.  NNY!)  It is "\x1b[K" Erase in Line!
makeNL = re.compile(r"\x1b\[[0-9;]*[JK]")


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 Observer(object):
    def __init__(self):
        self.dispatch = {}

    def emit(self, signal, message):
        """ emit a signal, return True if sent somewhere. """
        if signal in self.dispatch:
            # something to do
            ret = False
            for listener in self.dispatch[signal]:
                ret = True
                reactor.callLater(0, listener, message)
            return ret
        return False

    def connect(self, signal, func):
        """ Connect a signal to a given function. """
        if not signal in self.dispatch:
            self.dispatch[signal] = []

        self.dispatch[signal].append(func)

    def disconnect(self, signal, func):
        """ Disconnect a signal with a certain function. """
        if signal in self.dispatch:
            self.dispatch[signal].remove(func)
            if len(self.dispatch[signal]) == 0:
                self.dispatch.pop(signal)

    def get_funcs(self, signal):
        """ Gives a copy of the dispatch for a given signal. """
        if signal in self.dispatch:
            return list(self.dispatch[signal])
        else:
            return []

    def set_funcs(self, signal, funcs):
        """ Replaces the dispatch for a given signal. """
        if signal in self.dispatch:
            if len(funcs) == 0:
                self.dispatch.pop(signal)
            else:
                self.dispatch = list(funcs)
        else:
            if len(funcs) != 0:
                self.dispatch = list(funcs)


class MCP(object):
    def __init__(self, game):
        self.game = game
        self.queue_game = None
        # we don't have this .. yet!
        self.prompt = None
        self.observer = None
        self.keepalive = None

    def finishSetup(self):
        if self.queue_game is None:
            self.queue_game = self.game.queue_game
        if self.observer is None:
            self.observer = self.game.observer
        self.observer.connect("hotkey", self.activate)

    def stayAwake(self):
        """ Send a space to the game to keep it alive/don't timeout. """
        log.msg("Gameserver, stay awake.")
        self.game.queue_player.put(" ")

    def activate(self, _):
        log.msg("MCP menu called.")
        # We want the raw one, not the ANSI cleaned getPrompt.
        prompt = self.game.buffer
        if not self.prompt is None:
            # silly, we're already active
            log.msg("I think we're already active.  Ignoring request.")
            return
        # Or will the caller setup/restore the prompt?
        self.prompt = prompt
        # queue_game = to player
        self.displayMenu()
        self.observer.connect("player", self.fromPlayer)
        # TODO:  Add background "keepalive" event so the game doesn't time out on us.
        self.keepalive = task.LoopingCall(self.stayAwake)
        self.keepalive.start(30)

    def displayMenu(self):
        nl = "\n\r"
        c = merge(Style.BRIGHT + Fore.YELLOW + Back.BLUE)
        r = Style.RESET_ALL
        c1 = merge(Style.BRIGHT + Fore.BLUE)
        c2 = merge(Style.NORMAL + Fore.BLUE)

        self.queue_game.put(nl + c + "TradeWars Proxy active." + r + nl)
        self.queue_game.put(
            " " + c1 + "T" + c2 + " - " + c1 + "Display current Time" + nl
        )
        self.queue_game.put(" " + c1 + "P" + c2 + " - " + c1 + "Port CIM Report" + nl)
        self.queue_game.put(" " + c1 + "X" + c2 + " - " + c1 + "eXit" + nl)
        self.queue_game.put("   " + c + "-=>" + r + " ")

    def fromPlayer(self, chunk):
        """ Data from player (in bytes).  """
        chunk = chunk.decode("utf-8", "ignore")

        nl = "\n\r"
        c = merge(Style.BRIGHT + Fore.YELLOW + Back.BLUE)
        r = Style.RESET_ALL
        c1 = merge(Style.BRIGHT + Fore.BLUE)
        c2 = merge(Style.NORMAL + Fore.BLUE)

        key = chunk.upper()
        if key == "T":
            self.queue_game.put(c + key + r + nl)
            now = pendulum.now()
            log.msg("Time")
            self.queue_game.put(
                nl + c1 + "It is currently " + now.to_datetime_string() + "." + nl
            )
            self.displayMenu()
        elif key == "P":
            log.msg("Port")
            self.queue_game.put(c + key + r + nl)
            self.queue_game.put(nl + c + "NO, NOT YET!" + r + nl)
            self.displayMenu()
        elif key == "X":
            log.msg('"Quit, return to "normal".  (Whatever that means!)')
            self.queue_game.put(c + key + r + nl)
            self.observer.disconnect("player", self.fromPlayer)
            self.queue_game.put(nl + c1 + "Returning to game" + c2 + "..." + r + nl)
            self.queue_game.put(self.prompt)
            self.prompt = None
            self.keepalive.stop()
            self.keepalive = None
        else:
            if key.isprintable():
                self.queue_game.put(r + nl)
                self.queue_game.put("Excuse me? I don't understand '" + key + "'." + nl)
                self.displayMenu()


class Game(protocol.Protocol):
    def __init__(self):
        self.buffer = ""
        self.game = None
        self.usergame = (None, None)
        self.to_player = True
        self.mcp = MCP(self)

    def connectionMade(self):
        log.msg("Connected to Game Server")
        self.queue_player = self.factory.queue_player
        self.queue_game = self.factory.queue_game
        self.observer = self.factory.observer
        self.setPlayerReceived()
        self.observer.connect("user-game", self.show_game)
        self.mcp.finishSetup()

    def show_game(self, game):
        self.usergame = game
        log.msg("## User-Game:", game)

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

    def playerDataReceived(self, chunk):
        if chunk is False:
            self.queue_player = None
            log.msg("Player: disconnected, close connection to game")
            # I don't believe I need this if I'm using protocol.Factory
            self.factory.continueTrying = False
            self.transport.loseConnection()
        else:
            # Pass received data to the server
            if type(chunk) == str:
                self.transport.write(chunk.encode())
            else:
                self.transport.write(chunk)
            self.setPlayerReceived()

    def lineReceived(self, line):
        """ line received from the game. """
        if LOG_LINES:
            log.msg(">> [{0}]".format(line))

        if "TWGS v2.20b" in line and "www.eisonline.com" in line:
            self.queue_game.put(
                "TWGS Proxy build "
                + version
                + " 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
            if not self.game is None:
                # We were in a game.
                self.game = None
                self.observer.emit("user-game", (self.factory.player.user, self.game))

        if "Selection (? for menu): " in line:
            game = line[-1]
            if game >= "A" and game < "Q":
                self.game = game
                log.msg("Game: {0}".format(self.game))
                self.observer.emit("user-game", (self.factory.player.user, self.game))

        self.observer.emit("game-line", line)

    def getPrompt(self):
        """ Return the current prompt, stripped of ANSI. """
        return cleanANSI(self.buffer)

    def dataReceived(self, chunk):
        """ Data received from the Game. 
        
        Remove backspaces.
        Treat some ANSI codes as NewLine.        
        Remove ANSI.
        Break into lines.
        Trim out carriage returns.
        Call lineReceived().
        
        "Optionally" pass data to player.
        FUTURE: trigger on prompt. [cleanANSI(buffer)]
        """

        # Sequence error:
        # If I don't put the chunk(I received) to the player.
        # anything I display -- lineReceive() put() ... would
        # be out of order.  (I'd be responding -- before it
        # was displayed to the user.)

        if self.to_player:
            self.queue_game.put(chunk)

        self.buffer += chunk.decode("utf-8", "ignore")

        # Process any backspaces
        while "\x08" in self.buffer:
            part = self.buffer.partition("\x08")
            self.buffer = part[0][:-1] + part[2]

        # Treat some ANSI codes as a newline
        self.buffer = treatAsNL(self.buffer)

        # Break into lines
        while "\n" in self.buffer:
            part = self.buffer.partition("\n")
            line = part[0].replace("\r", "")
            # Clean ANSI codes from line
            line = cleanANSI(line)
            self.lineReceived(line)
            self.buffer = part[2]

        self.observer.emit("prompt", self.getPrompt())

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


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

    def __init__(self, player):
        self.player = player
        self.queue_player = player.queue_player
        self.queue_game = player.queue_game
        self.observer = player.observer

    def closeIt(self):
        log.msg("closeIt")
        self.queue_game.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_game.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.user = None
        self.observer = Observer()

    def connectionMade(self):
        """ connected, setup queues. 
        
        queue_player is data from player.
        queue_game is data to player. (possibly from game)
        """
        self.queue_player = defer.DeferredQueue()
        self.queue_game = defer.DeferredQueue()
        self.setGameReceived()

        # Connect GlueFactory to this Player object.
        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 gameDataReceived(self, chunk):
        """ Data received from the game. """
        if chunk is False:
            self.transport.loseConnection()
        else:
            if type(chunk) == bytes:
                self.transport.write(chunk)
            elif type(chunk) == str:
                self.transport.write(chunk.encode())
            else:
                log.err("gameDataReceived: type (%s) given!", type(chunk))
                self.transport.write(chunk)
            self.setGameReceived()

    def dataReceived(self, chunk):
        if self.user is None:
            self.buffer += chunk.decode("utf-8", "ignore")

            parts = self.buffer.split("\x00")
            if len(parts) >= 5:
                # rlogin we have the username
                self.user = parts[1]
                log.msg("User: {0}".format(self.user))
                zpos = self.buffer.rindex("\x00")
                self.buffer = self.buffer[zpos + 1 :]
                # but I don't need the buffer anymore, so:
                self.buffer = ""
                # Pass user value on to whatever needs it.
                self.observer.emit("user", self.user)
                # Unfortunately, the ones interested in this don't exist yet.

        if not self.observer.emit("player", chunk):
            # Was not dispatched.  Send to game.
            self.queue_player.put(chunk)

        if chunk == b"~":
            # Selection (? for menu):   (the game server menu)
            # Enter your choice:  (game menu)
            # Command [TL=00:00:00]:[1800] (?=Help)? :  <- YES!
            self.observer.emit("hotkey", None)

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

    def connectionFailed(self, why):
        log.msg("connectionFailed: %s" % why)


if __name__ == "__main__":
    if LOGFILE:
        log.startLogging(DailyLogFile("proxy.log", "."))
    else:
        log.startLogging(sys.stdout)

    log.msg("This is version: %s" % version)
    factory = protocol.Factory()
    factory.protocol = Player
    reactor.listenTCP(LISTEN_PORT, factory, interface=LISTEN_ON)
    reactor.run()