Last active
August 8, 2020 22:40
-
-
Save Cidolfas/0f51ca05d603840eb21a85c87b353679 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#! 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 :(") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment