schlafana/schlafana.py

#!/usr/bin/env python
import pdb
import json
import logging
from pprint import pprint
from websocket import create_connection
from prometheus_client import Counter, Gauge, Summary, start_http_server
from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily, REGISTRY
import time

MESSAGES = { "GameInfo" : 0x00,
             "WorldUpdate" : 0x01,
             "Tick" : 0x10,
             "BotSpawn" : 0x20,
             "BotKill" : 0x21,
             "BotMove" : 0x22,
             "BotStats" : 0x24,
        }


class BotDeathCollector(object):
    def collect(self):
        c = CounterMetricFamily("deaths", "Deaths of each player", labels=["name", "rev"])
        for user in game.users.users:
            user = game.users.users.get(user)
            c.add_metric([user.name, str(user.current_rev)], user.deaths)
        yield c


class BotKillCollector(object):
    def collect(self):
        c = CounterMetricFamily("kills", "Kills of each player", labels=["name", "rev"])
        for user in game.users.users:
            user = game.users.users.get(user)
            c.add_metric([user.name, str(user.current_rev)], user.kills)
        yield c


class BotRevisionCountCollector(object):
    def collect(self):
        c = CounterMetricFamily("code_revs", "Number of code revisions", labels=["name"])
        for user in game.users.users:
            user = game.users.users.get(user)
            c.add_metric([user.name], len(user.revs))
        yield c


class BotMassGaugeCollector(object):
    def collect(self):
        c = GaugeMetricFamily("mass", "Mass of each snake", labels=["name", "snake_id", "code_id"])
        for user in game.users.users:
            user = game.users.users.get(user)
            rev = user.revs.get(user.current_rev)
            c.add_metric([user.name, str(user.current_worm), str(user.current_rev)], rev.mass)
        yield c


class BotFoodCarrionCollector(object):
    def collect(self):
        c = CounterMetricFamily("carrion_food", "Carrion food consumed", labels=["name", "snake_id", "code_id"])
        for user in game.users.users:
            user = game.users.users.get(user)
            rev = user.revs.get(user.current_rev)
            c.add_metric([user.name, str(user.current_worm), str(user.current_rev)], rev.carrion_food_consumed)
        yield c


class BotFoodNaturalCollector(object):
    def collect(self):
        c = CounterMetricFamily("natural_food", "Natural food consumed", labels=["name", "snake_id", "code_id"])
        for user in game.users.users:
            user = game.users.users.get(user)
            rev = user.revs.get(user.current_rev)
            c.add_metric([user.name, str(user.current_worm), str(user.current_rev)], rev.natural_food_consumed)
        yield c


class BotFoodHuntedCollector(object):
    def collect(self):
        c = CounterMetricFamily("hunted_food", "Hunted food consumed", labels=["name", "snake_id", "code_id"])
        for user in game.users.users:
            user = game.users.users.get(user)
            rev = user.revs.get(user.current_rev)
            c.add_metric([user.name, str(user.current_worm), str(user.current_rev)], rev.hunted_food_consumed)
        yield c


def main():
    global game
    game = Game()
    global logger
    logger = logging.getLogger("schlafana")
    logger.setLevel(logging.DEBUG)

    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)

    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
    ch.setFormatter(formatter)

    logger.addHandler(ch)
    REGISTRY.register(BotDeathCollector())
    REGISTRY.register(BotKillCollector())
    REGISTRY.register(BotRevisionCountCollector())
    REGISTRY.register(BotMassGaugeCollector())
    REGISTRY.register(BotFoodCarrionCollector())
    REGISTRY.register(BotFoodNaturalCollector())
    REGISTRY.register(BotFoodHuntedCollector())
    start_http_server(9000)
    logger.info("Start prometheus scrape endpoint")

    while True:
        try:
            logger.warning("Trying to connect…")
            global ws
            ws = create_connection("wss://schlangen.bytewerk.org/websocket")
        except Exception as e:
            game.uptime.observe(0)
            time.sleep(5)
        try:
            while True:
                data = ws.recv()
                game.uptime.observe(1)
                game.parse_data(data)
        except Exception as e:
            print(e)
            time.sleep(0.01)


