blaseballPlayback.py

#! python

import asyncio
from datetime import datetime
import sys
import pickle
import aiohttp
from aiohttp import web
import json
import argparse

digits = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

class TColors:
	END = '\33[0m'

	BLACK   = '\33[30m'
	RED     = '\33[31m'
	GREEN   = '\33[32m'
	YELLOW  = '\33[33m'
	BLUE    = '\33[34m'
	VIOLET  = '\33[35m'
	BEIGE   = '\33[36m'
	WHITE   = '\33[37m'
	GREY    = '\33[90m'
	RED2    = '\33[91m'
	GREEN2  = '\33[92m'
	YELLOW2 = '\33[93m'
	BLUE2   = '\33[94m'
	VIOLET2 = '\33[95m'
	BEIGE2  = '\33[96m'
	WHITE2  = '\33[97m'

class BlaseballRecorder:
	def __init__(self, filepath=None, uri=None):
		self.filepath = filepath or "blaseballGame.stream"
		self.uri = uri or "wss://blaseball.com/socket.io/?EIO=3&transport=websocket"
		self.messages = []

	async def ping(self, ws, ping_time = 25):
		while True:
			await asyncio.sleep(ping_time)
			await ws.send_str("2")
			# print(f"{TColors.GREY}Ping{TColors.END}")

	async def listen(self, ws):
		start_time = datetime.now()
		async for message in ws:
			if message.type == aiohttp.WSMsgType.TEXT:
				now = datetime.now()
				time_since_start = now - start_time
				self.messages.append((time_since_start.total_seconds(), message.data))
				print(f"{TColors.BLUE2}{time_since_start.total_seconds():7.2f}{TColors.END}: {message.data[0:50]}")

	async def record(self):
		async with aiohttp.ClientSession() as session:
			async with session.ws_connect(self.uri) as ws:
				print(f"{TColors.RED}Begin Recording{TColors.END}")
				await asyncio.gather(
					self.ping(ws),
					self.listen(ws)
				)

	def write(self):
		file = open(self.filepath, 'bw')
		pickle.dump(self.messages, file)
		print(f"{TColors.RED}Recorded {TColors.YELLOW2}{len(self.messages)}{TColors.RED} messages to {TColors.GREEN2}{self.filepath}{TColors.END}")
		sys.exit(0)

	def start(self):
		try:
			asyncio.get_event_loop().run_until_complete(self.record())
		except KeyboardInterrupt:
			self.write()

class BlaseballStreamer:
	def __init__(self, filepath=None):
		filepath = filepath or "blaseballGame.stream"
		file = open(filepath, 'br')
		self.messages = pickle.load(file)
		print(f"{TColors.RED}Loaded {TColors.YELLOW2}{len(self.messages)}{TColors.RED} messages from {TColors.GREEN2}{filepath}{TColors.RED}, last at {TColors.BLUE2}{self.messages[-1][0]:7.2f}{TColors.END}")
		
		self.http = False
		self.ws = True

		self.webapp = web.Application()
		self.http_games = []

	async def start_http_playback(self):
		start_time = datetime.now()
		next_message = 0
		print(f"{TColors.RED}Started HTTP playback{TColors.END}")
		while len(self.messages) > next_message:
			now = datetime.now()
			time_since_start = now - start_time
			seconds = time_since_start.total_seconds()
			if seconds >= self.messages[next_message][0]:
				print(f"{TColors.GREEN2}HTTP: {TColors.BLUE2}{self.messages[next_message][0]:7.2f}{TColors.END}")
				message = self.messages[next_message][1]
				while len(message) > 0 and message[0] in digits:
					message = message[1:]
				if len(message) > 0:
					m = json.loads(message)
					if type(m) is list:
						if (m[0] == "gameDataUpdate"):
							self.http_games = m[1]["schedule"]
				next_message += 1
			else:
				await asyncio.sleep(0.05)
		print(f"{TColors.RED}Finished HTTP playback{TColors.END}")

	async def start_http_server(self, webapp):
		runner = web.AppRunner(webapp)
		await runner.setup()
		site = web.TCPSite(runner, 'localhost', 8080)
		await site.start()
		while True:
			await asyncio.sleep(100)

	async def handle_http_games(self, request):
		return web.json_response(self.http_games)

	async def handle_ws(self, request):
		ws = web.WebSocketResponse()
		await ws.prepare(request)

		start_time = datetime.now()
		next_message = 0
		print(f"{TColors.RED}Started playback due to WS connection{TColors.END}")

		while len(self.messages) > next_message:
			now = datetime.now()
			time_since_start = now - start_time
			seconds = time_since_start.total_seconds()
			if seconds >= self.messages[next_message][0]:
				print(f"{TColors.GREEN}WS: {TColors.BLUE2}{self.messages[next_message][0]:7.2f}{TColors.END}")
				await ws.send_str(self.messages[next_message][1])
				next_message += 1
			else:
				await asyncio.sleep(0.05)
		print(f"{TColors.RED}Finished WS playback{TColors.END}")

	async def start_http(self):
		print(f"{TColors.BEIGE2}Starting stream in HTTP mode...{TColors.END}")
		await asyncio.gather(
			self.start_http_playback(),
			self.start_http_server(self.webapp)
		)

	async def start_ws(self):
		print(f"{TColors.BEIGE2}Starting stream in WS mode...{TColors.END}")
		await asyncio.gather(
			self.start_http_server(self.webapp)
		)

	async def start_both(self):
		print(f"{TColors.BEIGE2}Starting stream in HTTP and WS mode...{TColors.END}")
		await asyncio.gather(
			self.start_http_playback(),
			self.start_http_server(self.webapp)
		)

	def start(self):
		try:
			if self.http and self.ws:
				self.webapp.add_routes([web.get('/games', self.handle_http_games), web.get('/ws', self.handle_ws)])
				asyncio.get_event_loop().run_until_complete(self.start_both())
			if self.http:
				self.webapp.add_routes([web.get('/games', self.handle_http_games)])
				asyncio.get_event_loop().run_until_complete(self.start_http())
			elif self.ws:
				self.webapp.add_routes([web.get('/ws', self.handle_ws)])
				asyncio.get_event_loop().run_until_complete(self.start_ws())
			else:
				print(f"The birds prevent any streaming from happening. Check your options.")
		except KeyboardInterrupt:
			print(f"{TColors.BEIGE2}Received interrupt, shutting down stream{TColors.END}")
			sys.exit(0)

def handle_record(args):
	bbr = BlaseballRecorder(args.filepath, args.uri)
	bbr.start()

def handle_stream(args):
	bbs = BlaseballStreamer(args.filepath)
	bbs.http = args.http
	bbs.ws = args.ws or not args.http
	bbs.start()

parser = argparse.ArgumentParser(description="Record and Stream Blaseball game data feeds")
subparsers = parser.add_subparsers(title="Modes", dest="mode", required=True, description="Available modes", help="Which mode the program should run in")

record_parser = subparsers.add_parser("record", help="Records a live websocket stream to a file")
record_parser.set_defaults(func=handle_record)
record_parser.add_argument("-f", "--file", dest="filepath", help="Path to save the recording to")
record_parser.add_argument("--uri", help="URI to connect to using a websocket")

stream_parser = subparsers.add_parser("stream", help="Streams a given recording back to websocket or HTTP polling")
stream_parser.set_defaults(func=handle_stream)
stream_parser.add_argument("-f", "--file", dest="filepath", help="Path to read the stream data from")
stream_parser.add_argument("--http", action="store_true", help="Use HTTP playback to localhost:8080/game for polling implementations")
stream_parser.add_argument("--ws", action="store_true", help="Use WS playback to localhost:8080/ws for websocket implementations. If neither this or --http is given, default to websockets")

args = parser.parse_args()
if args.func:
	args.func(args)
else:
	print(f"The Umpire incinerated your command :(")