mc-status-page/mc.py

"""
Application to ping Minecraft server and write all data to postgresql database.
"""

import os
import sys
import time
import datetime
import logging
import argparse

from mcstatus import JavaServer as MinecraftServer
import psycopg2
from typing import Tuple, Dict

from dotenv import load_dotenv

# read environment variables from .env file
load_dotenv()

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')

# get DB_HOST, DB_PROJ, DB_USER, DB_PASS, and DB_NAME from environment variables
DB_HOST = os.environ.get('DB_HOST', 'localhost')
DB_NAME = os.environ.get('DB_NAME', 'main')
DB_TABLE = os.environ.get('DB_TABLE', 'minecraft')
DB_PROJ = os.environ.get('DB_PROJ', 'default')
DB_USER = os.environ.get('DB_USER', 'user')
DB_PASS = os.environ.get('DB_PASS', 'password')
IP_ADD = os.environ.get('MINE_IP', 'localhost')


def parse_args() -> argparse.Namespace:
    # Set up command line arguments
    parser = argparse.ArgumentParser(description='Monitor Minecraft server activity and write it to postgresql database.')
    parser.add_argument('--dbhost', default=DB_HOST, help='Postgresql database host')
    parser.add_argument('--dbname', default=DB_NAME, help='Postgresql database name')
    parser.add_argument('--dbuser', default=DB_USER, help='Postgresql database user')
    parser.add_argument('--dbpass', default=DB_PASS, help='Postgresql database password')
    parser.add_argument('--project', default=DB_PROJ, help='Project name')
    parser.add_argument('--table', default=DB_TABLE, help='Table name')
    parser.add_argument('--ip', default=IP_ADD, help='IP Address of Minecraft Server')
    parser.add_argument('--interval', default=60, type=int, help='Interval in seconds between measurements')
    parser.add_argument('--clean', action='store_true', help='Clean database before starting (default: False)')
    args = parser.parse_args()
    return args

# Connect to database
def connect_to_database(args):
        
    try:
        conn = psycopg2.connect(
            host=args.dbhost,
            database=args.dbname,
            user=args.dbuser,
            password=args.dbpass,
            sslmode='require',
            options=f'project={args.project}'
        )
        cur = conn.cursor()
    except Exception as e:
        logging.error('Failed to connect to database: %s', e)
        sys.exit(1)

    table_name = vars(args).get('table', DB_TABLE)
    # if args.clean doesn't exist, it will be False
    if vars(args).get('clean', False):
        # remove table if it exists
        try:
            cur.execute(f'DROP TABLE IF EXISTS {table_name}')
            conn.commit()
        except Exception as e:
            logging.error('Failed to drop table: %s', e)
            sys.exit(1)

    # Create table if it doesn't exist
    try:
        cur.execute(f'CREATE TABLE IF NOT EXISTS {table_name} (timestamp TIMESTAMP, server_name TEXT, version TEXT, protocol INT, players INT, max_players INT, latency INT)')
        conn.commit()
    except Exception as e:
        logging.error('Failed to create table: %s', e)
        sys.exit(1)

    return conn


def get_minecraft_server_status(host: str) -> Dict[str, str]:
    # Get Minecraft server status
    try:
        server = MinecraftServer.lookup(host)
        status = server.status()
        logging.info('Successfully got Minecraft server status')
    except Exception as e:
        logging.error('Failed to get Minecraft server status: %s', e)
        sys.exit(1)

    # extract status as dictionary
    status = status.__dict__
    return status


def main(args, conn=None):
    cur = conn.cursor()
    try:
        # Get current time
        timestamp = datetime.datetime.now()

        # Get Minecraft server status
        ip_add = vars(args).get('ip', 'localhost:25571')
        status = get_minecraft_server_status(ip_add)
        table_name = vars(args).get('table', DB_TABLE)
        if status:
            # Write to postgres the status
            cur.execute(f'INSERT INTO {table_name} (timestamp, server_name, version, protocol, players, max_players, latency) VALUES (%s, %s, %s, %s, %s, %s, %s)', (timestamp, status['raw']['description'], status['raw']['version']['name'], status['raw']['version']['protocol'], status['raw']['players']['online'], status['raw']['players']['max'], status['latency']))
            conn.commit()
        # Sleep for interval
        time.sleep(args.interval)
    except KeyboardInterrupt:
        logging.info('Exiting...')
        sys.exit(0)
    except Exception as e:
        logging.error('Failed to write to database: %s', e)
        # sys.exit(1)


if __name__ == '__main__':
    # Main loop
    args = parse_args()
    conn = connect_to_database(args)
    while True:
        # if connection to server is lost, reconnect
        if not conn:
            logging.info('Reconnecting to database...')
            conn = connect_to_database(args)

        main(args, conn=conn)