class Game():
    def __init__(self):
        self.x = Gauge("x_resolution", "X resolution of the play field")
        self.y = Gauge("y_resolution", "Y resolution of the play field")
        self.food_decay_per_frame = 0
        self.users = UserList()
        self.uptime = Summary("api_up", "Status of the server API")
        self.frames = Counter("frames", "Frame counter")
        pass


    def parse_data(self, data):
        js = json.loads(data)
        msg_type = js.get("t")

        typ = MESSAGES.get(msg_type, -1)

        if typ == MESSAGES["BotSpawn"]:
            bot = js.get("bot")
            name = bot.get("name")
            self.users.update_by_name(name, data=bot)
        elif typ == MESSAGES["GameInfo"]:
            game.x.set(js.get("world_size_x"))
            game.y.set(js.get("world_size_y"))
        elif typ == MESSAGES["BotKill"]:
            try:
                victim = self.users.get_user_by_worm(js.get("victim_id"))
                killer = self.users.get_user_by_worm(js.get("killer_id"))
                if js.get("killer_id") == js.get("victim_id"):
                    #print("☕ {} killed himself".format(victim.name))
                    pass
                else:
                    #print("☠ {} fragged {}".format(killer.name, victim.name))
                    pass
                killer.kills += 1
                victim.deaths += 1
            except AttributeError as e:
                print(e)
                print("Error in kill: {}".format(js))
            pass
        elif typ == MESSAGES["BotStats"]:
            msg_data = js.get("data")
            for worm in msg_data:
                user = self.users.get_user_by_worm(int(worm))
                if user != None:
                    user.update({ "id" : int(worm), "data" : msg_data.get(worm)})
            pass
        elif typ == MESSAGES["WorldUpdate"]:
            bots = js.get("bots")
            for bot in bots:
                bot = bots.get(bot)
                self.users.update_by_name(bot.get("name"), data=bot)
        elif typ == MESSAGES["Tick"]:
            game.frames.inc()

        elif msg_type != "FoodSpawn" and msg_type != "FoodConsume" and msg_type != "BotMoveHead" and msg_type != "FoodDecay":
            logging.fatal("Unhandled msg_type: {}".format(msg_type))


class BotUser():
    def __init__(self, data):
        self.iterations = []
        self.db_id = None
        self.revs = {}
        self.current_rev = None
        self.__deaths = 0
        self.__kills = 0
        self.parse(data)

    @property
    def kills(self):
        return self.__kills

    @kills.setter
    def kills(self, kills):
        self.__kills = kills
        self.revs.get(self.current_rev).kills += 1

    @property
    def deaths(self):
        return self.__deaths

    @deaths.setter
    def deaths(self, deaths):
        self.__deaths = deaths
        self.revs.get(self.current_rev).deaths += 1

    def parse(self, data):
        self.db_id = data.get("db_id")
        self.name = data.get("name")
        self.dog_tag = data.get("dog_tag")
        self.face = data.get("face")
        self.current_rev = data.get("db_id")
        self.current_worm = data.get("id")
        self.revs.update({ self.current_rev : BotRevision(data) })

    def update(self, data):
        rev_id = data.get("db_id")
        worm_id = data.get("id")
        if rev_id in self.revs.keys(): # revision already known
            rev = self.revs.get(rev_id)
            rev.update(data)
            user = game.users.get_user_by_worm(worm_id)
        else:
            self.current_rev = rev_id
            self.revs.update({ self.current_rev : BotRevision(data) })


class BotRevision():
    def __init__(self, data):
        self.worms = {}
        self.current_worm = data.get("id")
        self.rev_id = data.get("db_id")
        self.worms.update({self.current_worm : [self.rev_id]})
        self.deaths = 0
        self.kills = 0
        self.mass = 0
        self.natural_food_consumed = 0
        self.carrion_food_consumed = 0
        self.hunted_food_consumed = 0
        pass

    def parse(self, data):
        pass

    def update(self, data):
        if "data" in data.keys():
            self.mass = data.get("data").get("m")
            self.carrion_food_consumed = data.get("data").get("c")
            self.natural_food_consumed = data.get("data").get("n")
            self.hunted_food_consumed = data.get("data").get("h")
        else:
            self.current_worm = data.get("id")
            self.rev_id = data.get("db_id")
            self.mass = data.get("mass")
            self.worms.update({self.current_worm : [self.rev_id]})
            pass


class UserList():
    def __init__(self):
        self.users = {}
        self.__usercount = Gauge('Users', 'Total count of active players')

    def update_by_name(self, name, data):
        name = data.get("name")
        user_obj = self.get_user(name)
        if user_obj != None:
            self.users.get(name).update(data)
            self.get_user(name).update(data)
        else:
            self.users.update({name : BotUser(data)})
            #print("{}({}): create user".format(self.users.get(db_id).name, db_id))
        self.__usercount.set(len(self.users))

    def get_user(self, username=""):
        return self.users.get(username, None)

    def get_user_by_rev(self, rev_id):
        for user in self.users:
            user = self.users.get(user)
            for rev in user.revs:
                rev = user.revs.get(rev)
                if rev.id == rev_id:
                    return user

    def get_user_by_worm(self, worm_id):
        for user in self.users:
            user = self.users.get(user)
            for rev in user.revs:
                rev = user.revs.get(rev)
                for worm in rev.worms:
                    if worm == worm_id:
                        return user
                else:
                    #print("No worm found")
                    pass


main()