init

qlpycon.py

@@ -0,0 +1,497 @@
+#!/usr/bin/env python3
+
+import sys
+import time
+import struct
+import argparse
+import uuid
+import threading
+import queue
+import json
+
+import logging
+logger = logging.getLogger('logger')
+logger.setLevel(logging.DEBUG)
+
+import zmq
+import curses
+import curses.textpad
+import signal
+
+import unittest
+
+import locale
+locale.setlocale(locale.LC_ALL, '')
+
+def _readSocketEvent( msg ):
+    event_id = struct.unpack( '<H', msg[:2] )[0]
+    event_names = {
+        zmq.EVENT_ACCEPTED : 'EVENT_ACCEPTED',
+        zmq.EVENT_ACCEPT_FAILED : 'EVENT_ACCEPT_FAILED',
+        zmq.EVENT_BIND_FAILED : 'EVENT_BIND_FAILED',
+        zmq.EVENT_CLOSED : 'EVENT_CLOSED',
+        zmq.EVENT_CLOSE_FAILED : 'EVENT_CLOSE_FAILED',
+        zmq.EVENT_CONNECTED : 'EVENT_CONNECTED',
+        zmq.EVENT_CONNECT_DELAYED : 'EVENT_CONNECT_DELAYED',
+        zmq.EVENT_CONNECT_RETRIED : 'EVENT_CONNECT_RETRIED',
+        zmq.EVENT_DISCONNECTED : 'EVENT_DISCONNECTED',
+        zmq.EVENT_LISTENING : 'EVENT_LISTENING',
+        zmq.EVENT_MONITOR_STOPPED : 'EVENT_MONITOR_STOPPED',
+    }
+    event_name = event_names[ event_id ] if event_id in event_names else '%d' % event_id
+    event_value = struct.unpack( '<I', msg[2:] )[0]
+    return ( event_id, event_name, event_value )
+
+def _checkMonitor( monitor ):
+    try:
+        event_monitor = monitor.recv( zmq.NOBLOCK )
+    except zmq.Again:
+        return
+
+    ( event_id, event_name, event_value ) = _readSocketEvent( event_monitor )
+    event_monitor_endpoint = monitor.recv( zmq.NOBLOCK )
+    logger.info( 'monitor: %s %d endpoint %s' % ( event_name, event_value, event_monitor_endpoint ) )
+    return ( event_id, event_value )
+
+class CursesHandler(logging.Handler):
+    def __init__(self, screen):
+        logging.Handler.__init__(self)
+        self.screen = screen
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            screen = self.screen
+            fs = "%s\n"
+            try:
+                PrintMessageFormatted(screen, fs % msg)
+                screen.refresh()
+            except UnicodeError:
+                PrintMessageFormatted(screen, fs % msg.encode("UTF-8"))
+                screen.refresh()
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except:
+            self.handleError(record)
+
+def setupInputQueue(window):
+    def waitStdin( q ):
+        while ( True ):
+            l = curses.textpad.Textbox(window).edit()
+            if len( l ) > 0 :
+                q.put( l )
+                window.clear()
+                window.refresh()
+    q = queue.Queue()
+    t = threading.Thread( target = waitStdin, args = ( q, ) )
+    t.daemon = True
+    t.start()
+    return q
+
+HOST = 'tcp://127.0.0.1:27961'
+POLL_TIMEOUT = 100
+
+def PrintMessageColored(window, message, attributes):
+    if not curses.has_colors:
+        window.addstr(message)
+        return
+
+    color = 0
+    parse_color = False
+    for ch in message:
+        val = ord( ch )
+        if parse_color:
+            if val >= ord('0') and val <= ord('7'):
+                color = val - ord('0')
+                if color == 7:
+                    color = 0
+            else:
+                window.addch('^', curses.color_pair(color) | attributes)
+                window.addch(ch, curses.color_pair(color) | attributes)
+            parse_color = False
+        elif ch == '^':
+            parse_color = True
+        else:
+            window.addch(ch, curses.color_pair(color) | attributes)
+
+    window.refresh()
+ 
+def PrintMessageFormatted(window, message):
+    attributes = 0
+
+    message = message.replace("\\n", "")
+    message = message.replace(chr(25), "")
+
+    if message[:10] == "broadcast:":
+        message = message[11:]
+        attributes = curses.A_BOLD
+
+    if message[:7] == "print \"":
+        message = message[7:-2] + "\n"
+
+    PrintMessageColored(window, message, attributes)
+
+def ParseGameEvent(message):
+    """Parse JSON game events and return formatted message"""
+    try:
+        jObject = json.loads(message)
+        if 'TYPE' not in jObject or 'DATA' not in jObject:
+            return None
+            
+        event_type = jObject['TYPE']
+        data = jObject['DATA']
+        
+        warmup = data.get('WARMUP', False)
+        warmup_suffix = " ^7(^3warmup^7)" if warmup else ""
+        
+        if event_type == 'PLAYER_SWITCHTEAM':
+            if 'KILLER' not in data:
+                return None
+            killer = data['KILLER']
+            team = killer.get('TEAM', '')
+            old_team = killer.get('OLD_TEAM', '')
+            
+            if team == old_team:
+                return None
+                
+            team_msg = {
+                'FREE': ' ^7Joined the ^2Fight^7',
+                'SPECTATOR': ' ^7Joined the ^3Spectators^7',
+                'RED': ' ^7Joined the ^1Red ^7team',
+                'BLUE': ' ^7Joined the ^4Blue ^7team'
+            }.get(team, ' ^7Joined team %s^7' % team)
+            
+            old_team_msg = {
+                'FREE': 'the ^2Fight^7',
+                'SPECTATOR': 'the ^3Spectators^7',
+                'RED': '^7the ^1Red ^7team',
+                'BLUE': '^7the ^4Blue ^7team'
+            }.get(old_team, 'team %s' % old_team)
+            
+            return "%s%s from %s%s\n" % (killer.get('NAME', 'Unknown'), team_msg, old_team_msg, warmup_suffix)
+        
+        elif event_type == 'PLAYER_DEATH':
+            if 'VICTIM' not in data:
+                return None
+            victim = data['VICTIM']
+            victim_name = victim.get('NAME', 'Unknown')
+            
+            if 'KILLER' not in data or not data['KILLER']:
+                mod = data.get('MOD', 'UNKNOWN')
+                death_msgs = {
+                    'FALLING': "%s ^1CRATERED^7" % victim_name,
+                    'HURT': "%s ^1WAS IN THE WRONG PLACE^7" % victim_name,
+                    'LAVA': "%s ^1DOES A BACKFLIP INTO THE LAVA^7" % victim_name,
+                    'WATER': "%s ^1SANK LIKE A ROCK^7" % victim_name,
+                    'SLIME': "%s ^1MELTED^7" % victim_name,
+                    'CRUSH': "%s ^1WAS CRUSHED^7" % victim_name
+                }
+                msg = death_msgs.get(mod, "%s ^1DIED FROM %s^7" % (victim_name, mod))
+                return "%s%s\n" % (msg, warmup_suffix)
+            
+            killer = data['KILLER']
+            killer_name = killer.get('NAME', 'Unknown')
+            
+            if killer_name != victim_name:
+                killer_team = killer.get('TEAM', '')
+                victim_team = victim.get('TEAM', '')
+                killer_color = '^7(^1RED^7)' if killer_team == '1' else '^7(^4BLUE^7)'
+                victim_color = '^7(^1RED^7)' if victim_team == '1' else '^7(^4BLUE^7)'
+                
+                weapon = killer.get('WEAPON', 'UNKNOWN')
+                weapon_names = {
+                    'ROCKET': 'the Rocket Launcher',
+                    'LIGHTNING': 'the Lightning Gun',
+                    'RAILGUN': 'the Railgun',
+                    'SHOTGUN': 'the Shotgun',
+                    'GAUNTLET': 'the Gauntlet',
+                    'GRENADE': 'the Grenade Launcher',
+                    'PLASMA': 'the Plasma Gun',
+                    'MACHINEGUN': 'the Machine Gun'
+                }
+                weapon_name = weapon_names.get(weapon, 'the %s' % weapon)
+                
+                return "%s %s ^5killed %s %s ^7with %s%s\n" % (killer_color, killer_name, victim_color, victim_name, weapon_name, warmup_suffix)
+            
+            else:
+                weapon = killer.get('WEAPON', 'OTHER_WEAPON')
+                if weapon != 'OTHER_WEAPON':
+                    weapon_names = {
+                        'ROCKET': 'Rocket Launcher',
+                        'PLASMA': 'Plasma Gun',
+                        'GRENADE': 'Grenade Launcher'
+                    }
+                    weapon_name = weapon_names.get(weapon, weapon)
+                    return "%s ^7committed suicide with the ^1%s%s\n" % (killer_name, weapon_name, warmup_suffix)
+        
+        elif event_type == 'PLAYER_MEDAL':
+            name = data.get('NAME', 'Unknown')
+            medal = data.get('MEDAL', 'UNKNOWN')
+            return "%s ^7got a ^6%s ^7medal%s\n" % (name, medal, warmup_suffix)
+            
+    except (json.JSONDecodeError, KeyError, TypeError) as e:
+        return None
+    
+    return None
+
+def InitWindows(screen, args):
+    logger.handlers = []
+    curses.endwin()
+
+    curses.initscr()
+    screen.nodelay(1)
+    curses.start_color()
+    curses.cbreak()
+    curses.setsyx(-1, -1)
+    screen.addstr("Quake Live rcon: %s" % args.host)
+    screen.refresh()
+    maxy, maxx = screen.getmaxyx()
+
+    for i in range(1,7):
+        curses.init_pair(i, i, 0)
+
+    curses.init_pair(5, 6, 0)
+    curses.init_pair(6, 5, 0)
+
+    begin_x = 2; width = maxx - 4
+    begin_y = 2; height = maxy - 5
+    output_window = curses.newwin(height, width, begin_y, begin_x)
+    screen.refresh()
+    output_window.scrollok(True)
+    output_window.idlok(True)
+    output_window.leaveok(True)
+    output_window.refresh()
+
+    begin_x = 4; width = maxx - 6
+    begin_y = maxy - 2; height = 1
+    input_window = curses.newwin(height, width, begin_y, begin_x)
+    screen.addstr(begin_y, begin_x - 2, '>')
+    screen.refresh()
+    input_window.idlok(True)
+    input_window.leaveok(False)
+    input_window.refresh()
+
+    begin_x = 2; width = maxx - 4
+    begin_y = maxy - 3; height = 1
+    divider_window = curses.newwin(height, width, begin_y, begin_x)
+    screen.refresh()
+    divider_window.hline(curses.ACS_HLINE, width)
+    divider_window.refresh()
+
+    mh = CursesHandler(output_window)
+    formatterDisplay = logging.Formatter('%(asctime)-8s|%(name)-12s|%(levelname)-6s|%(message)-s', '%H:%M:%S')
+    mh.setFormatter(formatterDisplay)
+    logger.addHandler(mh)
+
+    screen.refresh()
+
+    return input_window, output_window
+
+def main(screen):
+    parser = argparse.ArgumentParser( description = 'Verbose QuakeLive server statistics' )
+    parser.add_argument( '--host', default = HOST, help = 'ZMQ URI to connect to. Defaults to %s' % HOST )
+    parser.add_argument( '--password', required = False )
+    parser.add_argument( '--identity', default = uuid.uuid1().hex, help = 'Specify the socket identity. Random UUID used by default' )
+    parser.add_argument( '-v', '--verbose', action='store_true', help = 'Enable verbose logging' )
+    args = parser.parse_args()
+    
+    if not args.verbose:
+        logger.setLevel(logging.INFO)
+
+    input_window, output_window = InitWindows(screen, args)
+
+    logger.info('*** VERSION 1.0.2 STARTING ***')
+    logger.info('zmq python bindings %s, libzmq version %s' % ( repr( zmq.__version__ ), zmq.zmq_version() ) )
+
+    stats_port = None
+    stats_password = None
+    stats_context = None
+    stats_socket = None
+    stats_connected = False
+    stats_check_counter = 0
+
+    q = setupInputQueue(input_window)
+    try:
+        if args.verbose:
+            logger.info('Initializing ZMQ context...')
+        ctx = zmq.Context()
+        if args.verbose:
+            logger.info('Creating DEALER socket...')
+        socket = ctx.socket( zmq.DEALER )
+        if args.verbose:
+            logger.info('Setting up socket monitor...')
+        monitor = socket.get_monitor_socket( zmq.EVENT_ALL )
+        if ( args.password is not None ):
+            logger.info( 'Setting password for access' )
+            socket.plain_username = b'rcon'
+            socket.plain_password = args.password.encode('utf-8')
+            socket.zap_domain = b'rcon'
+        if args.verbose:
+            logger.info( 'Setting socket identity: %s' % args.identity )
+        socket.setsockopt( zmq.IDENTITY, args.identity.encode('utf-8') )
+        logger.info( 'Connecting to %s...' % args.host )
+        socket.connect( args.host )
+        if args.verbose:
+            logger.info( 'Connection initiated, waiting for events...' )
+        while ( True ):
+            event = socket.poll( POLL_TIMEOUT )
+            event_monitor = _checkMonitor( monitor )
+            if ( event_monitor is not None and event_monitor[0] == zmq.EVENT_CONNECTED ):
+                logger.info( 'Connected! Registering with the server...' )
+                socket.send( b'register' )
+                if args.verbose:
+                    logger.info( 'Registration message sent.' )
+                
+                logger.info( 'Requesting game stats connection info...' )
+                socket.send( b'zmq_stats_password' )
+                socket.send( b'net_port' )
+
+            while ( not q.empty() ):
+                l = q.get()
+                if args.verbose:
+                    logger.info( 'Sending command: %s' % repr( l.strip() ) )
+                socket.send( l.encode('utf-8') )
+
+            if stats_connected and stats_socket:
+                stats_check_counter += 1
+                
+                if stats_check_counter % 100 == 0:
+                    logger.info('Stats polling active (check #%d)' % (stats_check_counter // 100))
+                
+                try:
+                    stats_msg = stats_socket.recv(zmq.NOBLOCK)
+                    if len(stats_msg) > 0:
+                        logger.info('!!! STATS EVENT RECEIVED (%d bytes) !!!' % len(stats_msg))
+                        y,x = curses.getsyx()
+                        stats_str = stats_msg.decode('utf-8', errors='replace')
+                        
+                        logger.info('Stats JSON: %s' % repr(stats_str[:200]))
+                        
+                        parsed_event = ParseGameEvent(stats_str)
+                        if parsed_event:
+                            logger.info('Parsed successfully')
+                            PrintMessageFormatted(output_window, parsed_event)
+                        else:
+                            logger.info('Parse failed, showing raw')
+                            PrintMessageFormatted(output_window, stats_str + '\n')
+                        
+                        curses.setsyx(y,x)
+                        curses.doupdate()
+                except zmq.error.Again:
+                    pass
+                except Exception as e:
+                    logger.error('Stats socket error: %s' % e)
+                    import traceback
+                    logger.error(traceback.format_exc())
+
+            if ( event == 0 ):
+                continue
+
+            if args.verbose:
+                logger.debug( 'Socket has data available, reading messages...' )
+            msg_count = 0
+            while ( True ):
+                try:
+                    msg = socket.recv( zmq.NOBLOCK )
+                    msg_count += 1
+                except zmq.error.Again:
+                    if args.verbose and msg_count > 0:
+                        logger.debug( 'Read %d message(s) from socket.' % msg_count )
+                    break
+                except Exception as e:
+                    logger.error( 'Error receiving message: %s' % e )
+                    break
+                else:
+                    if len( msg ) > 0:
+                        if args.verbose:
+                            logger.debug( 'Received message (%d bytes): %s' % (len(msg), repr(msg[:100])) )
+                        y,x = curses.getsyx()
+                        msg_str = msg.decode('utf-8', errors='replace')
+                        
+                        if 'net_port' in msg_str and ' is:' in msg_str and '"net_port"' in msg_str:
+                            import re
+                            match = re.search(r' is:"([^"]+)"', msg_str)
+                            if match:
+                                port_str = match.group(1)
+                                port_str = port_str.strip()
+                                digit_match = re.search(r'(\d+)', port_str)
+                                if digit_match and stats_port is None:
+                                    stats_port = digit_match.group(1)
+                                    logger.info('>>> Got stats port: %s' % stats_port)
+                        
+                        if 'zmq_stats_password' in msg_str and ' is:' in msg_str and '"zmq_stats_password"' in msg_str:
+                            import re
+                            match = re.search(r' is:"([^"]+)"', msg_str)
+                            if match and stats_password is None:
+                                password_str = match.group(1)
+                                password_str = re.sub(r'\^\d', '', password_str)
+                                stats_password = password_str.strip()
+                                logger.info('>>> Got stats password: %s' % stats_password)
+                        
+                        if stats_port and stats_password and not stats_connected:
+                            try:
+                                logger.info('>>> CONNECTING TO STATS PORT <<<')
+                                logger.info('    Host: %s:%s' % (args.host.split('//')[1].split(':')[0], stats_port))
+                                logger.info('    Password: %s' % stats_password)
+                                
+                                host_ip = args.host.split('//')[1].split(':')[0]
+                                stats_host = 'tcp://%s:%s' % (host_ip, stats_port)
+                                
+                                time.sleep(0.5)
+                                
+                                stats_context = zmq.Context()
+                                stats_socket = stats_context.socket(zmq.SUB)
+                                logger.info('>>> Stats socket created (SUB type)')
+                                
+                                if stats_password and stats_password.strip():
+                                    logger.info('>>> Setting PLAIN authentication')
+                                    stats_socket.setsockopt(zmq.PLAIN_USERNAME, b'stats')
+                                    stats_socket.setsockopt(zmq.PLAIN_PASSWORD, stats_password.encode('utf-8'))
+                                    stats_socket.setsockopt_string(zmq.ZAP_DOMAIN, 'stats')
+                                    logger.info('>>> Auth configured')
+                                
+                                logger.info('>>> Connecting to %s' % stats_host)
+                                stats_socket.connect(stats_host)
+                                logger.info('>>> Connected')
+                                
+                                logger.info('>>> Setting ZMQ_SUBSCRIBE to empty (all messages)')
+                                stats_socket.setsockopt(zmq.SUBSCRIBE, b'')
+                                logger.info('>>> Subscribed')
+                                
+                                time.sleep(0.5)
+                                
+                                stats_connected = True
+                                logger.info('>>> STATS CONNECTION COMPLETE <<<')
+                                logger.info('>>> Waiting for game events...')
+                            except Exception as e:
+                                logger.error('!!! Stats connection FAILED: %s' % e)
+                                import traceback
+                                logger.error(traceback.format_exc())
+                        
+                        parsed_event = ParseGameEvent(msg_str)
+                        if parsed_event:
+                            PrintMessageFormatted(output_window, parsed_event)
+                        else:
+                            PrintMessageFormatted(output_window, msg_str)
+                        
+                        curses.setsyx(y,x)
+                        curses.doupdate()
+                    else:
+                        if args.verbose:
+                            logger.debug( 'Received empty message (possible keepalive or protocol frame)' )
+
+    except KeyboardInterrupt:
+        logger.info( 'Received keyboard interrupt, shutting down...' )
+    except Exception as e:
+        logger.error( 'Fatal error: %s' % e )
+        import traceback
+        logger.error( traceback.format_exc() )
+    finally:
+        if stats_socket:
+            stats_socket.close()
+        if stats_context:
+            stats_context.term()
+
+if ( __name__ == '__main__' ):
+    curses.wrapper(main)
+
+# Version: 1.0.2 - Python 3 ZMQ Rcon Client with Stats Support